Skip to content

Commit

Permalink
remove prev iteration of generator
Browse files Browse the repository at this point in the history
  • Loading branch information
@makslevental
makslevental committed Dec 30, 2024
1 parent dc70013 commit fa1c36c
Show file tree
Hide file tree
Showing 11 changed files with 151 additions and 6,071 deletions.
97 changes: 27 additions & 70 deletions projects/eudsl-py/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -180,10 +180,10 @@ eudslpygen(EUDSLGenDLTIDialect
NAMESPACES mlir::DLTIDialect mlir::DLTIDialect::detail
)

#eudslpygen(EUDSLGenemitc
# ${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
# NAMESPACES mlir::emitc mlir::emitc::detail
#)
eudslpygen(EUDSLGenemitc
${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
NAMESPACES mlir::emitc mlir::emitc::detail
)

eudslpygen(EUDSLGenfunc
${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
Expand Down Expand Up @@ -216,10 +216,10 @@ eudslpygen(EUDSLGenlinalg
NAMESPACES mlir::linalg mlir::linalg::detail
)

#eudslpygen(EUDSLGenLLVM
# ${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
# NAMESPACES mlir::LLVM mlir::LLVM::detail
#)
eudslpygen(EUDSLGenLLVM
${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
NAMESPACES mlir::LLVM mlir::LLVM::detail
)

eudslpygen(EUDSLGenmath
${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
Expand Down Expand Up @@ -323,6 +323,7 @@ eudslpygen(EUDSLGenspirv
NAMESPACES mlir::spirv mlir::spirv::detail
EXTRA_INCLUDES
mlir/Dialect/SPIRV/IR/SPIRVOps.h
mlir/IR/PatternMatch.h
)

eudslpygen(EUDSLGentensor
Expand Down Expand Up @@ -355,62 +356,17 @@ eudslpygen(EUDSLGenx86vector
NAMESPACES mlir::x86vector mlir::x86vector::detail
)

#eudslpygen(EUDSLGenxegpu
# ${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
# NAMESPACES mlir::xegpu mlir::xegpu::detail
#)
eudslpygen(EUDSLGenxegpu
${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
NAMESPACES mlir::xegpu mlir::xegpu::detail
)

nanobind_add_module(eudslpy_ext
NB_STATIC
NB_DOMAIN eudsl
src/eudslpy_ext.cpp
${EUDSLGenacc_GENERATED_SHARDS}
${EUDSLGenaffine_GENERATED_SHARDS}
${EUDSLGenamdgpu_GENERATED_SHARDS}
${EUDSLGenamx_GENERATED_SHARDS}
${EUDSLGenarith_GENERATED_SHARDS}
${EUDSLGenarm_neon_GENERATED_SHARDS}
${EUDSLGenarm_sme_GENERATED_SHARDS}
${EUDSLGenarm_sve_GENERATED_SHARDS}
${EUDSLGenasync_GENERATED_SHARDS}
${EUDSLGenbufferization_GENERATED_SHARDS}
${EUDSLGencf_GENERATED_SHARDS}
${EUDSLGencomplex_GENERATED_SHARDS}
${EUDSLGenDLTIDialect_GENERATED_SHARDS}
# ${EUDSLGenemitc_GENERATED_SHARDS}
${EUDSLGenfunc_GENERATED_SHARDS}
${EUDSLGengpu_GENERATED_SHARDS}
${EUDSLGenindex_GENERATED_SHARDS}
# ${EUDSLGenirdl_GENERATED_SHARDS}
# ${EUDSLGenlinalg_GENERATED_SHARDS}
# ${EUDSLGenLLVM_GENERATED_SHARDS}
${EUDSLGenmath_GENERATED_SHARDS}
${EUDSLGenmemref_GENERATED_SHARDS}
${EUDSLGenmesh_GENERATED_SHARDS}
${EUDSLGenml_program_GENERATED_SHARDS}
${EUDSLGenmpi_GENERATED_SHARDS}
${EUDSLGennvgpu_GENERATED_SHARDS}
${EUDSLGenNVVM_GENERATED_SHARDS}
# ${EUDSLGenomp_GENERATED_SHARDS}
${EUDSLGenpdl_GENERATED_SHARDS}
${EUDSLGenpdl_interp_GENERATED_SHARDS}
${EUDSLGenpolynomial_GENERATED_SHARDS}
${EUDSLGenptr_GENERATED_SHARDS}
${EUDSLGenquant_GENERATED_SHARDS}
${EUDSLGenROCDL_GENERATED_SHARDS}
# ${EUDSLGenscf_GENERATED_SHARDS}
# ${EUDSLGenshape_GENERATED_SHARDS}
${EUDSLGensparse_tensor_GENERATED_SHARDS}
# ${EUDSLGenspirv_GENERATED_SHARDS}
${EUDSLGentensor_GENERATED_SHARDS}
${EUDSLGentosa_GENERATED_SHARDS}
${EUDSLGentransform_GENERATED_SHARDS}
${EUDSLGenub_GENERATED_SHARDS}
# ${EUDSLGenvector_GENERATED_SHARDS}
${EUDSLGenx86vector_GENERATED_SHARDS}
# ${EUDSLGenxegpu_GENERATED_SHARDS}
)
add_dependencies(eudslpy_ext
)
target_link_libraries(eudslpy_ext PRIVATE
EUDSLGenacc
EUDSLGenaffine
EUDSLGenamdgpu
Expand All @@ -424,38 +380,38 @@ add_dependencies(eudslpy_ext
EUDSLGencf
EUDSLGencomplex
EUDSLGenDLTIDialect
# EUDSLGenemitc
EUDSLGenemitc
EUDSLGenfunc
EUDSLGengpu
EUDSLGenindex
# EUDSLGenirdl
# EUDSLGenlinalg
# EUDSLGenLLVM
# EUDSLGenirdl
EUDSLGenlinalg
EUDSLGenLLVM
EUDSLGenmath
EUDSLGenmemref
EUDSLGenmesh
EUDSLGenml_program
EUDSLGenmpi
EUDSLGennvgpu
EUDSLGenNVVM
# EUDSLGenomp
# EUDSLGenomp
EUDSLGenpdl
EUDSLGenpdl_interp
EUDSLGenpolynomial
EUDSLGenptr
EUDSLGenquant
EUDSLGenROCDL
# EUDSLGenscf
# EUDSLGenshape
EUDSLGenscf
EUDSLGenshape
EUDSLGensparse_tensor
# EUDSLGenspirv
EUDSLGenspirv
EUDSLGentensor
EUDSLGentosa
EUDSLGentransform
EUDSLGenub
# EUDSLGenvector
# EUDSLGenvector
EUDSLGenx86vector
# EUDSLGenxegpu
EUDSLGenxegpu
)
set_target_properties(eudslpy_ext
PROPERTIES
Expand Down Expand Up @@ -503,7 +459,8 @@ target_compile_options(eudslpy_ext PRIVATE ${nanobind_options})
target_compile_options(nanobind-static PRIVATE ${nanobind_options})

# note WORKING_DIRECTORY
set(NB_STUBGEN_CMD "${Python_EXECUTABLE}" "-m" "nanobind.stubgen" --module eudslpy_ext --recursive --include-private --output-dir .)
set(NB_STUBGEN_CMD "${Python_EXECUTABLE}" "-m" "nanobind.stubgen"
--module eudslpy_ext --recursive --include-private --output-dir .)
set(NB_STUBGEN_OUTPUTS "__init__.pyi")
add_custom_command(
OUTPUT ${NB_STUBGEN_OUTPUTS}
Expand Down
55 changes: 30 additions & 25 deletions projects/eudsl-py/cmake/EUDSLPYConfig.cmake
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,5 @@
# copy-pasta from AddMLIR.cmake/AddLLVM.cmake/TableGen.cmake

# Clear out any pre-existing compile_commands file before processing. This
# allows for generating a clean compile_commands on each configure.
file(REMOVE ${CMAKE_BINARY_DIR}/eudslpygen_compile_commands.yml)

# no clue why but with LLVM_LINK_LLVM_DYLIB even static libs depend on LLVM
get_property(MLIR_ALL_LIBS GLOBAL PROPERTY MLIR_ALL_LIBS)
foreach(_lib ${MLIR_ALL_LIBS})
Expand Down Expand Up @@ -52,7 +48,7 @@ function(eudslpygen target inputFile)
# Filter out empty items before prepending each entry with -I
list(REMOVE_ITEM eudslpygen_includes "")
list(TRANSFORM eudslpygen_includes PREPEND -I)

set(_gen_target_dir "${CMAKE_CURRENT_BINARY_DIR}/generated/${target}")
file(MAKE_DIRECTORY ${_gen_target_dir})
set(fullGenFile "${_gen_target_dir}/${target}.cpp.gen")
Expand All @@ -63,12 +59,16 @@ function(eudslpygen target inputFile)
# this could be smarter by asking people to list td targets or something but that's too onerous
file(GLOB_RECURSE global_tds "${MLIR_INCLUDE_DIR}/mlir/*.td")
# use cc -MM to collect all transitive headers
set(clang_command ${CMAKE_CXX_COMPILER} -v -xc++ "-std=c++${CMAKE_CXX_STANDARD}"
set(clang_command ${CMAKE_CXX_COMPILER}
# -v
-xc++ "-std=c++${CMAKE_CXX_STANDARD}"
-MM ${EUDSLPYGEN_TARGET_DEFINITIONS_ABSOLUTE}
-MT ${fullGenFile_rel}
${eudslpygen_includes}
-o ${_depfile})
execute_process(COMMAND ${clang_command} RESULT_VARIABLE had_error COMMAND_ECHO STDERR)
execute_process(COMMAND ${clang_command} RESULT_VARIABLE had_error
# COMMAND_ECHO STDERR
)
if(had_error OR NOT EXISTS "${_depfile}")
set(additional_cmdline -o "${fullGenFile_rel}")
else()
Expand Down Expand Up @@ -126,36 +126,41 @@ function(eudslpygen target inputFile)
DEPENDS ${ARG_DEPENDS} ${eudslpygen_depends} ${local_headers} ${global_tds}
COMMENT "EUDSLPY: Generating ${fullGenFile}..."
)
# this is the specific thing connected the dependencies...
set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${fullGenFile})

# epic hack to specify all shards that will be generated even though we don't know them before hand
# TODO(max): refactor eudslpy-gen into its own subproject so that we can do execute_process(CMAKE_COMMAND... )
set(_byproducts)
# lol spirv has 260 ops
set(_max_num_shards 30)
foreach(i RANGE ${_max_num_shards})
list(APPEND _byproducts "${fullGenFile}.shard.${i}.cpp")
endforeach()

add_custom_command(OUTPUT "${fullGenFile}.sharded.cpp"
COMMAND ${Python_EXECUTABLE} ${CMAKE_CURRENT_FUNCTION_LIST_DIR}/make_generated_registration.py
${fullGenFile} -t ${target} -I ${ARG_EXTRA_INCLUDES} ${EUDSLPYGEN_TARGET_DEFINITIONS_ABSOLUTE}
-m ${_max_num_shards}
BYPRODUCTS ${_byproducts}
WORKING_DIRECTORY ${CMAKE_BINARY_DIR}
DEPENDS ${fullGenFile}
COMMENT "EUDSLPY: Generating ${fullGenFile}.sharded.cpp..."
)

# this is the specific thing connected the dependencies...
set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${fullGenFile}.sharded.cpp")
file(GLOB _generated_shards "${_gen_target_dir}/*.shard.*")
list(APPEND _generated_shards "${fullGenFile}.sharded.cpp")
set(${target}_GENERATED_SHARDS ${_generated_shards} PARENT_SCOPE)

add_library(${target} STATIC "${fullGenFile}.sharded.cpp" ${_byproducts})
execute_process(
COMMAND "${Python_EXECUTABLE}" -m nanobind --include_dir
OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE nanobind_include_dir)
target_include_directories(${target} PRIVATE ${eudslpygen_includes}
${Python_INCLUDE_DIRS} ${nanobind_include_dir})

# `make clean' must remove all those generated files:
# TODO(max): clean up dep files
set_property(DIRECTORY APPEND PROPERTY ADDITIONAL_MAKE_CLEAN_FILES ${fullGenFile})
set_source_files_properties(${CMAKE_CURRENT_BINARY_DIR}/${fullGenFile} PROPERTIES
GENERATED 1)

# Append the includes used for this file to the pdll_compilation_commands
# file.
file(APPEND ${CMAKE_BINARY_DIR}/eudslpygen_compile_commands.yml
"--- !FileInfo:\n"
" filepath: \"${EUDSLPYGEN_TARGET_DEFINITIONS_ABSOLUTE}\"\n"
" includes: \"${CMAKE_CURRENT_SOURCE_DIR};${eudslpygen_includes}\"\n"
)

add_public_eudslpygen_target(${target} "${fullGenFile}.sharded.cpp;${_generated_shards}")
set_property(DIRECTORY ${CMAKE_CURRENT_LIST_DIR}
APPEND PROPERTY ADDITIONAL_MAKE_CLEAN_FILES ${_byproducts})
set_source_files_properties(${_byproducts} PROPERTIES GENERATED 1)
endfunction()

macro(add_eudslpygen target project)
Expand Down
19 changes: 16 additions & 3 deletions projects/eudsl-py/cmake/make_generated_registration.py
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
import argparse
import re
import sys
from pathlib import Path
from textwrap import dedent


def make_source_shards(filename: Path, target, extra_includes):
def make_source_shards(filename: Path, target, extra_includes, max_num_shards):
assert filename.name.endswith("cpp.gen"), "expected .cpp.gen file"
with open(filename) as f:
source = f.read()
Expand Down Expand Up @@ -33,6 +34,12 @@ def make_source_shards(filename: Path, target, extra_includes):
print(shar, file=f)
print("}", file=f)

if len(shards) > max_num_shards:
raise RuntimeError("expected less than 20 shards")
for i in range(len(shards), max_num_shards):
with open(f"{filename}.shard.{i}.cpp", "w") as f:
print(f"// dummy shard {i}", file=f)

with open(f"{filename}.sharded.cpp", "w") as f:
print(
dedent(
Expand All @@ -45,7 +52,10 @@ def make_source_shards(filename: Path, target, extra_includes):
file=f,
)
for i in range(len(shards)):
print(dedent(f"extern void populate{target}{i}Module(nb::module_ &m);"), file=f)
print(
dedent(f"extern void populate{target}{i}Module(nb::module_ &m);"),
file=f,
)
for i in range(len(shards)):
print(dedent(f"populate{target}{i}Module(m);"), file=f)

Expand All @@ -57,5 +67,8 @@ def make_source_shards(filename: Path, target, extra_includes):
parser.add_argument("filename")
parser.add_argument("-t", "--target")
parser.add_argument("-I", "--extra_includes", nargs="*")
parser.add_argument("-m", "--max-num-shards", type=int, default=20)
args = parser.parse_args()
make_source_shards(Path(args.filename), args.target, args.extra_includes)
make_source_shards(
Path(args.filename), args.target, args.extra_includes, args.max_num_shards + 1
)
Loading

0 comments on commit fa1c36c

Please sign in to comment.