Reland #18804

compiler/bindings/c/BUILD.bazel

@@ -14,6 +14,7 @@ cc_library(
     name = "headers",
     hdrs = [
         "iree/compiler/api_support.h",
+        "iree/compiler/dialects/iree_gpu.h",
         "iree/compiler/embedding_api.h",
         "iree/compiler/loader.h",
         "iree/compiler/mlir_interop.h",

compiler/bindings/c/CMakeLists.txt

@@ -22,6 +22,7 @@ iree_cc_library(
     headers
   HDRS
     "iree/compiler/api_support.h"
+    "iree/compiler/dialects/iree_gpu.h"
     "iree/compiler/embedding_api.h"
     "iree/compiler/loader.h"
     "iree/compiler/mlir_interop.h"

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

@@ -0,0 +1,59 @@
+// Copyright 2024 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_COMPILER_DIALECTS_IREE_GPU_H
+#define IREE_COMPILER_DIALECTS_IREE_GPU_H
+
+#include "mlir-c/IR.h"
+#include "mlir-c/Support.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+enum ireeGPUReorderWorkgroupsStrategyEnum {
+  ireeGPUReorderWorkgroupsStrategyEnumNone = 0,
+  ireeGPUReorderWorkgroupsStrategyEnumSwizzle = 1,
+  ireeGPUReorderWorkgroupsStrategyEnumTranspose = 2,
+};
+
+MLIR_CAPI_EXPORTED bool
+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,
+                              bool *noReduceSharedMemoryBankConflicts,
+                              MlirAttribute *reorderWorkgroupsStrategy);
+
+MLIR_CAPI_EXPORTED MlirAttribute
+ireeGPUPipelineOptionsAttrGetPrefetchSharedMemory(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED MlirAttribute
+ireeGPUPipelineOptionsAttrGetNoReduceSharedMemoryBankConflicts(
+    MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED MlirAttribute
+ireeGPUPipelineOptionsAttrGetReorderWorkgroupsStrategy(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED MlirTypeID ireeGPUPipelineOptionsAttrGetTypeID(void);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // IREE_COMPILER_DIALECTS_IREE_GPU_H

compiler/bindings/python/CMakeLists.txt

@@ -100,6 +100,15 @@ declare_mlir_dialect_python_bindings(
   DIALECT_NAME vm
 )
 
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT IREEPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler"
+  TD_FILE dialects/IREEGPUBinding.td
+  GEN_ENUM_BINDINGS
+  SOURCES dialects/iree_gpu.py
+  DIALECT_NAME iree_gpu
+)
+
 declare_mlir_python_sources(IREECompilerAPIPythonCore
   ADD_TO_PARENT IREEPythonSources
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler"
@@ -165,6 +174,17 @@ declare_mlir_python_extension(IREECompilerPythonExtensions.Registration
     LLVMSupport
 )
 
+declare_mlir_python_extension(IREECompilerPythonExtensions.CompilerDialects
+  MODULE_NAME _ireeCompilerDialects
+  ADD_TO_PARENT IREECompilerPythonExtensions
+  SOURCES
+    IREECompilerDialectsModule.cpp
+  EMBED_CAPI_LINK_LIBS
+    iree_compiler_API_SharedImpl
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+)
+
 ################################################################################
 # Generate packages and shared library
 ################################################################################

compiler/bindings/python/IREECompilerDialectsModule.cpp

@@ -0,0 +1,123 @@
+// Copyright 2024 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/dialects/iree_gpu.h"
+#include "mlir-c/BuiltinAttributes.h"
+#include "mlir-c/IR.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
+
+namespace py = pybind11;
+using namespace mlir::python::adaptors;
+
+PYBIND11_MODULE(_ireeCompilerDialects, m) {
+  m.doc() = "iree-compiler dialects python extension";
+
+  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
+  //===-------------------------------------------------------------------===//
+
+  mlir_attribute_subclass(iree_gpu_module, "PipelineOptionsAttr",
+                          ireeAttributeIsAGPUPipelineOptionsAttr,
+                          ireeGPUPipelineOptionsAttrGetTypeID)
+      .def_classmethod(
+          "get",
+          [](const py::object &, std::optional<bool> prefetchSharedMemory,
+             std::optional<bool> noReduceSharedMemoryBankConflicts,
+             std::optional<MlirAttribute> reorderWorkgroupsStrategy,
+             MlirContext ctx) {
+            return ireeGPUPipelineOptionsAttrGet(
+                ctx,
+                prefetchSharedMemory.has_value() ? &*prefetchSharedMemory
+                                                 : nullptr,
+                noReduceSharedMemoryBankConflicts.has_value()
+                    ? &*noReduceSharedMemoryBankConflicts
+                    : nullptr,
+                reorderWorkgroupsStrategy.has_value()
+                    ? &*reorderWorkgroupsStrategy
+                    : nullptr);
+          },
+          "cls"_a, "prefetch_shared_memory"_a = py::none(),
+          "no_reduce_shared_memory_bank_conflicts"_a = py::none(),
+          "reorder_workgroups_strategy"_a = py::none(), py::kw_only(),
+          "ctx"_a = py::none(), "Gets a gpu.pipeline_options from parameters.")
+      .def_property_readonly(
+          "prefetch_shared_memory",
+          [](MlirAttribute self) -> std::optional<bool> {
+            auto attr = ireeGPUPipelineOptionsAttrGetPrefetchSharedMemory(self);
+            if (!mlirAttributeIsNull(attr))
+              return mlirBoolAttrGetValue(attr);
+            return std::nullopt;
+          })
+      .def_property_readonly(
+          "no_reduce_shared_memory_bank_conflicts",
+          [](MlirAttribute self) -> std::optional<bool> {
+            auto attr =
+                ireeGPUPipelineOptionsAttrGetNoReduceSharedMemoryBankConflicts(
+                    self);
+            if (!mlirAttributeIsNull(attr))
+              return mlirBoolAttrGetValue(attr);
+            return std::nullopt;
+          })
+      .def_property_readonly(
+          "reorder_workgroups_strategy",
+          [](MlirAttribute self) -> std::optional<MlirAttribute> {
+            auto attr =
+                ireeGPUPipelineOptionsAttrGetReorderWorkgroupsStrategy(self);
+            if (!mlirAttributeIsNull(attr))
+              return attr;
+            return std::nullopt;
+          });
+}

compiler/bindings/python/iree/compiler/dialects/IREEGPUBinding.td

@@ -0,0 +1,12 @@
+// Copyright 2024 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef PYTHON_BINDINGS_IREEGPU_OPS
+#define PYTHON_BINDINGS_IREEGPU_OPS
+
+include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.td"
+
+#endif // PYTHON_BINDINGS_IREEGPU_OPS

compiler/bindings/python/iree/compiler/dialects/iree_gpu.py

@@ -0,0 +1,9 @@
+# Copyright 2024 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._iree_gpu_ops_gen import *
+from ._iree_gpu_enum_gen import *
+from .._mlir_libs._ireeCompilerDialects.iree_gpu import *

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

@@ -7,10 +7,69 @@
 from iree.compiler import ir
 
 # Make sure that our dialects import.
-from iree.compiler.dialects import (
-    flow,
-    hal,
-    stream,
-    vm,
-    util,
-)
+from iree.compiler.dialects import flow, hal, stream, vm, util, iree_gpu
+
+
+@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 = iree_gpu.ReorderWorkgroupsStrategyAttr.get(
+                iree_gpu.ReorderWorkgroupsStrategy.Swizzle, ctx
+            )
+            gpu_attr = iree_gpu.PipelineOptionsAttr.get(
+                True,
+                False,
+                reorder_attr,
+            )
+            assert type(gpu_attr) is iree_gpu.PipelineOptionsAttr
+            assert gpu_attr.prefetch_shared_memory
+            assert not gpu_attr.no_reduce_shared_memory_bank_conflicts
+
+            gpu_attr = iree_gpu.PipelineOptionsAttr.get(
+                False,
+                True,
+                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.PipelineOptionsAttr.get()
+            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.PipelineOptionsAttr.get(True)
+            assert gpu_attr.prefetch_shared_memory
+            assert (
+                gpu_attr.no_reduce_shared_memory_bank_conflicts is None
+                and gpu_attr.reorder_workgroups_strategy is None
+            )
+
+            gpu_attr = iree_gpu.PipelineOptionsAttr.get(True, False)
+            assert gpu_attr.reorder_workgroups_strategy is None
+
+            gpu_attr = iree_gpu.PipelineOptionsAttr.get(
+                no_reduce_shared_memory_bank_conflicts=False
+            )
+            assert (
+                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.PipelineOptionsAttr.get(
+                reorder_workgroups_strategy=reorder_attr
+            )
+            assert gpu_attr.reorder_workgroups_strategy is not None
+            assert (
+                gpu_attr.reorder_workgroups_strategy.value
+                # unfortunately not `is`
+                == iree_gpu.ReorderWorkgroupsStrategy.Swizzle
+            )

compiler/src/iree/compiler/API/CMakeLists.txt

@@ -78,6 +78,7 @@ set(_EXPORT_OBJECT_LIBS
   obj.MLIRCAPITransformDialectTransforms
   iree_compiler_API_Internal_CompilerDriver.objects
   iree_compiler_API_Internal_IREECompileToolEntryPoint.objects
+  iree_compiler_API_Internal_IREEGPUDialectCAPI.objects
   iree_compiler_API_Internal_IREEMLIRLSPServerToolEntryPoint.objects
   iree_compiler_API_Internal_IREEOptToolEntryPoint.objects
   iree_compiler_API_Internal_IREEReduceToolEntryPoint.objects

compiler/src/iree/compiler/API/Internal/BUILD.bazel

@@ -125,3 +125,16 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:Support",
     ],
 )
+
+iree_compiler_cc_library(
+    name = "IREEGPUDialectCAPI",
+    srcs = [
+        "IREEGPUDialectCAPI.cpp",
+    ],
+    deps = [
+        "//compiler/bindings/c:headers",
+        "//compiler/src/iree/compiler/Codegen/Dialect/GPU/IR:IREEGPUDialect",
+        "@llvm-project//mlir:CAPIIR",
+        "@llvm-project//mlir:CAPIIRHeaders",
+    ],
+)

compiler/src/iree/compiler/API/Internal/CMakeLists.txt

@@ -103,6 +103,19 @@ iree_cc_library(
   PUBLIC
 )
 
+iree_cc_library(
+  NAME
+    IREEGPUDialectCAPI
+  SRCS
+    "IREEGPUDialectCAPI.cpp"
+  DEPS
+    IREELLVMIncludeSetup
+    MLIRCAPIIR
+    iree::compiler::Codegen::Dialect::GPU::IR::IREEGPUDialect
+    iree::compiler::bindings::c::headers
+  PUBLIC
+)
+
 ### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
 
 set(_lld_copts)