map ReorderWorkgroupsStrategy correctly

Signed-off-by: makslevental <maksim.levental@gmail.com>

compiler/bindings/c/iree/compiler/dialects/iree_gpu.h

@@ -9,14 +9,31 @@
 
 #include "mlir-c/IR.h"
 #include "mlir-c/Support.h"
-#include "mlir/CAPI/Wrap.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
+enum ireeGPUReorderWorkgroupsStrategyEnum {
+  ireeGPUReorderWorkgroupsStrategyEnumNone = 0,
+  ireeGPUReorderWorkgroupsStrategyEnumSwizzle = 1,
+  ireeGPUReorderWorkgroupsStrategyEnumTranspose = 2,
+};
+
 MLIR_CAPI_EXPORTED bool
-ireeAttributeIsAGPUPipelineOptionsAttr(MlirAttribute attr);
+ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED MlirTypeID
+ireeGPUReorderWorkgroupsStrategyAttrGetTypeID(void);
+
+MLIR_CAPI_EXPORTED MlirAttribute ireeGPUReorderWorkgroupsStrategyAttrGet(
+    MlirContext mlirCtx, ireeGPUReorderWorkgroupsStrategyEnum value);
+
+MLIR_CAPI_EXPORTED ireeGPUReorderWorkgroupsStrategyEnum
+ireeGPUReorderWorkgroupsStrategyAttrGetValue(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED
+bool ireeAttributeIsAGPUPipelineOptionsAttr(MlirAttribute attr);
 
 MLIR_CAPI_EXPORTED MlirAttribute
 ireeGPUPipelineOptionsAttrGet(MlirContext mlirCtx, bool *prefetchSharedMemory,

compiler/bindings/python/IREECompilerDialectsModule.cpp

@@ -18,6 +18,53 @@ PYBIND11_MODULE(_ireeCompilerDialects, m) {
   auto iree_gpu_module =
       m.def_submodule("iree_gpu", "iree_gpu dialect bindings");
 
+  //===-------------------------------------------------------------------===//
+  // GPUReorderWorkgroupsStrategyAttr
+  //===-------------------------------------------------------------------===//
+
+  auto strategyEnum =
+      py::enum_<ireeGPUReorderWorkgroupsStrategyEnum>(
+          iree_gpu_module, "ReorderWorkgroupsStrategy", py::module_local())
+          .value("None_", ireeGPUReorderWorkgroupsStrategyEnumNone)
+          .value("Swizzle", ireeGPUReorderWorkgroupsStrategyEnumSwizzle)
+          .value("Transpose", ireeGPUReorderWorkgroupsStrategyEnumTranspose)
+          .def(
+              "__str__",
+              [](ireeGPUReorderWorkgroupsStrategyEnum &self) {
+                switch (self) {
+                case ireeGPUReorderWorkgroupsStrategyEnumNone:
+                  return "None";
+                case ireeGPUReorderWorkgroupsStrategyEnumSwizzle:
+                  return "Swizzle";
+                case ireeGPUReorderWorkgroupsStrategyEnumTranspose:
+                  return "Transpose";
+                default:
+                  llvm::report_fatal_error(
+                      "unknown ReorderWorkgroupsStrategy variant");
+                }
+              },
+              // pybind overloads are tried in the order they were registered.
+              // As a result, enums used the default __str__ method instead of
+              // the custom one. Adding py::prepend() fixes this issue.
+              py::prepend());
+
+  mlir_attribute_subclass(iree_gpu_module, "ReorderWorkgroupsStrategyAttr",
+                          ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr,
+                          ireeGPUReorderWorkgroupsStrategyAttrGetTypeID)
+      .def_classmethod(
+          "get",
+          [](const py::object &, ireeGPUReorderWorkgroupsStrategyEnum value,
+             MlirContext ctx) {
+            return ireeGPUReorderWorkgroupsStrategyAttrGet(ctx, value);
+          },
+          "cls"_a, "value"_a, "ctx"_a = py::none(),
+          "Gets a gpu.reorder_workgroups_strategy from parameters.")
+      .def_property_readonly(
+          "value",
+          [](MlirAttribute self) -> ireeGPUReorderWorkgroupsStrategyEnum {
+            return ireeGPUReorderWorkgroupsStrategyAttrGetValue(self);
+          });
+
   //===-------------------------------------------------------------------===//
   // GPUPipelineOptionsAttr
   //===-------------------------------------------------------------------===//

compiler/bindings/python/test/ir/dialects_test.py

@@ -8,75 +8,68 @@
 
 # Make sure that our dialects import.
 from iree.compiler.dialects import flow, hal, stream, vm, util, iree_gpu
-from iree.compiler.dialects._iree_gpu_enum_gen import (
-    _ireegpu_reorderworkgroupsstrategyattr as ReorderWorkgroupsStrategyAttr,
-)
 
 
 @lambda _: _()
 def gpu_pipeline_options_attr():
     with ir.Context() as ctx, ir.Location.unknown():
         module = ir.Module.create()
         with ir.InsertionPoint(module.body):
-            reorder_attr = ReorderWorkgroupsStrategyAttr(
+            reorder_attr = iree_gpu.ReorderWorkgroupsStrategyAttr.get(
                 iree_gpu.ReorderWorkgroupsStrategy.Swizzle, ctx
             )
             gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get(
                 True,
                 False,
                 reorder_attr,
             )
-            assert (
-                str(gpu_attr.reorder_workgroups_strategy)
-                == "#iree_gpu.reorder_workgroups_strategy<Swizzle>"
-            )
-            assert (
-                str(gpu_attr)
-                == "#iree_gpu.pipeline_options<prefetch_shared_memory = true, no_reduce_shared_memory_bank_conflicts = false, reorder_workgroups_strategy = <Swizzle>>"
-            )
             assert type(gpu_attr) is iree_gpu.GPUPipelineOptionsAttr
             assert gpu_attr.prefetch_shared_memory
             assert not gpu_attr.no_reduce_shared_memory_bank_conflicts
 
             gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get(
                 False,
                 True,
-                ReorderWorkgroupsStrategyAttr(
+                iree_gpu.ReorderWorkgroupsStrategyAttr.get(
                     iree_gpu.ReorderWorkgroupsStrategy.Transpose, ctx
                 ),
             )
             assert not gpu_attr.prefetch_shared_memory
             assert gpu_attr.no_reduce_shared_memory_bank_conflicts
 
             gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get()
-            assert str(gpu_attr) == "#iree_gpu.pipeline_options<>"
+            assert (
+                gpu_attr.prefetch_shared_memory is None
+                and gpu_attr.no_reduce_shared_memory_bank_conflicts is None
+                and gpu_attr.reorder_workgroups_strategy is None
+            )
 
             gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get(True)
+            assert gpu_attr.prefetch_shared_memory
             assert (
-                str(gpu_attr)
-                == "#iree_gpu.pipeline_options<prefetch_shared_memory = true>"
+                gpu_attr.no_reduce_shared_memory_bank_conflicts is None
+                and gpu_attr.reorder_workgroups_strategy is None
             )
 
             gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get(True, False)
-            assert (
-                str(gpu_attr)
-                == "#iree_gpu.pipeline_options<prefetch_shared_memory = true, no_reduce_shared_memory_bank_conflicts = false>"
-            )
+            assert gpu_attr.reorder_workgroups_strategy is None
 
             gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get(
                 no_reduce_shared_memory_bank_conflicts=False
             )
             assert (
-                str(gpu_attr)
-                == "#iree_gpu.pipeline_options<no_reduce_shared_memory_bank_conflicts = false>"
+                gpu_attr.no_reduce_shared_memory_bank_conflicts is not None
+                and not gpu_attr.no_reduce_shared_memory_bank_conflicts
             )
             assert gpu_attr.prefetch_shared_memory is None
             assert gpu_attr.reorder_workgroups_strategy is None
 
             gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get(
                 reorder_workgroups_strategy=reorder_attr
             )
+            assert gpu_attr.reorder_workgroups_strategy is not None
             assert (
-                str(gpu_attr)
-                == "#iree_gpu.pipeline_options<reorder_workgroups_strategy = <Swizzle>>"
+                gpu_attr.reorder_workgroups_strategy.value
+                # unfortunately not `is`
+                == iree_gpu.ReorderWorkgroupsStrategy.Swizzle
             )

compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel

@@ -82,17 +82,18 @@ iree_compiler_cc_library(
         ":IREEGPUEnums",
         ":IREEGPUInterfaces",
         ":IREEGPUOpsGen",
+        "//compiler/bindings/c:headers",
         "//compiler/src/iree/compiler/Codegen/Common:TileSwizzle",
         "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
         "//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect",
         "//compiler/src/iree/compiler/Codegen/Utils:VectorOpUtils",
         "//compiler/src/iree/compiler/Dialect/LinalgExt/IR",
-        "//compiler/src/iree/compiler/bindings/c:headers",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AMDGPUDialect",
         "@llvm-project//mlir:AffineDialect",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:ArithUtils",
+        "@llvm-project//mlir:CAPIIR",
         "@llvm-project//mlir:ControlFlowInterfaces",
         "@llvm-project//mlir:DialectUtils",
         "@llvm-project//mlir:IR",

compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt

@@ -51,6 +51,7 @@ iree_cc_library(
     MLIRAffineDialect
     MLIRArithDialect
     MLIRArithUtils
+    MLIRCAPIIR
     MLIRControlFlowInterfaces
     MLIRIR
     MLIRLinalgDialect

compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp

@@ -37,7 +37,6 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/OpDefinition.h"
-#include "mlir/IR/TypeUtilities.h"
 
 #define DEBUG_TYPE "iree-gpu-attrs"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
@@ -1752,3 +1751,51 @@ MlirTypeID ireeGPUPipelineOptionsAttrGetTypeID() {
   return wrap(
       mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr::getTypeID());
 }
+
+static_assert(
+    static_cast<uint32_t>(ireeGPUReorderWorkgroupsStrategyEnumNone) ==
+            static_cast<uint32_t>(mlir::iree_compiler::IREE::GPU::
+                                      ReorderWorkgroupsStrategy::None) &&
+        static_cast<uint32_t>(ireeGPUReorderWorkgroupsStrategyEnumSwizzle) ==
+            static_cast<uint32_t>(mlir::iree_compiler::IREE::GPU::
+                                      ReorderWorkgroupsStrategy::Swizzle) &&
+        static_cast<uint32_t>(ireeGPUReorderWorkgroupsStrategyEnumTranspose) ==
+            static_cast<uint32_t>(mlir::iree_compiler::IREE::GPU::
+                                      ReorderWorkgroupsStrategy::Transpose) &&
+        static_cast<uint32_t>(ireeGPUReorderWorkgroupsStrategyEnumTranspose) ==
+            mlir::iree_compiler::IREE::GPU::
+                getMaxEnumValForReorderWorkgroupsStrategy(),
+    "ireeGPUReorderWorkgroupsStrategyEnum and "
+    "mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategy definitions "
+    "have diverged");
+
+bool ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr(MlirAttribute attr) {
+  return llvm::isa<
+      mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategyAttr>(
+      unwrap(attr));
+}
+
+MlirTypeID ireeGPUReorderWorkgroupsStrategyAttrGetTypeID() {
+  return wrap(mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategyAttr::
+                  getTypeID());
+}
+
+MlirAttribute ireeGPUReorderWorkgroupsStrategyAttrGet(
+    MlirContext mlirCtx, ireeGPUReorderWorkgroupsStrategyEnum value) {
+  mlir::MLIRContext *ctx = unwrap(mlirCtx);
+  return wrap(
+      mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategyAttr::get(
+          ctx, static_cast<
+                   mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategy>(
+                   value)));
+}
+
+ireeGPUReorderWorkgroupsStrategyEnum
+ireeGPUReorderWorkgroupsStrategyAttrGetValue(MlirAttribute attr) {
+  assert(ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr(attr) &&
+         "attr is not a GPUReorderWorkgroupsStrategyAttr");
+  return static_cast<ireeGPUReorderWorkgroupsStrategyEnum>(
+      llvm::cast<mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategyAttr>(
+          unwrap(attr))
+          .getValue());
+}