[Mlir-commits] [mlir] 6995183 - [mlir][python] Register LLVM translations in the RegisterEverything for python (#70428)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 30 14:46:26 PDT 2023
Author: Jungwook Park
Date: 2023-10-30T14:46:21-07:00
New Revision: 6995183e174280f3987858bd13a4eca9905f6365
URL: https://github.com/llvm/llvm-project/commit/6995183e174280f3987858bd13a4eca9905f6365
DIFF: https://github.com/llvm/llvm-project/commit/6995183e174280f3987858bd13a4eca9905f6365.diff
LOG: [mlir][python] Register LLVM translations in the RegisterEverything for python (#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
Added:
mlir/test/python/dialects/gpu/dialect.py
mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
mlir/test/python/dialects/gpu/module-to-binary-rocdl.py
Modified:
mlir/lib/Bindings/Python/RegisterEverything.cpp
mlir/python/mlir/_mlir_libs/__init__.py
Removed:
mlir/test/python/dialects/gpu.py
################################################################################
diff --git a/mlir/lib/Bindings/Python/RegisterEverything.cpp b/mlir/lib/Bindings/Python/RegisterEverything.cpp
index fed5c36a625bff6..6b2f6b0a6a3b866 100644
--- a/mlir/lib/Bindings/Python/RegisterEverything.cpp
+++ b/mlir/lib/Bindings/Python/RegisterEverything.cpp
@@ -7,20 +7,17 @@
//===----------------------------------------------------------------------===//
#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_llvm_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..71c074bc955e8c3 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,11 @@ 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_llvm_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..70c08ceb7a6f2d3
--- /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..fad088cbd6d893b
--- /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