[Mlir-commits] [mlir] [mlir][python] Register all translations in the RegisterEverything for python (PR #70428)
Jungwook Park
llvmlistbot at llvm.org
Fri Oct 27 12:28:10 PDT 2023
https://github.com/jungpark-mlir updated https://github.com/llvm/llvm-project/pull/70428
>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 1/3] 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 = "{{.*}}">]
+
>From 3adb8edc1a4cdd81372ab4cc224c85b97bde03f7 Mon Sep 17 00:00:00 2001
From: jungpark-mlir <jungwook at jungwook-22.04>
Date: Fri, 27 Oct 2023 14:26:10 +0100
Subject: [PATCH 2/3] Format.
---
mlir/lib/Bindings/Python/RegisterEverything.cpp | 5 ++---
mlir/python/mlir/_mlir_libs/__init__.py | 4 +++-
mlir/test/python/dialects/gpu/module-to-binary-nvvm.py | 4 ++--
mlir/test/python/dialects/gpu/module-to-binary-rocdl.py | 4 ++--
4 files changed, 9 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Bindings/Python/RegisterEverything.cpp b/mlir/lib/Bindings/Python/RegisterEverything.cpp
index cdb4a30228a8ef4..6c0bdb43eebbf7c 100644
--- a/mlir/lib/Bindings/Python/RegisterEverything.cpp
+++ b/mlir/lib/Bindings/Python/RegisterEverything.cpp
@@ -15,9 +15,8 @@ PYBIND11_MODULE(_mlirRegisterEverything, m) {
m.def("register_dialects", [](MlirDialectRegistry registry) {
mlirRegisterAllDialects(registry);
});
- m.def("register_translations", [](MlirContext context) {
- mlirRegisterAllLLVMTranslations(context);
- });
+ m.def("register_translations",
+ [](MlirContext context) { mlirRegisterAllLLVMTranslations(context); });
// Register all passes on load.
mlirRegisterAllPasses();
diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py
index 68b841fdf21aa41..2702062024930a8 100644
--- a/mlir/python/mlir/_mlir_libs/__init__.py
+++ b/mlir/python/mlir/_mlir_libs/__init__.py
@@ -104,7 +104,9 @@ def __init__(self, *args, **kwargs):
# 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)
+ 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/module-to-binary-nvvm.py b/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
index af0938715ff370a..70c08ceb7a6f2d3 100644
--- a/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
+++ b/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
@@ -38,6 +38,7 @@ def testGPUToLLVMBin():
# CHECK-LABEL:gpu.binary @kernel_module1
# CHECK:[#gpu.object<#nvvm.target<chip = "sm_70">, offload = "{{.*}}">]
+
# CHECK-LABEL: testGPUToASMBin
@run
def testGPUToASMBin():
@@ -57,8 +58,7 @@ def testGPUToASMBin():
)
pm = PassManager("any")
pm.add("gpu-module-to-binary{format=isa}")
- pm.run(module.operation)
+ 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
index 3f0071d1cc6ea7f..fad088cbd6d893b 100644
--- a/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py
+++ b/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py
@@ -38,6 +38,7 @@ def testGPUToLLVMBin():
# CHECK-LABEL:gpu.binary @kernel_module1
# CHECK:[#gpu.object<#rocdl.target<chip = "gfx90a">, offload = "{{.*}}">]
+
# CHECK-LABEL: testGPUToASMBin
@run
def testGPUToASMBin():
@@ -57,8 +58,7 @@ def testGPUToASMBin():
)
pm = PassManager("any")
pm.add("gpu-module-to-binary{format=isa}")
- pm.run(module.operation)
+ 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 = "{{.*}}">]
-
>From 983a88dfe9cf15ee52e7b331b092b8da9fb992ff Mon Sep 17 00:00:00 2001
From: jungpark-mlir <jungwook at jungwook-22.04>
Date: Fri, 27 Oct 2023 20:27:06 +0100
Subject: [PATCH 3/3] rename register function
---
mlir/lib/Bindings/Python/RegisterEverything.cpp | 2 +-
mlir/python/mlir/_mlir_libs/__init__.py | 2 +-
2 files changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Bindings/Python/RegisterEverything.cpp b/mlir/lib/Bindings/Python/RegisterEverything.cpp
index 6c0bdb43eebbf7c..6b2f6b0a6a3b866 100644
--- a/mlir/lib/Bindings/Python/RegisterEverything.cpp
+++ b/mlir/lib/Bindings/Python/RegisterEverything.cpp
@@ -15,7 +15,7 @@ PYBIND11_MODULE(_mlirRegisterEverything, m) {
m.def("register_dialects", [](MlirDialectRegistry registry) {
mlirRegisterAllDialects(registry);
});
- m.def("register_translations",
+ m.def("register_llvm_translations",
[](MlirContext context) { mlirRegisterAllLLVMTranslations(context); });
// Register all passes on load.
diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py
index 2702062024930a8..71c074bc955e8c3 100644
--- a/mlir/python/mlir/_mlir_libs/__init__.py
+++ b/mlir/python/mlir/_mlir_libs/__init__.py
@@ -107,7 +107,7 @@ def __init__(self, *args, **kwargs):
logger.debug(
"Registering translations from initializer %r", init_module
)
- init_module.register_translations(self)
+ init_module.register_llvm_translations(self)
ir.Context = Context
More information about the Mlir-commits
mailing list