types working

This fixes test breakage on clang bots. See comment
llvm/llvm-project#121081 (comment)

Author: makslevental

projects/eudsl-py/CMakeLists.txt

@@ -107,10 +107,11 @@ add_eudslpygen(eudslpy-gen EUDSLPY
   src/eudslpy-gen.cpp
 )
 
-eudslpygen(EUDSLGenacc
-  ${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
-  NAMESPACES mlir::acc mlir::acc::detail
-)
+# too big
+# eudslpygen(EUDSLGenacc
+#   ${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
+#   NAMESPACES mlir::acc mlir::acc::detail
+# )
 
 eudslpygen(EUDSLGenaffine
   ${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
@@ -140,10 +141,11 @@ eudslpygen(EUDSLGenarm_neon
   NAMESPACES mlir::arm_neon mlir::arm_neon::detail
 )
 
-eudslpygen(EUDSLGenarm_sme
-  ${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
-  NAMESPACES mlir::arm_sme mlir::arm_sme::detail
-)
+# too big
+# eudslpygen(EUDSLGenarm_sme
+#   ${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
+#   NAMESPACES mlir::arm_sme mlir::arm_sme::detail
+# )
 
 eudslpygen(EUDSLGenarm_sve
   ${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
@@ -180,10 +182,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
@@ -216,10 +218,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
@@ -318,12 +320,14 @@ eudslpygen(EUDSLGensparse_tensor
   NAMESPACES mlir::sparse_tensor mlir::sparse_tensor::detail
 )
 
-eudslpygen(EUDSLGenspirv
-  ${MLIR_INCLUDE_DIR}/mlir/Dialect/SPIRV/IR/SPIRVOps.h
-  NAMESPACES mlir::spirv mlir::spirv::detail
-  EXTRA_INCLUDES
-    mlir/Dialect/SPIRV/IR/SPIRVOps.h
-)
+# nb::detail::nb_func_new("get_vce_triple_attr_name"): mismatched static/instance method flags in function overloads!
+# eudslpygen(EUDSLGenspirv
+#   ${MLIR_INCLUDE_DIR}/mlir/Dialect/SPIRV/IR/SPIRVOps.h
+#   NAMESPACES mlir::spirv mlir::spirv::detail
+#   EXTRA_INCLUDES
+#     mlir/Dialect/SPIRV/IR/SPIRVOps.h
+#     mlir/IR/PatternMatch.h
+# )
 
 eudslpygen(EUDSLGentensor
   ${MLIR_INCLUDE_DIR}/mlir/InitAllDialects.h
@@ -355,107 +359,62 @@ 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
-  EUDSLGenacc
+)
+target_link_libraries(eudslpy_ext PRIVATE
+  # EUDSLGenacc
   EUDSLGenaffine
   EUDSLGenamdgpu
   EUDSLGenamx
   EUDSLGenarith
   EUDSLGenarm_neon
-  EUDSLGenarm_sme
+  # EUDSLGenarm_sme
   EUDSLGenarm_sve
   EUDSLGenasync
   EUDSLGenbufferization
   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
   EUDSLGentensor
   EUDSLGentosa
   EUDSLGentransform
   EUDSLGenub
-#  EUDSLGenvector
+  #  EUDSLGenvector
   EUDSLGenx86vector
-#  EUDSLGenxegpu
+  EUDSLGenxegpu
 )
 set_target_properties(eudslpy_ext
   PROPERTIES
@@ -503,7 +462,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}

projects/eudsl-py/cmake/EUDSLPYConfig.cmake

@@ -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})
@@ -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")
@@ -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()
@@ -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)

projects/eudsl-py/cmake/make_generated_registration.py

@@ -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()
@@ -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(
@@ -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)
 
@@ -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
+    )