[Mlir-commits] [mlir] [mlir] enable python bindings for nvgpu transforms (PR #68088)
Oleksandr Alex Zinenko
llvmlistbot at llvm.org
Tue Oct 3 05:28:16 PDT 2023
https://github.com/ftynse updated https://github.com/llvm/llvm-project/pull/68088
>From aa08d7d8d3dcb46978199bed4303de82675a4f4b Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko at google.com>
Date: Tue, 3 Oct 2023 10:55:48 +0000
Subject: [PATCH 1/2] [mlir] enable python bindings for nvgpu transforms
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Expose the autogenerated bindings.
Co-authored-by: Martin Lücke <mluecke at google.com>
---
mlir/python/CMakeLists.txt | 9 +++++++
.../python/mlir/dialects/NVGPUTransformOps.td | 20 ++++++++++++++
mlir/python/mlir/dialects/transform/nvgpu.py | 5 ++++
.../python/dialects/transform_nvgpu_ext.py | 27 +++++++++++++++++++
.../mlir/python/BUILD.bazel | 20 ++++++++++++++
5 files changed, 81 insertions(+)
create mode 100644 mlir/python/mlir/dialects/NVGPUTransformOps.td
create mode 100644 mlir/python/mlir/dialects/transform/nvgpu.py
create mode 100644 mlir/test/python/dialects/transform_nvgpu_ext.py
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 9368cb4c2f1657f..f09c79a0a88e718 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -200,6 +200,15 @@ declare_mlir_dialect_extension_python_bindings(
DIALECT_NAME transform
EXTENSION_NAME memref_transform)
+declare_mlir_dialect_extension_python_bindings(
+ADD_TO_PARENT MLIRPythonSources.Dialects
+ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+ TD_FILE dialects/NVGPUTransformOps.td
+ SOURCES
+ dialects/transform/nvgpu.py
+ DIALECT_NAME transform
+ EXTENSION_NAME nvgpu_transform)
+
declare_mlir_dialect_extension_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
diff --git a/mlir/python/mlir/dialects/NVGPUTransformOps.td b/mlir/python/mlir/dialects/NVGPUTransformOps.td
new file mode 100644
index 000000000000000..1f504e3220e92de
--- /dev/null
+++ b/mlir/python/mlir/dialects/NVGPUTransformOps.td
@@ -0,0 +1,20 @@
+//===-- NVGPUTransformOps.td -------------------------------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Entry point of the Python bindings generator for the transform ops provided
+// by the NVGPU dialect.
+//
+//===----------------------------------------------------------------------===//
+
+
+#ifndef PYTHON_BINDINGS_NVGPU_TRANSFORM_OPS
+#define PYTHON_BINDINGS_NVGPU_TRANSFORM_OPS
+
+include "mlir/Dialect/NVGPU/TransformOps/NVGPUTransformOps.td"
+
+#endif // PYTHON_BINDINGS_NVGPU_TRANSFORM_OPS
diff --git a/mlir/python/mlir/dialects/transform/nvgpu.py b/mlir/python/mlir/dialects/transform/nvgpu.py
new file mode 100644
index 000000000000000..74ba4c9aeb6c107
--- /dev/null
+++ b/mlir/python/mlir/dialects/transform/nvgpu.py
@@ -0,0 +1,5 @@
+# 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
+
+from .._nvgpu_transform_ops_gen import *
diff --git a/mlir/test/python/dialects/transform_nvgpu_ext.py b/mlir/test/python/dialects/transform_nvgpu_ext.py
new file mode 100644
index 000000000000000..1de8b25bab7a5c6
--- /dev/null
+++ b/mlir/test/python/dialects/transform_nvgpu_ext.py
@@ -0,0 +1,27 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+from mlir.dialects import transform
+from mlir.dialects.transform import nvgpu
+
+
+def run(f):
+ with Context(), Location.unknown():
+ module = Module.create()
+ with InsertionPoint(module.body):
+ print("\nTEST:", f.__name__)
+ f()
+ print(module)
+ return f
+
+
+ at run
+def testCreateAsyncGroups():
+ sequence = transform.SequenceOp(
+ transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()
+ )
+ with InsertionPoint(sequence.body):
+ nvgpu.CreateAsyncGroupsOp(transform.AnyOpType.get(), sequence.bodyTarget)
+ transform.YieldOp()
+ # CHECK-LABEL: TEST: testCreateAsyncGroups
+ # CHECK: transform.nvgpu.create_async_groups
diff --git a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
index b8b70f4685feaef..949bfa4fbda3a6c 100644
--- a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
@@ -1209,6 +1209,25 @@ gentbl_filegroup(
],
)
+gentbl_filegroup(
+ name = "NVGPUTransformOpsPyGen",
+ tbl_outs = [
+ (
+ [
+ "-gen-python-op-bindings",
+ "-bind-dialect=transform",
+ "-dialect-extension=nvgpu_transform",
+ ],
+ "mlir/dialects/_nvgpu_transform_ops_gen.py",
+ ),
+ ],
+ tblgen = "//mlir:mlir-tblgen",
+ td_file = "mlir/dialects/NVGPUTransformOps.td",
+ deps = [
+ "//mlir:NVGPUTransformOpsTdFiles",
+ ],
+)
+
gentbl_filegroup(
name = "PDLTransformOpsPyGen",
tbl_outs = [
@@ -1327,6 +1346,7 @@ filegroup(
":GPUTransformOpsPyGen",
":LoopTransformOpsPyGen",
":MemRefTransformOpsPyGen",
+ ":NVGPUTransformOpsPyGen",
":PDLTransformOpsPyGen",
":SparseTensorTransformOpsPyGen",
":StructureTransformEnumPyGen",
>From 9d2bdc45a198bfee345779c91522a12c7bb1bfc0 Mon Sep 17 00:00:00 2001
From: "Oleksandr \"Alex\" Zinenko" <ftynse at gmail.com>
Date: Tue, 3 Oct 2023 14:28:05 +0200
Subject: [PATCH 2/2] Fix whitespace
---
mlir/python/CMakeLists.txt | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index f09c79a0a88e718..088d9a765b97730 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -201,8 +201,8 @@ declare_mlir_dialect_extension_python_bindings(
EXTENSION_NAME memref_transform)
declare_mlir_dialect_extension_python_bindings(
-ADD_TO_PARENT MLIRPythonSources.Dialects
-ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+ ADD_TO_PARENT MLIRPythonSources.Dialects
+ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
TD_FILE dialects/NVGPUTransformOps.td
SOURCES
dialects/transform/nvgpu.py
More information about the Mlir-commits
mailing list