[WIP] iree_gpu Python bindings

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,37 @@
+// 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
+ireeGPUObjectAttrGet(MlirContext mlirCtx, bool prefetchSharedMemory,
+                     bool noReduceSharedMemoryBankConflicts);
+
+MLIR_CAPI_EXPORTED bool
+ireeGPUObjectAttrGetPrefetchSharedMemory(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED bool
+ireeGPUObjectAttrGetNoReduceSharedMemoryBankConflicts(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED MlirTypeID ireeGPUObjectAttrGetTypeID(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,43 @@
+// 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/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,
+                          ireeGPUObjectAttrGetTypeID)
+      .def_classmethod(
+          "get",
+          [](const py::object &, bool prefetchSharedMemory,
+             bool noReduceSharedMemoryBankConflicts, MlirContext ctx) {
+            return ireeGPUObjectAttrGet(ctx, prefetchSharedMemory,
+                                        noReduceSharedMemoryBankConflicts);
+          },
+          "cls"_a, "prefetch_shared_memory"_a,
+          "no_reduce_shared_memory_bank_conflicts"_a, "ctx"_a = py::none(),
+          "Gets a gpu.pipeline_options from parameters.")
+      .def_property_readonly("prefetch_shared_memory",
+                             [](MlirAttribute self) {
+                               return ireeGPUObjectAttrGetPrefetchSharedMemory(
+                                   self);
+                             })
+      .def_property_readonly(
+          "no_reduce_shared_memory_bank_conflicts", [](MlirAttribute self) {
+            return ireeGPUObjectAttrGetNoReduceSharedMemoryBankConflicts(self);
+          });
+}

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,23 @@
 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
+
+with ir.Context() as ctx, ir.Location.unknown():
+    module = ir.Module.create()
+    with ir.InsertionPoint(module.body):
+        gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get(True, False)
+        assert (
+            repr(gpu_attr)
+            == "GPUPipelineOptionsAttr(#iree_gpu.pipeline_options<prefetch_shared_memory = true, no_reduce_shared_memory_bank_conflicts = false>)"
+        )
+        assert (
+            repr(type(gpu_attr))
+            == "<class 'importlib._bootstrap.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)
+        assert not gpu_attr.prefetch_shared_memory
+        assert gpu_attr.no_reduce_shared_memory_bank_conflicts

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,35 @@ 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 ireeGPUObjectAttrGet(MlirContext mlirCtx,
+                                   bool prefetchSharedMemory,
+                                   bool noReduceSharedMemoryBankConflicts) {
+  mlir::MLIRContext *ctx = unwrap(mlirCtx);
+  return wrap(mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr::get(
+      ctx, prefetchSharedMemory, noReduceSharedMemoryBankConflicts));
+}
+
+bool ireeGPUObjectAttrGetPrefetchSharedMemory(MlirAttribute attr) {
+  auto gpuAttr =
+      llvm::cast<mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr>(
+          unwrap(attr));
+  return gpuAttr.getPrefetchSharedMemory().getValue();
+}
+
+bool ireeGPUObjectAttrGetNoReduceSharedMemoryBankConflicts(MlirAttribute attr) {
+  auto gpuAttr =
+      llvm::cast<mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr>(
+          unwrap(attr));
+  return gpuAttr.getNoReduceSharedMemoryBankConflicts().getValue();
+}
+
+MlirTypeID ireeGPUObjectAttrGetTypeID(void) {
+  return wrap(
+      mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr::getTypeID());
+}

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

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