[llvm] [mlir] [mlir][py] Make Python enum gen consistent with C++ type (PR #119407)

Maksim Levental via llvm-commits llvm-commits at lists.llvm.org
Tue Dec 10 08:14:23 PST 2024


https://github.com/makslevental created https://github.com/llvm/llvm-project/pull/119407

In C++ generation we take the main file as the one for which we are generating methods for. E.g.,

mlir-tblgen -gen-type-constraint-decls Foo.td

generates decls for Foo.td rather than anything transitively included. It makes what is acted upon more explicit. Unfortunately enums didn't get the same attention. So one would generate for enums even where files included others. Python side there is an added layer of indirection which is not needed here. Pass the enum file to be generated directly to the tool and filter on main file to avoid special casing.

Now if there were multiple files with enums (instead of common case 1), one would have to invoke the command multiple times and include all. That is less painful Python side than C++ side (given the common import all in the wrapper dialect.py file).

Note: this is just a partial example for discussion.

>From ab53eb773331e25270c716536287cc64ed4075e4 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Fri, 6 Dec 2024 07:19:32 +0000
Subject: [PATCH] [mlir][py] Make Python enum gen consistent with C++ type

In C++ generation we take the main file as the one for which we are generating methods for. E.g.,

mlir-tblgen -gen-type-constraint-decls Foo.td

generates decls for Foo.td rather than anything transitively included. It makes what is acted upon more explicit. Unfortunately enums didn't get the same attention. So one would generate for enums even where files included others. Python side there is an added layer of indirection which is not needed here. Pass the enum file to be generated directly to the tool and filter on main file to avoid special casing.

Now if there were multiple files with enums (instead of common case 1), one would have to invoke the command multiple times and include all. That is less painful Python side than C++ side (given the common import all in the wrapper dialect.py file).

Note: this is just a partial example for discussion.
---
 mlir/include/mlir/Dialect/GPU/IR/GPUBase.td   | 164 ++++++++++++++++++
 mlir/include/mlir/Dialect/GPU/IR/GPUOps.td    | 164 ------------------
 mlir/python/CMakeLists.txt                    |   4 +-
 .../mlir-tblgen/EnumPythonBindingGen.cpp      |   7 +
 .../llvm-project-overlay/mlir/BUILD.bazel     | 133 +++++++++-----
 .../mlir/python/BUILD.bazel                   |   6 +-
 6 files changed, 269 insertions(+), 209 deletions(-)

diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td b/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td
index 860f8933672038..ff1f5d624303d9 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td
@@ -99,6 +99,170 @@ def GPU_AddressSpaceEnum : GPU_I32Enum<
 def GPU_AddressSpaceAttr :
   GPU_I32EnumAttr<"address_space", GPU_AddressSpaceEnum>;
 
+def GPU_Dimension : I32EnumAttr<"Dimension",
+    "a dimension, either 'x', 'y', or 'z'",
+    [
+      I32EnumAttrCase<"x", 0>,
+      I32EnumAttrCase<"y", 1>,
+      I32EnumAttrCase<"z", 2>
+    ]>{
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::gpu";
+}
+def GPU_DimensionAttr : EnumAttr<GPU_Dialect, GPU_Dimension, "dim">;
+
+// These mirror the reduction combining kinds from the vector dialect.
+def GPU_AllReduceOpAdd : I32EnumAttrCase<"ADD", 0, "add">;
+def GPU_AllReduceOpMul : I32EnumAttrCase<"MUL", 1, "mul">;
+def GPU_AllReduceOpMinUI : I32EnumAttrCase<"MINUI", 2, "minui">;
+def GPU_AllReduceOpMinSI : I32EnumAttrCase<"MINSI", 3, "minsi">;
+// Follows the `arith.minnumf` semantics.
+def GPU_AllReduceOpMinnumF : I32EnumAttrCase<"MINNUMF", 4, "minnumf">;
+def GPU_AllReduceOpMaxUI : I32EnumAttrCase<"MAXUI", 5, "maxui">;
+def GPU_AllReduceOpMaxSI : I32EnumAttrCase<"MAXSI", 6, "maxsi">;
+// Follows the `arith.maxnumf` semantics.
+def GPU_AllReduceOpMaxnumF : I32EnumAttrCase<"MAXNUMF", 7, "maxnumf">;
+def GPU_AllReduceOpAnd : I32EnumAttrCase<"AND", 8, "and">;
+def GPU_AllReduceOpOr  : I32EnumAttrCase<"OR",  9, "or">;
+def GPU_AllReduceOpXor : I32EnumAttrCase<"XOR", 10, "xor">;
+// Follows the `arith.minimumf` semantics.
+def GPU_AllReduceOpMinimumF : I32EnumAttrCase<"MINIMUMF", 11, "minimumf">;
+// Follows the `arith.maximumf` semantics.
+def GPU_AllReduceOpMaximumF : I32EnumAttrCase<"MAXIMUMF", 12, "maximumf">;
+
+def GPU_AllReduceOperation : I32EnumAttr<"AllReduceOperation",
+    "built-in reduction operations supported by gpu.allreduce.",
+    [
+      GPU_AllReduceOpAdd,
+      GPU_AllReduceOpMul,
+      GPU_AllReduceOpMinUI,
+      GPU_AllReduceOpMinSI,
+      GPU_AllReduceOpMinnumF,
+      GPU_AllReduceOpMaxUI,
+      GPU_AllReduceOpMaxSI,
+      GPU_AllReduceOpMaxnumF,
+      GPU_AllReduceOpAnd,
+      GPU_AllReduceOpOr,
+      GPU_AllReduceOpXor,
+      GPU_AllReduceOpMinimumF,
+      GPU_AllReduceOpMaximumF
+    ]>{
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::gpu";
+}
+def GPU_AllReduceOperationAttr : EnumAttr<GPU_Dialect, GPU_AllReduceOperation,
+                                          "all_reduce_op">;
+
+def GPU_ShuffleOpXor  : I32EnumAttrCase<"XOR",  0, "xor">;
+def GPU_ShuffleOpDown : I32EnumAttrCase<"DOWN", 1, "down">;
+def GPU_ShuffleOpUp   : I32EnumAttrCase<"UP",   2, "up">;
+def GPU_ShuffleOpIdx  : I32EnumAttrCase<"IDX",  3, "idx">;
+
+def GPU_ShuffleMode : I32EnumAttr<"ShuffleMode",
+    "Indexing modes supported by gpu.shuffle.",
+    [
+      GPU_ShuffleOpXor, GPU_ShuffleOpUp, GPU_ShuffleOpDown, GPU_ShuffleOpIdx,
+    ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::gpu";
+}
+def GPU_ShuffleModeAttr : EnumAttr<GPU_Dialect, GPU_ShuffleMode,
+                                   "shuffle_mode">;
+
+def GPU_ElementwiseOpAddF  : I32EnumAttrCase<"ADDF", 0, "addf">;
+def GPU_ElementwiseOpMulF  : I32EnumAttrCase<"MULF", 1, "mulf">;
+def GPU_ElementwiseOpSUBF  : I32EnumAttrCase<"SUBF", 2, "subf">;
+def GPU_ElementwiseOpMaxF : I32EnumAttrCase<"MAXF", 3, "maxf">;
+def GPU_ElementwiseOpMinF : I32EnumAttrCase<"MINF", 4, "minf">;
+def GPU_ElementwiseOpDivF : I32EnumAttrCase<"DIVF", 5, "divf">;
+def GPU_ElementwiseOpAddI  : I32EnumAttrCase<"ADDI", 6, "addi">;
+def GPU_ElementwiseOpMulI  : I32EnumAttrCase<"MULI", 7, "muli">;
+def GPU_ElementwiseOpSUBI  : I32EnumAttrCase<"SUBI", 8, "subi">;
+def GPU_ElementwiseOpDivS : I32EnumAttrCase<"DIVS", 9, "divs">;
+def GPU_ElementwiseOpDivU : I32EnumAttrCase<"DIVU", 10, "divu">;
+def GPU_ElementwiseOpNEGF : I32EnumAttrCase<"NEGATEF", 11, "negatef">;
+def GPU_ElementwiseOpNEGS : I32EnumAttrCase<"NEGATES", 12, "negates">;
+def GPU_ElementwiseOpEXTF : I32EnumAttrCase<"EXTF", 13, "extf">;
+
+def MMAElementWise : I32EnumAttr<"MMAElementwiseOp",
+  "elementwise operation to apply to mma matrix", [
+    GPU_ElementwiseOpAddF,
+    GPU_ElementwiseOpMulF,
+    GPU_ElementwiseOpSUBF,
+    GPU_ElementwiseOpMaxF,
+    GPU_ElementwiseOpMinF,
+    GPU_ElementwiseOpDivF,
+    GPU_ElementwiseOpAddI,
+    GPU_ElementwiseOpMulI,
+    GPU_ElementwiseOpSUBI,
+    GPU_ElementwiseOpDivS,
+    GPU_ElementwiseOpDivU,
+    GPU_ElementwiseOpNEGF,
+    GPU_ElementwiseOpNEGS,
+    GPU_ElementwiseOpEXTF
+  ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::gpu";
+}
+def MMAElementWiseAttr : EnumAttr<GPU_Dialect, MMAElementWise,
+                                  "mma_element_wise">;
+
+def GPU_Prune2To4SpMatFlag : I32EnumAttr<"Prune2To4SpMatFlag",
+  "pruning strategy for 2:4 sparse matrix",
+  [
+    I32EnumAttrCase<"NONE", 0>,
+    I32EnumAttrCase<"PRUNE_ONLY", 1>,
+    I32EnumAttrCase<"PRUNE_AND_CHECK", 2>,
+  ]> {
+    let genSpecializedAttr = 0;
+    let cppNamespace = GPU_Dialect.cppNamespace;
+}
+
+def GPU_Prune2To4SpMatFlagAttr : EnumAttr<GPU_Dialect, GPU_Prune2To4SpMatFlag,
+                                   "prune_2to4_spmat_flag">{
+  let defaultValue = "Prune2To4SpMatFlag::PRUNE_AND_CHECK";
+}
+
+// To avoid coupling this dialect with cusparse.h specifics, we hardcoded magic
+// literals in this enum. Note that this should be kept in sync with
+// cusparseOperation_t in cusparse.h:
+// typedef enum {
+// CUSPARSE_OPERATION_NON_TRANSPOSE       = 0,
+// CUSPARSE_OPERATION_TRANSPOSE           = 1,
+// CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE = 2
+// } cusparseOperation_t;
+// TODO: find a proper way to keep them in sync?
+def GPU_TransposeMode : I32EnumAttr<"TransposeMode",
+    "transpose mode of sparse matrix supported by sparse tensor ops",
+    [
+      I32EnumAttrCase<"NON_TRANSPOSE", 0>,
+      I32EnumAttrCase<"TRANSPOSE", 1>,
+      I32EnumAttrCase<"CONJUGATE_TRANSPOSE", 2>,
+    ]> {
+      let genSpecializedAttr = 0;
+      let cppNamespace = GPU_Dialect.cppNamespace;
+}
+
+def GPU_TransposeModeAttr : EnumAttr<GPU_Dialect, GPU_TransposeMode,
+                                   "mat_transpose_mode">{
+  let defaultValue = "TransposeMode::NON_TRANSPOSE";
+}
+
+def GPU_SpGEMMWorkEstimationOrComputeKind : I32EnumAttr<"SpGEMMWorkEstimationOrComputeKind",
+    "choose whether spgemm_work_estimation_or_compute does work estimation or compute",
+    [
+      I32EnumAttrCase<"WORK_ESTIMATION", 0>,
+      I32EnumAttrCase<"COMPUTE", 1>,
+    ]> {
+      let genSpecializedAttr = 0;
+      let cppNamespace = GPU_Dialect.cppNamespace;
+}
+
+def GPU_SpGEMMWorkEstimationOrComputeKindAttr : EnumAttr<GPU_Dialect,
+    GPU_SpGEMMWorkEstimationOrComputeKind,
+    "spgemm_work_estimation_or_compute_kind"> {}
+
+
 //===----------------------------------------------------------------------===//
 // GPU Types.
 //===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
index d08e7ceb9e6c69..4c5ae7c2e8c6c8 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
@@ -38,18 +38,6 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
 class GPU_Op<string mnemonic, list<Trait> traits = []> :
     Op<GPU_Dialect, mnemonic, traits>;
 
-def GPU_Dimension : I32EnumAttr<"Dimension",
-    "a dimension, either 'x', 'y', or 'z'",
-    [
-      I32EnumAttrCase<"x", 0>,
-      I32EnumAttrCase<"y", 1>,
-      I32EnumAttrCase<"z", 2>
-    ]>{
-  let genSpecializedAttr = 0;
-  let cppNamespace = "::mlir::gpu";
-}
-def GPU_DimensionAttr : EnumAttr<GPU_Dialect, GPU_Dimension, "dim">;
-
 class GPU_IndexOp<string mnemonic, list<Trait> traits = []> :
     GPU_Op<mnemonic, !listconcat(traits, [
         Pure,
@@ -1104,51 +1092,8 @@ def GPU_YieldOp : GPU_Op<"yield", [Pure, ReturnLike, Terminator]>,
   let assemblyFormat = "attr-dict ($values^ `:` type($values))?";
 }
 
-// These mirror the reduction combining kinds from the vector dialect.
-def GPU_AllReduceOpAdd : I32EnumAttrCase<"ADD", 0, "add">;
-def GPU_AllReduceOpMul : I32EnumAttrCase<"MUL", 1, "mul">;
-def GPU_AllReduceOpMinUI : I32EnumAttrCase<"MINUI", 2, "minui">;
-def GPU_AllReduceOpMinSI : I32EnumAttrCase<"MINSI", 3, "minsi">;
-// Follows the `arith.minnumf` semantics.
-def GPU_AllReduceOpMinnumF : I32EnumAttrCase<"MINNUMF", 4, "minnumf">;
-def GPU_AllReduceOpMaxUI : I32EnumAttrCase<"MAXUI", 5, "maxui">;
-def GPU_AllReduceOpMaxSI : I32EnumAttrCase<"MAXSI", 6, "maxsi">;
-// Follows the `arith.maxnumf` semantics.
-def GPU_AllReduceOpMaxnumF : I32EnumAttrCase<"MAXNUMF", 7, "maxnumf">;
-def GPU_AllReduceOpAnd : I32EnumAttrCase<"AND", 8, "and">;
-def GPU_AllReduceOpOr  : I32EnumAttrCase<"OR",  9, "or">;
-def GPU_AllReduceOpXor : I32EnumAttrCase<"XOR", 10, "xor">;
-// Follows the `arith.minimumf` semantics.
-def GPU_AllReduceOpMinimumF : I32EnumAttrCase<"MINIMUMF", 11, "minimumf">;
-// Follows the `arith.maximumf` semantics.
-def GPU_AllReduceOpMaximumF : I32EnumAttrCase<"MAXIMUMF", 12, "maximumf">;
-
-def GPU_AllReduceOperation : I32EnumAttr<"AllReduceOperation",
-    "built-in reduction operations supported by gpu.allreduce.",
-    [
-      GPU_AllReduceOpAdd,
-      GPU_AllReduceOpMul,
-      GPU_AllReduceOpMinUI,
-      GPU_AllReduceOpMinSI,
-      GPU_AllReduceOpMinnumF,
-      GPU_AllReduceOpMaxUI,
-      GPU_AllReduceOpMaxSI,
-      GPU_AllReduceOpMaxnumF,
-      GPU_AllReduceOpAnd,
-      GPU_AllReduceOpOr,
-      GPU_AllReduceOpXor,
-      GPU_AllReduceOpMinimumF,
-      GPU_AllReduceOpMaximumF
-    ]>{
-  let genSpecializedAttr = 0;
-  let cppNamespace = "::mlir::gpu";
-}
-
 def AnyIntegerOrFloat : AnyTypeOf<[AnySignlessInteger, AnyFloat], "Integer or Float">;
 
-def GPU_AllReduceOperationAttr : EnumAttr<GPU_Dialect, GPU_AllReduceOperation,
-                                          "all_reduce_op">;
-
 def GPU_AllReduceOp : GPU_Op<"all_reduce",
     [SameOperandsAndResultType, IsolatedFromAbove]> {
   let summary = "Reduce values among workgroup.";
@@ -1276,22 +1221,6 @@ def GPU_SubgroupReduceOp : GPU_Op<"subgroup_reduce", [SameOperandsAndResultType]
   let hasVerifier = 1;
 }
 
-def GPU_ShuffleOpXor  : I32EnumAttrCase<"XOR",  0, "xor">;
-def GPU_ShuffleOpDown : I32EnumAttrCase<"DOWN", 1, "down">;
-def GPU_ShuffleOpUp   : I32EnumAttrCase<"UP",   2, "up">;
-def GPU_ShuffleOpIdx  : I32EnumAttrCase<"IDX",  3, "idx">;
-
-def GPU_ShuffleMode : I32EnumAttr<"ShuffleMode",
-    "Indexing modes supported by gpu.shuffle.",
-    [
-      GPU_ShuffleOpXor, GPU_ShuffleOpUp, GPU_ShuffleOpDown, GPU_ShuffleOpIdx,
-    ]> {
-  let genSpecializedAttr = 0;
-  let cppNamespace = "::mlir::gpu";
-}
-def GPU_ShuffleModeAttr : EnumAttr<GPU_Dialect, GPU_ShuffleMode,
-                                   "shuffle_mode">;
-
 def GPU_ShuffleOp : GPU_Op<
     "shuffle", [Pure, AllTypesMatch<["value", "shuffleResult"]>]>,
     Arguments<(ins AnyIntegerOrFloatOr1DVector:$value, I32:$offset, I32:$width,
@@ -1914,44 +1843,6 @@ def GPU_SubgroupMmaConstantMatrixOp : GPU_Op<"subgroup_mma_constant_matrix",
   }];
 }
 
-def GPU_ElementwiseOpAddF  : I32EnumAttrCase<"ADDF", 0, "addf">;
-def GPU_ElementwiseOpMulF  : I32EnumAttrCase<"MULF", 1, "mulf">;
-def GPU_ElementwiseOpSUBF  : I32EnumAttrCase<"SUBF", 2, "subf">;
-def GPU_ElementwiseOpMaxF : I32EnumAttrCase<"MAXF", 3, "maxf">;
-def GPU_ElementwiseOpMinF : I32EnumAttrCase<"MINF", 4, "minf">;
-def GPU_ElementwiseOpDivF : I32EnumAttrCase<"DIVF", 5, "divf">;
-def GPU_ElementwiseOpAddI  : I32EnumAttrCase<"ADDI", 6, "addi">;
-def GPU_ElementwiseOpMulI  : I32EnumAttrCase<"MULI", 7, "muli">;
-def GPU_ElementwiseOpSUBI  : I32EnumAttrCase<"SUBI", 8, "subi">;
-def GPU_ElementwiseOpDivS : I32EnumAttrCase<"DIVS", 9, "divs">;
-def GPU_ElementwiseOpDivU : I32EnumAttrCase<"DIVU", 10, "divu">;
-def GPU_ElementwiseOpNEGF : I32EnumAttrCase<"NEGATEF", 11, "negatef">;
-def GPU_ElementwiseOpNEGS : I32EnumAttrCase<"NEGATES", 12, "negates">;
-def GPU_ElementwiseOpEXTF : I32EnumAttrCase<"EXTF", 13, "extf">;
-
-def MMAElementWise : I32EnumAttr<"MMAElementwiseOp",
-  "elementwise operation to apply to mma matrix", [
-    GPU_ElementwiseOpAddF,
-    GPU_ElementwiseOpMulF,
-    GPU_ElementwiseOpSUBF,
-    GPU_ElementwiseOpMaxF,
-    GPU_ElementwiseOpMinF,
-    GPU_ElementwiseOpDivF,
-    GPU_ElementwiseOpAddI,
-    GPU_ElementwiseOpMulI,
-    GPU_ElementwiseOpSUBI,
-    GPU_ElementwiseOpDivS,
-    GPU_ElementwiseOpDivU,
-    GPU_ElementwiseOpNEGF,
-    GPU_ElementwiseOpNEGS,
-    GPU_ElementwiseOpEXTF
-  ]> {
-  let genSpecializedAttr = 0;
-  let cppNamespace = "::mlir::gpu";
-}
-def MMAElementWiseAttr : EnumAttr<GPU_Dialect, MMAElementWise,
-                                  "mma_element_wise">;
-
 def GPU_SubgroupMmaElementwiseOp : GPU_Op<"subgroup_mma_elementwise",
     [Pure,
      AllTypesMatch<["args"]>]>{
@@ -2262,22 +2153,6 @@ def GPU_CreateBsrOp : GPU_Op<"create_bsr", [GPU_AsyncOpInterface]> {
   }];
 }
 
-def GPU_Prune2To4SpMatFlag : I32EnumAttr<"Prune2To4SpMatFlag",
-  "pruning strategy for 2:4 sparse matrix",
-  [
-    I32EnumAttrCase<"NONE", 0>,
-    I32EnumAttrCase<"PRUNE_ONLY", 1>,
-    I32EnumAttrCase<"PRUNE_AND_CHECK", 2>,
-  ]> {
-    let genSpecializedAttr = 0;
-    let cppNamespace = GPU_Dialect.cppNamespace;
-}
-
-def GPU_Prune2To4SpMatFlagAttr : EnumAttr<GPU_Dialect, GPU_Prune2To4SpMatFlag,
-                                   "prune_2to4_spmat_flag">{
-  let defaultValue = "Prune2To4SpMatFlag::PRUNE_AND_CHECK";
-}
-
 
 def GPU_Create2To4SpMatOp : GPU_Op<"create_2to4_spmat", [GPU_AsyncOpInterface]> {
   let summary = "Create sparse matrix with 2:4 sparsity operation";
@@ -2340,31 +2215,6 @@ def GPU_DestroySpMatOp : GPU_Op<"destroy_sp_mat", [GPU_AsyncOpInterface]> {
   }];
 }
 
-// To avoid coupling this dialect with cusparse.h specifics, we hardcoded magic
-// literals in this enum. Note that this should be kept in sync with
-// cusparseOperation_t in cusparse.h:
-// typedef enum {
-// CUSPARSE_OPERATION_NON_TRANSPOSE       = 0,
-// CUSPARSE_OPERATION_TRANSPOSE           = 1,
-// CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE = 2
-// } cusparseOperation_t;
-// TODO: find a proper way to keep them in sync?
-def GPU_TransposeMode : I32EnumAttr<"TransposeMode",
-    "transpose mode of sparse matrix supported by sparse tensor ops",
-    [
-      I32EnumAttrCase<"NON_TRANSPOSE", 0>,
-      I32EnumAttrCase<"TRANSPOSE", 1>,
-      I32EnumAttrCase<"CONJUGATE_TRANSPOSE", 2>,
-    ]> {
-      let genSpecializedAttr = 0;
-      let cppNamespace = GPU_Dialect.cppNamespace;
-}
-
-def GPU_TransposeModeAttr : EnumAttr<GPU_Dialect, GPU_TransposeMode,
-                                   "mat_transpose_mode">{
-  let defaultValue = "TransposeMode::NON_TRANSPOSE";
-}
-
 def GPU_SpMVBufferSizeOp : GPU_Op<"spmv_buffer_size", [GPU_AsyncOpInterface]> {
   let summary = "Precompute buffersize for SpMV operation";
   let description = [{
@@ -2677,20 +2527,6 @@ def GPU_SDDMMOp : GPU_Op<"sddmm", [GPU_AsyncOpInterface]> {
   }];
 }
 
-def GPU_SpGEMMWorkEstimationOrComputeKind : I32EnumAttr<"SpGEMMWorkEstimationOrComputeKind",
-    "choose whether spgemm_work_estimation_or_compute does work estimation or compute",
-    [
-      I32EnumAttrCase<"WORK_ESTIMATION", 0>,
-      I32EnumAttrCase<"COMPUTE", 1>,
-    ]> {
-      let genSpecializedAttr = 0;
-      let cppNamespace = GPU_Dialect.cppNamespace;
-}
-
-def GPU_SpGEMMWorkEstimationOrComputeKindAttr : EnumAttr<GPU_Dialect,
-    GPU_SpGEMMWorkEstimationOrComputeKind,
-    "spgemm_work_estimation_or_compute_kind"> {}
-
 def GPU_SpGEMMCreateDescrOp : GPU_Op<"spgemm_create_descr", [GPU_AsyncOpInterface]> {
   let summary = "SpGEMM Create Descr operation";
   let description = [{
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 23187f256455bb..7eaabae3be9ec1 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -60,7 +60,7 @@ declare_mlir_python_sources(MLIRPythonCAPI.HeaderSources
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
-  TD_FILE dialects/AffineOps.td
+  TD_FILE "${CMAKE_CURRENT_SOURCE_DIR}/include/mlir/dialects/Affine/IR/AffineOps.td"
   SOURCES
     dialects/affine.py
   DIALECT_NAME affine
@@ -320,7 +320,7 @@ declare_mlir_dialect_python_bindings(
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
-  TD_FILE dialects/NVGPUOps.td
+  TD_FILE "${CMAKE_CURRENT_SOURCE_DIR}/include/mlir/dialects/NVGPU/NVGPUOps.td"
   SOURCES
     dialects/nvgpu.py
   DIALECT_NAME nvgpu
diff --git a/mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp
index 3f660ae151c749..d815a4e62ca786 100644
--- a/mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp
+++ b/mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp
@@ -15,6 +15,7 @@
 #include "mlir/TableGen/AttrOrTypeDef.h"
 #include "mlir/TableGen/Attribute.h"
 #include "mlir/TableGen/Dialect.h"
+#include "llvm/TableGen/Error.h"
 #include "mlir/TableGen/GenInfo.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/TableGen/Record.h"
@@ -136,12 +137,18 @@ static bool emitPythonEnums(const RecordKeeper &records, raw_ostream &os) {
   os << fileHeader;
   for (const Record *it :
        records.getAllDerivedDefinitionsIfDefined("EnumAttrInfo")) {
+    if (llvm::SrcMgr.FindBufferContainingLoc(it->getLoc()[0]) !=
+        llvm::SrcMgr.getMainFileID())
+      continue;
     EnumAttr enumAttr(*it);
     emitEnumClass(enumAttr, os);
     emitAttributeBuilder(enumAttr, os);
   }
   for (const Record *it :
        records.getAllDerivedDefinitionsIfDefined("EnumAttr")) {
+    if (llvm::SrcMgr.FindBufferContainingLoc(it->getLoc()[0]) !=
+        llvm::SrcMgr.getMainFileID())
+      continue;
     AttrOrTypeDef attr(&*it);
     if (!attr.getMnemonic()) {
       llvm::errs() << "enum case " << attr
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 179fed2f5e9a00..2e3cccf8090f5b 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -1647,14 +1647,6 @@ gentbl_cc_library(
             ["-gen-dialect-defs"],
             "include/mlir/Dialect/EmitC/IR/EmitCDialect.cpp.inc",
         ),
-        (
-            ["-gen-enum-decls"],
-            "include/mlir/Dialect/EmitC/IR/EmitCEnums.h.inc",
-        ),
-        (
-            ["-gen-enum-defs"],
-            "include/mlir/Dialect/EmitC/IR/EmitCEnums.cpp.inc",
-        ),
         (
             ["-gen-op-decls"],
             "include/mlir/Dialect/EmitC/IR/EmitC.h.inc",
@@ -1677,6 +1669,23 @@ gentbl_cc_library(
     deps = [":EmitCTdFiles"],
 )
 
+gentbl_cc_library(
+    name = "EmitCEnumsIncGen",
+    tbl_outs = [
+        (
+            ["-gen-enum-decls"],
+            "include/mlir/Dialect/EmitC/IR/EmitCEnums.h.inc",
+        ),
+        (
+            ["-gen-enum-defs"],
+            "include/mlir/Dialect/EmitC/IR/EmitCEnums.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/EmitC/IR/EmitCAttributes.td",
+    deps = [":EmitCTdFiles"],
+)
+
 gentbl_cc_library(
     name = "EmitCPassIncGen",
     tbl_outs = [
@@ -3997,6 +4006,7 @@ cc_library(
         ":CastInterfaces",
         ":ControlFlowInterfaces",
         ":EmitCAttributesIncGen",
+        ":EmitCEnumsIncGen",
         ":EmitCOpsIncGen",
         ":FunctionInterfaces",
         ":IR",
@@ -5436,6 +5446,7 @@ cc_library(
         ":InferTypeOpInterface",
         ":InliningUtils",
         ":LLVMDialectInterfaceIncGen",
+        ":LLVMEnumsIncGen",
         ":LLVMIntrinsicOpsIncGen",
         ":LLVMOpsIncGen",
         ":LLVMTypesIncGen",
@@ -5598,14 +5609,6 @@ gentbl_cc_library(
             ["-gen-op-defs"],
             "include/mlir/Dialect/GPU/IR/GPUOps.cpp.inc",
         ),
-        (
-            ["-gen-enum-decls"],
-            "include/mlir/Dialect/GPU/IR/GPUOpsEnums.h.inc",
-        ),
-        (
-            ["-gen-enum-defs"],
-            "include/mlir/Dialect/GPU/IR/GPUOpsEnums.cpp.inc",
-        ),
         (
             ["-gen-attrdef-decls"],
             "include/mlir/Dialect/GPU/IR/GPUOpsAttributes.h.inc",
@@ -5623,6 +5626,26 @@ gentbl_cc_library(
     ],
 )
 
+gentbl_cc_library(
+    name = "GPUEnumsIncGen",
+    tbl_outs = [
+        (
+            ["-gen-enum-decls"],
+            "include/mlir/Dialect/GPU/IR/GPUOpsEnums.h.inc",
+        ),
+        (
+            ["-gen-enum-defs"],
+            "include/mlir/Dialect/GPU/IR/GPUOpsEnums.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/GPU/IR/GPUBase.td",
+    deps = [
+        ":DLTIDialectTdFiles",
+        ":GPUOpsTdFiles",
+    ],
+)
+
 cc_library(
     name = "GPUDialect",
     srcs = glob(
@@ -5641,6 +5664,7 @@ cc_library(
         ":FunctionInterfaces",
         ":GPUBaseIncGen",
         ":GPUCompilationAttrInterfacesIncGen",
+        ":GPUEnumsIncGen",
         ":GPUOpsIncGen",
         ":IR",
         ":InferIntRangeInterface",
@@ -6246,14 +6270,6 @@ gentbl_cc_library(
             ["-gen-dialect-defs"],
             "include/mlir/Dialect/LLVMIR/LLVMOpsDialect.cpp.inc",
         ),
-        (
-            ["-gen-enum-decls"],
-            "include/mlir/Dialect/LLVMIR/LLVMOpsEnums.h.inc",
-        ),
-        (
-            ["-gen-enum-defs"],
-            "include/mlir/Dialect/LLVMIR/LLVMOpsEnums.cpp.inc",
-        ),
         (
             [
                 "--gen-attrdef-decls",
@@ -6274,6 +6290,23 @@ gentbl_cc_library(
     deps = [":LLVMOpsTdFiles"],
 )
 
+gentbl_cc_library(
+    name = "LLVMEnumsIncGen",
+    tbl_outs = [
+        (
+            ["-gen-enum-decls"],
+            "include/mlir/Dialect/LLVMIR/LLVMOpsEnums.h.inc",
+        ),
+        (
+            ["-gen-enum-defs"],
+            "include/mlir/Dialect/LLVMIR/LLVMOpsEnums.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/LLVMIR/LLVMEnums.td",
+    deps = [":LLVMOpsTdFiles"],
+)
+
 gentbl_cc_library(
     name = "LLVMTypesIncGen",
     tbl_outs = [
@@ -7108,14 +7141,6 @@ gentbl_cc_library(
             ["-gen-op-doc"],
             "g3doc/Dialects/SPIRV/SPIRVOps.md",
         ),
-        (
-            ["-gen-enum-decls"],
-            "include/mlir/Dialect/SPIRV/IR/SPIRVEnums.h.inc",
-        ),
-        (
-            ["-gen-enum-defs"],
-            "include/mlir/Dialect/SPIRV/IR/SPIRVEnums.cpp.inc",
-        ),
         (
             ["-gen-spirv-enum-avail-decls"],
             "include/mlir/Dialect/SPIRV/IR/SPIRVEnumAvailability.h.inc",
@@ -7134,6 +7159,23 @@ gentbl_cc_library(
     deps = [":SPIRVOpsTdFiles"],
 )
 
+gentbl_cc_library(
+    name = "SPIRVEnumsIncGen",
+    tbl_outs = [
+        (
+            ["-gen-enum-decls"],
+            "include/mlir/Dialect/SPIRV/IR/SPIRVEnums.h.inc",
+        ),
+        (
+            ["-gen-enum-defs"],
+            "include/mlir/Dialect/SPIRV/IR/SPIRVEnums.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/SPIRV/IR/SPIRVBase.td",
+    deps = [":SPIRVOpsTdFiles"],
+)
+
 gentbl_cc_library(
     name = "SPIRVAttributesIncGen",
     tbl_outs = [
@@ -7237,6 +7279,7 @@ cc_library(
         ":SPIRVAttributesIncGen",
         ":SPIRVAvailabilityIncGen",
         ":SPIRVCanonicalizationIncGen",
+        ":SPIRVEnumsIncGen",
         ":SPIRVOpsIncGen",
         ":SideEffectInterfaces",
         ":Support",
@@ -10584,14 +10627,6 @@ gentbl_cc_library(
             ["-gen-op-defs"],
             "include/mlir/Dialect/OpenMP/OpenMPOps.cpp.inc",
         ),
-        (
-            ["-gen-enum-decls"],
-            "include/mlir/Dialect/OpenMP/OpenMPOpsEnums.h.inc",
-        ),
-        (
-            ["-gen-enum-defs"],
-            "include/mlir/Dialect/OpenMP/OpenMPOpsEnums.cpp.inc",
-        ),
         (
             [
                 "-gen-dialect-decls",
@@ -10634,6 +10669,23 @@ gentbl_cc_library(
     deps = [":OpenMPOpsTdFiles"],
 )
 
+gentbl_cc_library(
+    name = "OpenMPEnumsIncGen",
+    tbl_outs = [
+        (
+            ["-gen-enum-decls"],
+            "include/mlir/Dialect/OpenMP/OpenMPOpsEnums.h.inc",
+        ),
+        (
+            ["-gen-enum-defs"],
+            "include/mlir/Dialect/OpenMP/OpenMPOpsEnums.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/OpenMP/OpenMPEnums.td",
+    deps = [":OpenMPOpsTdFiles"],
+)
+
 gentbl_cc_library(
     name = "OpenMPTypeInterfacesIncGen",
     tbl_outs = [
@@ -10690,6 +10742,7 @@ cc_library(
         ":IR",
         ":LLVMDialect",
         ":OpenACCMPOpsInterfacesIncGen",
+        ":OpenMPEnumsIncGen",
         ":OpenMPInterfacesIncGen",
         ":OpenMPOpsIncGen",
         ":OpenMPTypeInterfacesIncGen",
diff --git a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
index 254cab0db4a5d6..f0000766aa3590 100644
--- a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
@@ -124,7 +124,7 @@ gentbl_filegroup(
         ),
     ],
     tblgen = "//mlir:mlir-tblgen",
-    td_file = "mlir/dialects/AffineOps.td",
+    td_file = "//mlir:include/mlir/Dialect/Affine/IR/AffineOps.td",
     deps = [
         ":AffineOpsPyTdFiles",
     ],
@@ -821,7 +821,7 @@ gentbl_filegroup(
         ),
     ],
     tblgen = "//mlir:mlir-tblgen",
-    td_file = "mlir/dialects/GPUOps.td",
+    td_file = "//mlir:include/mlir/Dialect/GPU/IR/GPUBase.td",
     deps = [
         ":GPUOpsPyTdFiles",
     ],
@@ -877,7 +877,7 @@ gentbl_filegroup(
         ),
     ],
     tblgen = "//mlir:mlir-tblgen",
-    td_file = "mlir/dialects/NVGPUOps.td",
+    td_file = "//mlir:include/mlir/Dialect/NVGPU/IR/NVGPU.td",
     deps = [
         ":NVGPUOpsPyTdFiles",
     ],



More information about the llvm-commits mailing list