[Mlir-commits] [mlir] Register all translations in the RegisterEverything for python (PR #70428)
Jungwook Park
llvmlistbot at llvm.org
Fri Oct 27 01:55:39 PDT 2023
https://github.com/jungpark-mlir created https://github.com/llvm/llvm-project/pull/70428
Added missing register_translations in python to replicate the same in the C-API
Cleaned up the current calls to register passes where the other calls are already embedded in the mlirRegisterAllPasses.
found here,
https://discourse.llvm.org/t/opencl-example/74187
>From 227c7b2c17a246e09f77afb41e725b9d34ce4c74 Mon Sep 17 00:00:00 2001
From: jungwook <jungwook at jungwook-22.04>
Date: Thu, 26 Oct 2023 16:31:52 +0100
Subject: [PATCH] Register all translations in the RegisterEverything for
python Clean up the current calls to register passes where the other calls
are already embedded in the mlirRegisterAllPasses. Added
register_translations found here,
https://discourse.llvm.org/t/opencl-example/74187
---
.../Bindings/Python/RegisterEverything.cpp | 10 ++-
mlir/python/mlir/_mlir_libs/__init__.py | 6 +-
.../dialects/{gpu.py => gpu/dialect.py} | 0
.../dialects/gpu/module-to-binary-nvvm.py | 64 +++++++++++++++++++
.../dialects/gpu/module-to-binary-rocdl.py | 64 +++++++++++++++++++
5 files changed, 137 insertions(+), 7 deletions(-)
rename mlir/test/python/dialects/{gpu.py => gpu/dialect.py} (100%)
create mode 100644 mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
create mode 100644 mlir/test/python/dialects/gpu/module-to-binary-rocdl.py
diff --git a/mlir/lib/Bindings/Python/RegisterEverything.cpp b/mlir/lib/Bindings/Python/RegisterEverything.cpp
index fed5c36a625bff6..cdb4a30228a8ef4 100644
--- a/mlir/lib/Bindings/Python/RegisterEverything.cpp
+++ b/mlir/lib/Bindings/Python/RegisterEverything.cpp
@@ -7,20 +7,18 @@
//===----------------------------------------------------------------------===//
#include "mlir-c/RegisterEverything.h"
-#include "mlir-c/Conversion.h"
-#include "mlir-c/Transforms.h"
-
#include "mlir/Bindings/Python/PybindAdaptors.h"
PYBIND11_MODULE(_mlirRegisterEverything, m) {
- m.doc() = "MLIR All Upstream Dialects and Passes Registration";
+ m.doc() = "MLIR All Upstream Dialects, Translations and Passes Registration";
m.def("register_dialects", [](MlirDialectRegistry registry) {
mlirRegisterAllDialects(registry);
});
+ m.def("register_translations", [](MlirContext context) {
+ mlirRegisterAllLLVMTranslations(context);
+ });
// Register all passes on load.
mlirRegisterAllPasses();
- mlirRegisterConversionPasses();
- mlirRegisterTransformsPasses();
}
diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py
index 03fcb10130c3ae8..68b841fdf21aa41 100644
--- a/mlir/python/mlir/_mlir_libs/__init__.py
+++ b/mlir/python/mlir/_mlir_libs/__init__.py
@@ -83,7 +83,8 @@ def process_initializer_module(module_name):
# If _mlirRegisterEverything is built, then include it as an initializer
# module.
- process_initializer_module("_mlirRegisterEverything")
+ if process_initializer_module("_mlirRegisterEverything"):
+ init_module = importlib.import_module(f"._mlirRegisterEverything", __name__)
# Load all _site_initialize_{i} modules, where 'i' is a number starting
# at 0.
@@ -102,6 +103,9 @@ def __init__(self, *args, **kwargs):
# all dialects. It is being done here in order to preserve existing
# behavior. See: https://github.com/llvm/llvm-project/issues/56037
self.load_all_available_dialects()
+ if init_module:
+ logger.debug("Registering translations from initializer %r", init_module)
+ init_module.register_translations(self)
ir.Context = Context
diff --git a/mlir/test/python/dialects/gpu.py b/mlir/test/python/dialects/gpu/dialect.py
similarity index 100%
rename from mlir/test/python/dialects/gpu.py
rename to mlir/test/python/dialects/gpu/dialect.py
diff --git a/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py b/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
new file mode 100644
index 000000000000000..af0938715ff370a
--- /dev/null
+++ b/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
@@ -0,0 +1,64 @@
+# REQUIRES: host-supports-nvptx
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+import mlir.dialects.gpu as gpu
+import mlir.dialects.gpu.passes
+from mlir.passmanager import *
+
+
+def run(f):
+ print("\nTEST:", f.__name__)
+ with Context(), Location.unknown():
+ f()
+ return f
+
+
+# CHECK-LABEL: testGPUToLLVMBin
+ at run
+def testGPUToLLVMBin():
+ with Context():
+ module = Module.parse(
+ r"""
+module attributes {gpu.container_module} {
+ gpu.module @kernel_module1 [#nvvm.target<chip = "sm_70">] {
+ llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr<f32>,
+ %arg2: !llvm.ptr<f32>, %arg3: i64, %arg4: i64,
+ %arg5: i64) attributes {gpu.kernel} {
+ llvm.return
+ }
+ }
+}
+ """
+ )
+ pm = PassManager("any")
+ pm.add("gpu-module-to-binary{format=llvm}")
+ pm.run(module.operation)
+ print(module)
+ # CHECK-LABEL:gpu.binary @kernel_module1
+ # CHECK:[#gpu.object<#nvvm.target<chip = "sm_70">, offload = "{{.*}}">]
+
+# CHECK-LABEL: testGPUToASMBin
+ at run
+def testGPUToASMBin():
+ with Context():
+ module = Module.parse(
+ r"""
+module attributes {gpu.container_module} {
+ gpu.module @kernel_module2 [#nvvm.target<flags = {fast}>, #nvvm.target] {
+ llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr<f32>,
+ %arg2: !llvm.ptr<f32>, %arg3: i64, %arg4: i64,
+ %arg5: i64) attributes {gpu.kernel} {
+ llvm.return
+ }
+ }
+}
+ """
+ )
+ 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 = "{{.*}}">]
+
diff --git a/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py b/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py
new file mode 100644
index 000000000000000..3f0071d1cc6ea7f
--- /dev/null
+++ b/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py
@@ -0,0 +1,64 @@
+# REQUIRES: host-supports-amdgpu
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+import mlir.dialects.gpu as gpu
+import mlir.dialects.gpu.passes
+from mlir.passmanager import *
+
+
+def run(f):
+ print("\nTEST:", f.__name__)
+ with Context(), Location.unknown():
+ f()
+ return f
+
+
+# CHECK-LABEL: testGPUToLLVMBin
+ at run
+def testGPUToLLVMBin():
+ with Context():
+ module = Module.parse(
+ r"""
+module attributes {gpu.container_module} {
+ gpu.module @kernel_module1 [#rocdl.target<chip = "gfx90a">] {
+ llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr<f32>,
+ %arg2: !llvm.ptr<f32>, %arg3: i64, %arg4: i64,
+ %arg5: i64) attributes {gpu.kernel} {
+ llvm.return
+ }
+ }
+}
+ """
+ )
+ pm = PassManager("any")
+ pm.add("gpu-module-to-binary{format=llvm}")
+ pm.run(module.operation)
+ print(module)
+ # CHECK-LABEL:gpu.binary @kernel_module1
+ # CHECK:[#gpu.object<#rocdl.target<chip = "gfx90a">, offload = "{{.*}}">]
+
+# CHECK-LABEL: testGPUToASMBin
+ at run
+def testGPUToASMBin():
+ with Context():
+ module = Module.parse(
+ r"""
+module attributes {gpu.container_module} {
+ gpu.module @kernel_module2 [#rocdl.target<flags = {fast}>, #rocdl.target] {
+ llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr<f32>,
+ %arg2: !llvm.ptr<f32>, %arg3: i64, %arg4: i64,
+ %arg5: i64) attributes {gpu.kernel} {
+ llvm.return
+ }
+ }
+}
+ """
+ )
+ 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<#rocdl.target<flags = {fast}>, assembly = "{{.*}}">, #gpu.object<#rocdl.target, assembly = "{{.*}}">]
+
More information about the Mlir-commits
mailing list