iree_gpu Python bindings (GPUPipelineOptionsAttr)

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,42 @@
+// 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"
+#include "mlir/CAPI/Wrap.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+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,73 @@
+// 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";
+
+  //===-------------------------------------------------------------------===//
+  // GPUPipelineOptionsAttr
+  //===-------------------------------------------------------------------===//
+
+  mlir_attribute_subclass(m, "GPUPipelineOptionsAttr",
+                          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 import *

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

@@ -7,10 +7,76 @@
 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
+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(
+                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.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<>"
+
+            gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get(True)
+            assert (
+                str(gpu_attr)
+                == "#iree_gpu.pipeline_options<prefetch_shared_memory = true>"
+            )
+
+            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>"
+            )
+
+            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>"
+            )
+            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 (
+                str(gpu_attr)
+                == "#iree_gpu.pipeline_options<reorder_workgroups_strategy = <Swizzle>>"
+            )

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

@@ -87,6 +87,7 @@ iree_compiler_cc_library(
         "//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",

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

@@ -67,6 +67,7 @@ iree_cc_library(
     iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect
     iree::compiler::Codegen::Utils::VectorOpUtils
     iree::compiler::Dialect::LinalgExt::IR
+    iree::compiler::bindings::c::headers
   PUBLIC
 )
 

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

@@ -14,12 +14,16 @@
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h"
 #include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
 #include "iree/compiler/Codegen/Utils/VectorOpUtils.h"
+#include "iree/compiler/dialects/iree_gpu.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/STLForwardCompat.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "mlir-c/IR.h"
+#include "mlir/CAPI/IR.h"
+#include "mlir/CAPI/Support.h"
 #include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
@@ -1687,3 +1691,64 @@ void IREEGPUDialect::registerAttributes() {
 }
 
 } // namespace mlir::iree_compiler::IREE::GPU
+
+bool ireeAttributeIsAGPUPipelineOptionsAttr(MlirAttribute attr) {
+  return llvm::isa<mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr>(
+      unwrap(attr));
+}
+
+MlirAttribute
+ireeGPUPipelineOptionsAttrGet(MlirContext mlirCtx, bool *prefetchSharedMemory,
+                              bool *noReduceSharedMemoryBankConflicts,
+                              MlirAttribute *reorderWorkgroupsStrategy) {
+  mlir::MLIRContext *ctx = unwrap(mlirCtx);
+  mlir::Builder b(ctx);
+  auto prefetchSharedMemoryAttr = mlir::BoolAttr();
+  if (prefetchSharedMemory) {
+    prefetchSharedMemoryAttr = b.getBoolAttr(*prefetchSharedMemory);
+  }
+  auto noReduceSharedMemoryBankConflictsAttr = mlir::BoolAttr();
+  if (noReduceSharedMemoryBankConflicts) {
+    noReduceSharedMemoryBankConflictsAttr =
+        b.getBoolAttr(*noReduceSharedMemoryBankConflicts);
+  }
+  auto strategyAttr =
+      mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategyAttr();
+  if (reorderWorkgroupsStrategy) {
+    strategyAttr = llvm::dyn_cast<
+        mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategyAttr>(
+        unwrap(*reorderWorkgroupsStrategy));
+  }
+  return wrap(mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr::get(
+      ctx, prefetchSharedMemoryAttr, noReduceSharedMemoryBankConflictsAttr,
+      strategyAttr));
+}
+
+MlirAttribute
+ireeGPUPipelineOptionsAttrGetPrefetchSharedMemory(MlirAttribute attr) {
+  auto gpuAttr =
+      llvm::cast<mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr>(
+          unwrap(attr));
+  return wrap(gpuAttr.getPrefetchSharedMemory());
+}
+
+MlirAttribute ireeGPUPipelineOptionsAttrGetNoReduceSharedMemoryBankConflicts(
+    MlirAttribute attr) {
+  auto gpuAttr =
+      llvm::cast<mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr>(
+          unwrap(attr));
+  return wrap(gpuAttr.getNoReduceSharedMemoryBankConflicts());
+}
+
+MlirAttribute
+ireeGPUPipelineOptionsAttrGetReorderWorkgroupsStrategy(MlirAttribute attr) {
+  auto gpuAttr =
+      llvm::cast<mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr>(
+          unwrap(attr));
+  return wrap(gpuAttr.getReorderWorkgroupsStrategy());
+}
+
+MlirTypeID ireeGPUPipelineOptionsAttrGetTypeID() {
+  return wrap(
+      mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr::getTypeID());
+}

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

@@ -481,8 +481,8 @@ def IREEGPU_LaneIdAttr : AttrDef<IREEGPU_Dialect, "LaneId", [
 //===----------------------------------------------------------------------===//
 
 def IREEGPU_ReorderWorkgroupsStrategyAttr :
-    EnumAttr<IREEGPU_Dialect, IREEGPU_ReorderWorkgroupsStrategy, ""> {
-  let assemblyFormat = "``$value";
+    EnumAttr<IREEGPU_Dialect, IREEGPU_ReorderWorkgroupsStrategy, "reorder_workgroups_strategy"> {
+  let assemblyFormat = "`<` $value `>`";
   let cppNamespace = "::mlir::iree_compiler::IREE::GPU";
 }