[python][tuner] Set up bindings for iree_codegen (#19108)

Starting with the `DispatchLoweringPassPipeline` attribute.

compiler/bindings/c/BUILD.bazel

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

compiler/bindings/c/CMakeLists.txt

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

compiler/bindings/c/iree/compiler/dialects/iree_codegen.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_CODEGEN_H
+#define IREE_COMPILER_DIALECTS_IREE_CODEGEN_H
+
+#include "mlir-c/IR.h"
+#include "mlir-c/Support.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// The following C API is **NOT STABLE** and likely to change in the future.
+// It mirrors the IREE Codegen Dialect which is not stable itself.
+
+MLIR_CAPI_EXPORTED bool
+ireeAttributeIsACodegenDispatchLoweringPassPipelineAttr(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED MlirTypeID
+ireeCodegenDispatchLoweringPassPipelineAttrGetTypeID(void);
+
+MLIR_CAPI_EXPORTED MlirAttribute ireeCodegenDispatchLoweringPassPipelineAttrGet(
+    MlirContext mlirCtx, uint32_t value);
+
+MLIR_CAPI_EXPORTED
+uint32_t
+ireeCodegenDispatchLoweringPassPipelineAttrGetValue(MlirAttribute attr);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // IREE_COMPILER_DIALECTS_IREE_CODEGEN_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/IREECodegenBinding.td
+  GEN_ENUM_BINDINGS
+  SOURCES dialects/iree_codegen.py
+  DIALECT_NAME iree_codegen
+)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT IREEPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler"

compiler/bindings/python/IREECompilerDialectsModule.cpp

@@ -5,11 +5,14 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 #include <cstdint>
+#include "iree/compiler/dialects/iree_codegen.h"
 #include "iree/compiler/dialects/iree_gpu.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/IR.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 
+static const char *kCodegenModuleImportPath =
+    MAKE_MLIR_PYTHON_QUALNAME("dialects.iree_codegen");
 static const char *kGpuModuleImportPath =
     MAKE_MLIR_PYTHON_QUALNAME("dialects.iree_gpu");
 
@@ -19,6 +22,34 @@ using namespace mlir::python::adaptors;
 PYBIND11_MODULE(_ireeCompilerDialects, m) {
   m.doc() = "iree-compiler dialects python extension";
 
+  auto iree_codegen_module =
+      m.def_submodule("iree_codegen", "iree_codegen dialect bindings");
+
+  //===-------------------------------------------------------------------===//
+  // CodegenDispatchLoweringPassPipelineAttr
+  //===-------------------------------------------------------------------===//
+
+  mlir_attribute_subclass(
+      iree_codegen_module, "DispatchLoweringPassPipelineAttr",
+      ireeAttributeIsACodegenDispatchLoweringPassPipelineAttr,
+      ireeCodegenDispatchLoweringPassPipelineAttrGetTypeID)
+      .def_classmethod(
+          "get",
+          [](const py::object &, uint32_t value, MlirContext ctx) {
+            return ireeCodegenDispatchLoweringPassPipelineAttrGet(ctx, value);
+          },
+          "cls"_a, "value"_a, "ctx"_a = py::none(),
+          "Gets an #iree_codegen.dispatch_lowering_pass_pipeline from "
+          "parameters.")
+      .def_property_readonly(
+          "raw_value", ireeCodegenDispatchLoweringPassPipelineAttrGetValue)
+      .def_property_readonly("value", [](MlirAttribute self) -> py::object {
+        uint32_t rawValue =
+            ireeCodegenDispatchLoweringPassPipelineAttrGetValue(self);
+        return py::module_::import(kCodegenModuleImportPath)
+            .attr("DispatchLoweringPassPipeline")(rawValue);
+      });
+
   auto iree_gpu_module =
       m.def_submodule("iree_gpu", "iree_gpu dialect bindings");
 
@@ -35,7 +66,7 @@ PYBIND11_MODULE(_ireeCompilerDialects, m) {
             return ireeGPUReorderWorkgroupsStrategyAttrGet(ctx, value);
           },
           "cls"_a, "value"_a, "ctx"_a = py::none(),
-          "Gets a gpu.reorder_workgroups_strategy from parameters.")
+          "Gets an #iree_gpu.reorder_workgroups_strategy from parameters.")
       .def_property_readonly("raw_value",
                              ireeGPUReorderWorkgroupsStrategyAttrGetValue)
       .def_property_readonly("value", [](MlirAttribute self) -> py::object {
@@ -75,7 +106,8 @@ PYBIND11_MODULE(_ireeCompilerDialects, m) {
           "no_reduce_shared_memory_bank_conflicts"_a = py::none(),
           "use_igemm_convolution"_a = py::none(),
           "reorder_workgroups_strategy"_a = py::none(), py::kw_only(),
-          "ctx"_a = py::none(), "Gets a gpu.pipeline_options from parameters.")
+          "ctx"_a = py::none(),
+          "Gets an #iree_gpu.pipeline_options from parameters.")
       .def_property_readonly(
           "prefetch_shared_memory",
           [](MlirAttribute self) -> std::optional<bool> {
@@ -125,7 +157,7 @@ PYBIND11_MODULE(_ireeCompilerDialects, m) {
             return ireeGPUMMAIntrinsicAttrGet(ctx, value);
           },
           "cls"_a, "value"_a, "ctx"_a = py::none(),
-          "Gets a gpu.mma_intrinsic from parameters.")
+          "Gets an #iree_gpu.mma_intrinsic from parameters.")
       .def_property_readonly("raw_value", ireeGPUMMAIntrinsicAttrGetValue)
       .def_property_readonly("value",
                              [](MlirAttribute self) -> py::object {
@@ -151,7 +183,7 @@ PYBIND11_MODULE(_ireeCompilerDialects, m) {
             return ireeGPUMMAAttrGet(ctx, value);
           },
           "cls"_a, "value"_a, "ctx"_a = py::none(),
-          "Gets a gpu.mma from parameters.")
+          "Gets an #iree_gpu.mma from parameters.")
       .def_property_readonly(
           "abc_element_types",
           [](MlirAttribute self) -> py::tuple {
@@ -185,7 +217,7 @@ PYBIND11_MODULE(_ireeCompilerDialects, m) {
             return ireeGPULoweringConfigAttrGet(ctx, attributeDictionary);
           },
           "cls"_a, "value"_a, "ctx"_a = py::none(),
-          "Gets a gpu.lowering_config from parameters.")
+          "Gets an #iree_gpu.lowering_config from parameters.")
       .def_property_readonly("attributes",
                              ireeGPULoweringConfigAttrGetAttributes);
 }

compiler/bindings/python/iree/compiler/dialects/IREECodegenBinding.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_IREECODEGEN_OPS
+#define PYTHON_BINDINGS_IREECODEGEN_OPS
+
+include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td"
+
+#endif // PYTHON_BINDINGS_IREECODEGEN_OPS

compiler/bindings/python/iree/compiler/dialects/iree_codegen.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_codegen_ops_gen import *
+from ._iree_codegen_enum_gen import *
+from .._mlir_libs._ireeCompilerDialects.iree_codegen import *

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

@@ -7,7 +7,7 @@
 from iree.compiler import ir
 
 # Make sure that our dialects import.
-from iree.compiler.dialects import flow, hal, stream, vm, util, iree_gpu
+from iree.compiler.dialects import flow, hal, stream, vm, util, iree_codegen, iree_gpu
 
 
 def run(fn):
@@ -19,6 +19,32 @@ def run(fn):
     return fn
 
 
+# ======================================================================
+# IREE Codegen Dialect
+# ======================================================================
+
+
+@run
+def codegen_dispatch_lowering_pass_pipeline():
+    pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUTileAndFuse
+    )
+    assert pipeline_attr is not None
+    assert (
+        pipeline_attr.value
+        == iree_codegen.DispatchLoweringPassPipeline.LLVMGPUTileAndFuse
+    )
+    assert pipeline_attr.raw_value == int(
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUTileAndFuse
+    )
+    assert "LLVMGPUTileAndFuse" in str(pipeline_attr)
+
+
+# ======================================================================
+# IREE GPU Dialect
+# ======================================================================
+
+
 @run
 def gpu_pipeline_options_attr():
     reorder_attr = iree_gpu.ReorderWorkgroupsStrategyAttr.get(

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_IREECodegenDialectCAPI.objects
   iree_compiler_API_Internal_IREEGPUDialectCAPI.objects
   iree_compiler_API_Internal_IREEMLIRLSPServerToolEntryPoint.objects
   iree_compiler_API_Internal_IREEOptToolEntryPoint.objects

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

@@ -129,6 +129,20 @@ iree_compiler_cc_library(
     ],
 )
 
+iree_compiler_cc_library(
+    name = "IREECodegenDialectCAPI",
+    srcs = [
+        "IREECodegenDialectCAPI.cpp",
+    ],
+    deps = [
+        "//compiler/bindings/c:headers",
+        "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
+        "@llvm-project//mlir:CAPIIR",
+        "@llvm-project//mlir:CAPIIRHeaders",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
 iree_compiler_cc_library(
     name = "IREEGPUDialectCAPI",
     srcs = [

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

@@ -106,6 +106,20 @@ iree_cc_library(
   PUBLIC
 )
 
+iree_cc_library(
+  NAME
+    IREECodegenDialectCAPI
+  SRCS
+    "IREECodegenDialectCAPI.cpp"
+  DEPS
+    IREELLVMIncludeSetup
+    MLIRCAPIIR
+    MLIRIR
+    iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
+    iree::compiler::bindings::c::headers
+  PUBLIC
+)
+
 iree_cc_library(
   NAME
     IREEGPUDialectCAPI

compiler/src/iree/compiler/API/Internal/IREECodegenDialectCAPI.cpp

@@ -0,0 +1,44 @@
+// 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 <cstdint>
+#include <type_traits>
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
+#include "iree/compiler/dialects/iree_codegen.h"
+#include "mlir-c/IR.h"
+#include "mlir/CAPI/IR.h"
+#include "mlir/CAPI/Support.h"
+
+using mlir::iree_compiler::IREE::Codegen::DispatchLoweringPassPipeline;
+using mlir::iree_compiler::IREE::Codegen::DispatchLoweringPassPipelineAttr;
+
+bool ireeAttributeIsACodegenDispatchLoweringPassPipelineAttr(
+    MlirAttribute attr) {
+  return llvm::isa<DispatchLoweringPassPipelineAttr>(unwrap(attr));
+}
+
+MlirTypeID ireeCodegenDispatchLoweringPassPipelineAttrGetTypeID() {
+  return wrap(DispatchLoweringPassPipelineAttr::getTypeID());
+}
+
+static_assert(
+    std::is_same_v<uint32_t,
+                   std::underlying_type_t<DispatchLoweringPassPipeline>>,
+    "Enum type changed");
+
+MlirAttribute
+ireeCodegenDispatchLoweringPassPipelineAttrGet(MlirContext mlirCtx,
+                                               uint32_t value) {
+  mlir::MLIRContext *ctx = unwrap(mlirCtx);
+  return wrap(DispatchLoweringPassPipelineAttr::get(
+      ctx, static_cast<DispatchLoweringPassPipeline>(value)));
+}
+
+uint32_t
+ireeCodegenDispatchLoweringPassPipelineAttrGetValue(MlirAttribute attr) {
+  return static_cast<uint32_t>(
+      llvm::cast<DispatchLoweringPassPipelineAttr>(unwrap(attr)).getValue());
+}

compiler/src/iree/compiler/API/api_exports.c

@@ -10,10 +10,15 @@
 
 #include <stdint.h>
 
+extern void ireeAttributeIsACodegenDispatchLoweringPassPipelineAttr();
+extern void ireeAttributeIsAGPULoweringConfigAttr();
 extern void ireeAttributeIsAGPUMMAAttr();
 extern void ireeAttributeIsAGPUMMAIntrinsicAttr();
 extern void ireeAttributeIsAGPUPipelineOptionsAttr();
 extern void ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr();
+extern void ireeCodegenDispatchLoweringPassPipelineAttrGet();
+extern void ireeCodegenDispatchLoweringPassPipelineAttrGetTypeID();
+extern void ireeCodegenDispatchLoweringPassPipelineAttrGetValue();
 extern void ireeCompilerEnumeratePlugins();
 extern void ireeCompilerEnumerateRegisteredHALTargetBackends();
 extern void ireeCompilerErrorDestroy();
@@ -65,6 +70,9 @@ extern void ireeCompilerSourceDestroy();
 extern void ireeCompilerSourceOpenFile();
 extern void ireeCompilerSourceSplit();
 extern void ireeCompilerSourceWrapBuffer();
+extern void ireeGPULoweringConfigAttrGet();
+extern void ireeGPULoweringConfigAttrGetAttributes();
+extern void ireeGPULoweringConfigAttrGetTypeID();
 extern void ireeGPUMMAAttrGet();
 extern void ireeGPUMMAAttrGetInfo();
 extern void ireeGPUMMAAttrGetTypeID();
@@ -856,10 +864,15 @@ extern void mlirVectorTypeIsScalable();
 
 uintptr_t __iree_compiler_hidden_force_extern() {
   uintptr_t x = 0;
+  x += (uintptr_t)&ireeAttributeIsACodegenDispatchLoweringPassPipelineAttr;
+  x += (uintptr_t)&ireeAttributeIsAGPULoweringConfigAttr;
   x += (uintptr_t)&ireeAttributeIsAGPUMMAAttr;
   x += (uintptr_t)&ireeAttributeIsAGPUMMAIntrinsicAttr;
   x += (uintptr_t)&ireeAttributeIsAGPUPipelineOptionsAttr;
   x += (uintptr_t)&ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr;
+  x += (uintptr_t)&ireeCodegenDispatchLoweringPassPipelineAttrGet;
+  x += (uintptr_t)&ireeCodegenDispatchLoweringPassPipelineAttrGetTypeID;
+  x += (uintptr_t)&ireeCodegenDispatchLoweringPassPipelineAttrGetValue;
   x += (uintptr_t)&ireeCompilerEnumeratePlugins;
   x += (uintptr_t)&ireeCompilerEnumerateRegisteredHALTargetBackends;
   x += (uintptr_t)&ireeCompilerErrorDestroy;
@@ -911,6 +924,9 @@ uintptr_t __iree_compiler_hidden_force_extern() {
   x += (uintptr_t)&ireeCompilerSourceOpenFile;
   x += (uintptr_t)&ireeCompilerSourceSplit;
   x += (uintptr_t)&ireeCompilerSourceWrapBuffer;
+  x += (uintptr_t)&ireeGPULoweringConfigAttrGet;
+  x += (uintptr_t)&ireeGPULoweringConfigAttrGetAttributes;
+  x += (uintptr_t)&ireeGPULoweringConfigAttrGetTypeID;
   x += (uintptr_t)&ireeGPUMMAAttrGet;
   x += (uintptr_t)&ireeGPUMMAAttrGetInfo;
   x += (uintptr_t)&ireeGPUMMAAttrGetTypeID;