[mlir][python] add binding to #gpu.object

[makslevental]

The goal here is to enable extracting the compiled artifact for downstream (GPU) compilers (currently you have to do some unholy things...).

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/88992.diff

7 Files Affected:

• (modified) mlir/include/mlir-c/Dialect/GPU.h (+25)
• (added) mlir/lib/Bindings/Python/DialectGPU.cpp (+69)
• (modified) mlir/lib/CAPI/Dialect/GPU.cpp (+53-1)
• (modified) mlir/python/CMakeLists.txt (+11)
• (modified) mlir/python/mlir/dialects/gpu/init.py (+1)
• (modified) mlir/test/python/dialects/gpu/dialect.py (+26)
• (modified) mlir/test/python/dialects/gpu/module-to-binary-nvvm.py (+26-4)

diff --git a/mlir/include/mlir-c/Dialect/GPU.h b/mlir/include/mlir-c/Dialect/GPU.h
index 1a18d82c01d53e..2adf73ddff6eae 100644
--- a/mlir/include/mlir-c/Dialect/GPU.h
+++ b/mlir/include/mlir-c/Dialect/GPU.h
@@ -19,6 +19,31 @@ extern "C" {
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(GPU, gpu);
 
+//===---------------------------------------------------------------------===//
+// ObjectAttr
+//===---------------------------------------------------------------------===//
+
+MLIR_CAPI_EXPORTED bool mlirAttributeIsAGPUObjectAttr(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirGPUObjectAttrGet(MlirContext mlirCtx, MlirAttribute target, uint32_t format,
+                     MlirStringRef objectStrRef, MlirAttribute mlirObjectProps);
+
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirGPUObjectAttrGetTarget(MlirAttribute mlirObjectAttr);
+
+MLIR_CAPI_EXPORTED uint32_t
+mlirGPUObjectAttrGetFormat(MlirAttribute mlirObjectAttr);
+
+MLIR_CAPI_EXPORTED MlirStringRef
+mlirGPUObjectAttrGetObject(MlirAttribute mlirObjectAttr);
+
+MLIR_CAPI_EXPORTED bool
+mlirGPUObjectAttrHasProperties(MlirAttribute mlirObjectAttr);
+
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirGPUObjectAttrGetProperties(MlirAttribute mlirObjectAttr);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/lib/Bindings/Python/DialectGPU.cpp b/mlir/lib/Bindings/Python/DialectGPU.cpp
new file mode 100644
index 00000000000000..243b57c5444841
--- /dev/null
+++ b/mlir/lib/Bindings/Python/DialectGPU.cpp
@@ -0,0 +1,69 @@
+//===- DialectGPU.cpp - Pybind module for the GPU passes ------------------===//
+//
+// Part of the LLVM Project, 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 "mlir-c/Dialect/GPU.h"
+#include "mlir-c/IR.h"
+#include "mlir-c/Support.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
+
+#include <pybind11/detail/common.h>
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+using namespace mlir;
+using namespace mlir::python;
+using namespace mlir::python::adaptors;
+
+// -----------------------------------------------------------------------------
+// Module initialization.
+// -----------------------------------------------------------------------------
+
+PYBIND11_MODULE(_mlirDialectsGPU, m) {
+  m.doc() = "MLIR GPU Dialect";
+
+  //===-------------------------------------------------------------------===//
+  // ObjectAttr
+  //===-------------------------------------------------------------------===//
+
+  mlir_attribute_subclass(m, "ObjectAttr", mlirAttributeIsAGPUObjectAttr)
+      .def_classmethod(
+          "get",
+          [](py::object cls, MlirAttribute target, uint32_t format,
+             py::bytes object, std::optional<MlirAttribute> mlirObjectProps,
+             std::optional<MlirContext> context) {
+            py::buffer_info info(py::buffer(object).request());
+            MlirStringRef objectStrRef =
+                mlirStringRefCreate(static_cast<char *>(info.ptr), info.size);
+            return cls(mlirGPUObjectAttrGet(
+                context.has_value() ? *context
+                                    : mlirAttributeGetContext(target),
+                target, format, objectStrRef,
+                mlirObjectProps.has_value() ? *mlirObjectProps
+                                            : MlirAttribute{nullptr}));
+          },
+          "cls"_a, "target"_a, "format"_a, "object"_a,
+          "properties"_a = py::none(), "context"_a = py::none(),
+          "Gets a gpu.object from parameters.")
+      .def_property_readonly(
+          "target",
+          [](MlirAttribute self) { return mlirGPUObjectAttrGetTarget(self); })
+      .def_property_readonly(
+          "format",
+          [](MlirAttribute self) { return mlirGPUObjectAttrGetFormat(self); })
+      .def_property_readonly(
+          "object",
+          [](MlirAttribute self) {
+            MlirStringRef stringRef = mlirGPUObjectAttrGetObject(self);
+            return py::bytes(stringRef.data, stringRef.length);
+          })
+      .def_property_readonly("properties", [](MlirAttribute self) {
+        if (mlirGPUObjectAttrHasProperties(self))
+          return py::cast(mlirGPUObjectAttrGetProperties(self));
+        return py::none().cast<py::object>();
+      });
+}
diff --git a/mlir/lib/CAPI/Dialect/GPU.cpp b/mlir/lib/CAPI/Dialect/GPU.cpp
index cd58f0e249a9e2..c23422a8c9b06f 100644
--- a/mlir/lib/CAPI/Dialect/GPU.cpp
+++ b/mlir/lib/CAPI/Dialect/GPU.cpp
@@ -1,4 +1,4 @@
-//===- GPUc.cpp - C Interface for GPU dialect ----------------------------===//
+//===- GPU.cpp - C Interface for GPU dialect ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -9,5 +9,57 @@
 #include "mlir-c/Dialect/GPU.h"
 #include "mlir/CAPI/Registration.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "llvm/Support/Casting.h"
 
 MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(GPU, gpu, mlir::gpu::GPUDialect)
+
+//===---------------------------------------------------------------------===//
+// ObjectAttr
+//===---------------------------------------------------------------------===//
+
+bool mlirAttributeIsAGPUObjectAttr(MlirAttribute attr) {
+  return llvm::isa<mlir::gpu::ObjectAttr>(unwrap(attr));
+}
+
+MlirAttribute mlirGPUObjectAttrGet(MlirContext mlirCtx, MlirAttribute target,
+                                   uint32_t format, MlirStringRef objectStrRef,
+                                   MlirAttribute mlirObjectProps) {
+  mlir::MLIRContext *ctx = unwrap(mlirCtx);
+  llvm::StringRef object = unwrap(objectStrRef);
+  mlir::DictionaryAttr objectProps =
+      llvm::cast<mlir::DictionaryAttr>(unwrap(mlirObjectProps));
+  return wrap(mlir::gpu::ObjectAttr::get(
+      ctx, unwrap(target), static_cast<mlir::gpu::CompilationTarget>(format),
+      mlir::StringAttr::get(ctx, object), objectProps));
+}
+
+MlirAttribute mlirGPUObjectAttrGetTarget(MlirAttribute mlirObjectAttr) {
+  mlir::gpu::ObjectAttr objectAttr =
+      llvm::cast<mlir::gpu::ObjectAttr>(unwrap(mlirObjectAttr));
+  return wrap(objectAttr.getTarget());
+}
+
+uint32_t mlirGPUObjectAttrGetFormat(MlirAttribute mlirObjectAttr) {
+  mlir::gpu::ObjectAttr objectAttr =
+      llvm::cast<mlir::gpu::ObjectAttr>(unwrap(mlirObjectAttr));
+  return static_cast<uint32_t>(objectAttr.getFormat());
+}
+
+MlirStringRef mlirGPUObjectAttrGetObject(MlirAttribute mlirObjectAttr) {
+  mlir::gpu::ObjectAttr objectAttr =
+      llvm::cast<mlir::gpu::ObjectAttr>(unwrap(mlirObjectAttr));
+  llvm::StringRef object = objectAttr.getObject();
+  return mlirStringRefCreate(object.data(), object.size());
+}
+
+bool mlirGPUObjectAttrHasProperties(MlirAttribute mlirObjectAttr) {
+  mlir::gpu::ObjectAttr objectAttr =
+      llvm::cast<mlir::gpu::ObjectAttr>(unwrap(mlirObjectAttr));
+  return objectAttr.getProperties() != nullptr;
+}
+
+MlirAttribute mlirGPUObjectAttrGetProperties(MlirAttribute mlirObjectAttr) {
+  mlir::gpu::ObjectAttr objectAttr =
+      llvm::cast<mlir::gpu::ObjectAttr>(unwrap(mlirObjectAttr));
+  return wrap(objectAttr.getProperties());
+}
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index c27ee688a04087..d31bad34afa82c 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -498,6 +498,17 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
     MLIRCAPILinalg
 )
 
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
+  MODULE_NAME _mlirDialectsGPU
+  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
+  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  SOURCES
+    DialectGPU.cpp
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIIR
+    MLIRCAPIGPU
+)
+
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
   MODULE_NAME _mlirDialectsLLVM
   ADD_TO_PARENT MLIRPythonSources.Dialects.llvm
diff --git a/mlir/python/mlir/dialects/gpu/__init__.py b/mlir/python/mlir/dialects/gpu/__init__.py
index 033386b0f803b2..4cd80aa8b7ca85 100644
--- a/mlir/python/mlir/dialects/gpu/__init__.py
+++ b/mlir/python/mlir/dialects/gpu/__init__.py
@@ -4,3 +4,4 @@
 
 from .._gpu_ops_gen import *
 from .._gpu_enum_gen import *
+from ..._mlir_libs._mlirDialectsGPU import *
diff --git a/mlir/test/python/dialects/gpu/dialect.py b/mlir/test/python/dialects/gpu/dialect.py
index 2f49e2e053999b..c2cbd33cd00df7 100644
--- a/mlir/test/python/dialects/gpu/dialect.py
+++ b/mlir/test/python/dialects/gpu/dialect.py
@@ -30,3 +30,29 @@ def testMMAElementWiseAttr():
     # CHECK: %block_dim_y = gpu.block_dim  y
     print(module)
     pass
+
+
+# CHECK-LABEL: testObjectAttr
+@run
+def testObjectAttr():
+    module = Module.create()
+    target = Attribute.parse("#nvvm.target")
+    format = gpu.CompilationTarget.Fatbin
+    object = b"BC\xc0\xde5\x14\x00\x00\x05\x00\x00\x00b\x0c0$MY\xbef"
+    properties = DictAttr.get({"O": IntegerAttr.get(IntegerType.get_signless(32), 2)})
+    o = gpu.ObjectAttr.get(target, format, object, properties)
+    # CHECK: #gpu.object<#nvvm.target, properties = {O = 2 : i32}, "BC\C0\DE5\14\00\00\05\00\00\00b\0C0$MY\BEf">
+    print(o)
+    assert o.object == object
+
+    o = gpu.ObjectAttr.get(target, format, object)
+    # CHECK: #gpu.object<#nvvm.target, "BC\C0\DE5\14\00\00\05\00\00\00b\0C0$MY\BEf">
+    print(o)
+
+    object = (
+        b"//\n// Generated by LLVM NVPTX Back-End\n//\n\n.version 6.0\n.target sm_50"
+    )
+    o = gpu.ObjectAttr.get(target, format, object)
+    # CHECK: #gpu.object<#nvvm.target, "//\0A// Generated by LLVM NVPTX Back-End\0A//\0A\0A.version 6.0\0A.target sm_50">
+    print(o)
+    assert o.object == object
diff --git a/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py b/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
index 1c2eb652e71f91..e0225911b21443 100644
--- a/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
+++ b/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
@@ -34,9 +34,20 @@ def testGPUToLLVMBin():
     pm = PassManager("any")
     pm.add("gpu-module-to-binary{format=llvm}")
     pm.run(module.operation)
+    # CHECK-LABEL: gpu.binary @kernel_module1
     print(module)
-    # CHECK-LABEL:gpu.binary @kernel_module1
-    # CHECK:[#gpu.object<#nvvm.target<chip = "sm_70">, offload = "{{.*}}">]
+
+    o = gpu.ObjectAttr(module.body.operations[0].objects[0])
+    # CHECK: #gpu.object<#nvvm.target<chip = "sm_70">, offload = "{{.*}}">
+    print(o)
+    # CHECK: #nvvm.target<chip = "sm_70">
+    print(o.target)
+    # CHECK: offload
+    print(gpu.CompilationTarget(o.format))
+    # CHECK: b'{{.*}}'
+    print(o.object)
+    # CHECK: None
+    print(o.properties)
 
 
 # CHECK-LABEL: testGPUToASMBin
@@ -59,6 +70,17 @@ def testGPUToASMBin():
     pm = PassManager("any")
     pm.add("gpu-module-to-binary{format=isa}")
     pm.run(module.operation)
-    print(module)
     # CHECK-LABEL:gpu.binary @kernel_module2
-    # CHECK:[#gpu.object<#nvvm.target<flags = {fast}>, properties = {O = 2 : i32}, assembly = "{{.*}}">, #gpu.object<#nvvm.target, properties = {O = 2 : i32}, assembly = "{{.*}}">]
+    print(module)
+
+    o = gpu.ObjectAttr(module.body.operations[0].objects[0])
+    # CHECK: #gpu.object<#nvvm.target<flags = {fast}>
+    print(o)
+    # CHECK: #nvvm.target<flags = {fast}>
+    print(o.target)
+    # CHECK: assembly
+    print(gpu.CompilationTarget(o.format))
+    # CHECK: b'//\n// Generated by LLVM NVPTX Back-End{{.*}}'
+    print(o.object)
+    # CHECK: {O = 2 : i32}
+    print(o.properties)

[grypp]

Not an expert in py bindings fields. But PR makes setting target nicer and simpler