[Mlir-commits] [mlir] [mlir][py] Add NVGPU's `TensorMapDescriptorType` in py bindings (PR #88855)
Guray Ozen
llvmlistbot at llvm.org
Wed Apr 17 01:17:51 PDT 2024
https://github.com/grypp updated https://github.com/llvm/llvm-project/pull/88855
>From d0d37b7675b14547be615b47f7a5203b8ac3a035 Mon Sep 17 00:00:00 2001
From: grypp <guray.ozen at gmail.com>
Date: Tue, 16 Apr 2024 08:26:49 +0000
Subject: [PATCH 1/3] [mlir][py] Add NVGPU's `TensorMapDescriptorType` in py
bindings
This PR adds NVGPU dialects' TensorMapDescriptorType in the py bindings.
---
mlir/include/mlir-c/Dialect/NVGPU.h | 11 ++++++
mlir/lib/Bindings/Python/DialectNVGPU.cpp | 46 +++++++++++++++++++++++
mlir/lib/CAPI/Dialect/NVGPU.cpp | 18 +++++++++
mlir/python/CMakeLists.txt | 13 +++++++
mlir/python/mlir/dialects/nvgpu.py | 1 +
mlir/test/python/dialects/nvgpu.py | 13 +++++++
6 files changed, 102 insertions(+)
create mode 100644 mlir/lib/Bindings/Python/DialectNVGPU.cpp
diff --git a/mlir/include/mlir-c/Dialect/NVGPU.h b/mlir/include/mlir-c/Dialect/NVGPU.h
index 580d566794c09f..143284ee32b5ae 100644
--- a/mlir/include/mlir-c/Dialect/NVGPU.h
+++ b/mlir/include/mlir-c/Dialect/NVGPU.h
@@ -11,6 +11,7 @@
#define MLIR_C_DIALECT_NVGPU_H
#include "mlir-c/IR.h"
+#include "mlir-c/Support.h"
#ifdef __cplusplus
extern "C" {
@@ -18,6 +19,16 @@ extern "C" {
MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(NVGPU, nvgpu);
+//===---------------------------------------------------------------------===//
+// TensorMapDescriptorType
+//===---------------------------------------------------------------------===//
+
+MLIR_CAPI_EXPORTED bool mlirTypeIsANVGPUTensorMapDescriptorType(MlirType type);
+
+MLIR_CAPI_EXPORTED MlirType mlirNVGPUTensorMapDescriptorTypeGet(
+ MlirContext ctx, MlirType tensorType, int swizzle, int l2promo, int oob,
+ int interleave);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/lib/Bindings/Python/DialectNVGPU.cpp b/mlir/lib/Bindings/Python/DialectNVGPU.cpp
new file mode 100644
index 00000000000000..d0f7b2500085d1
--- /dev/null
+++ b/mlir/lib/Bindings/Python/DialectNVGPU.cpp
@@ -0,0 +1,46 @@
+//===- DialectLinalg.cpp - Pybind module for Nvgpu dialect API support --===//
+//
+// 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 "IRModule.h"
+#include "mlir-c/Dialect/NVGPU.h"
+#include "mlir-c/IR.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
+#include "llvm/Support/raw_ostream.h"
+#include <cstdint>
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+using namespace llvm;
+using namespace mlir;
+using namespace mlir::python;
+using namespace mlir::python::adaptors;
+
+static void populateDialectNvgpuSubmodule(const pybind11::module &m) {
+ auto nvgpuTensorMapDescriptorType = mlir_type_subclass(
+ m, "TensorMapDescriptorType", mlirTypeIsANVGPUTensorMapDescriptorType);
+
+ nvgpuTensorMapDescriptorType.def_classmethod(
+ "get",
+ [](py::object cls, MlirType tensorType, int swizzle, int l2promo, int oob,
+ int interleave, MlirContext ctx) {
+ return cls(mlirNVGPUTensorMapDescriptorTypeGet(
+ ctx, tensorType, swizzle, l2promo, oob, interleave));
+ },
+ "Gets an instance of RangeType in the same context as the provided "
+ "element type.",
+ py::arg("cls"), py::arg("tensor_type"), py::arg("swizzle"),
+ py::arg("l2promo"), py::arg("oob"), py::arg("interleave"),
+ py::arg("ctx") = py::none());
+}
+
+PYBIND11_MODULE(_mlirDialectsNvgpu, m) {
+ m.doc() = "MLIR NVGPU dialect.";
+
+ populateDialectNvgpuSubmodule(m);
+}
diff --git a/mlir/lib/CAPI/Dialect/NVGPU.cpp b/mlir/lib/CAPI/Dialect/NVGPU.cpp
index 02d10954a03776..7536a525e98778 100644
--- a/mlir/lib/CAPI/Dialect/NVGPU.cpp
+++ b/mlir/lib/CAPI/Dialect/NVGPU.cpp
@@ -9,5 +9,23 @@
#include "mlir-c/Dialect/NVGPU.h"
#include "mlir/CAPI/Registration.h"
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
+#include "mlir/IR/BuiltinTypes.h"
+
+using namespace mlir;
+using namespace mlir::nvgpu;
MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(NVGPU, nvgpu, mlir::nvgpu::NVGPUDialect)
+
+bool mlirTypeIsANVGPUTensorMapDescriptorType(MlirType type) {
+ return isa<nvgpu::TensorMapDescriptorType>(unwrap(type));
+}
+
+MlirType mlirNVGPUTensorMapDescriptorTypeGet(MlirContext ctx,
+ MlirType tensorType, int swizzle,
+ int l2promo, int oob,
+ int interleave) {
+ return wrap(nvgpu::TensorMapDescriptorType::get(
+ unwrap(ctx), cast<MemRefType>(unwrap(tensorType)),
+ TensorMapSwizzleKind(swizzle), TensorMapL2PromoKind(l2promo),
+ TensorMapOOBKind(oob), TensorMapInterleaveKind(interleave)));
+}
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index c27ee688a04087..0a2dc0754c09d0 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -524,6 +524,19 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
MLIRCAPIQuant
)
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
+ MODULE_NAME _mlirDialectsNvgpu
+ ADD_TO_PARENT MLIRPythonSources.Dialects.nvgpu
+ ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ SOURCES
+ DialectNVGPU.cpp
+ PRIVATE_LINK_LIBS
+ LLVMSupport
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIIR
+ MLIRCAPINVGPU
+)
+
declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
MODULE_NAME _mlirDialectsPDL
ADD_TO_PARENT MLIRPythonSources.Dialects.pdl
diff --git a/mlir/python/mlir/dialects/nvgpu.py b/mlir/python/mlir/dialects/nvgpu.py
index 2f6993b768ca53..e93fc956a10dd5 100644
--- a/mlir/python/mlir/dialects/nvgpu.py
+++ b/mlir/python/mlir/dialects/nvgpu.py
@@ -4,3 +4,4 @@
from ._nvgpu_ops_gen import *
from ._nvgpu_enum_gen import *
+from .._mlir_libs._mlirDialectsNvgpu import *
\ No newline at end of file
diff --git a/mlir/test/python/dialects/nvgpu.py b/mlir/test/python/dialects/nvgpu.py
index 3158388f0e6869..2b2fec587352af 100644
--- a/mlir/test/python/dialects/nvgpu.py
+++ b/mlir/test/python/dialects/nvgpu.py
@@ -14,6 +14,19 @@ def constructAndPrintInModule(f):
print(module)
return f
+# CHECK-LABEL: testTypes
+ at constructAndPrintInModule
+def testTypes():
+ tensorType = MemRefType.get((128,64), F16Type.get(), memory_space=Attribute.parse("3"))
+ # CHECK: !nvgpu.tensormap.descriptor<tensor = memref<128x64xf16, 3>, swizzle = swizzle_128b, l2promo = l2promo_256b, oob = nan, interleave = none>
+ tma_desc = nvgpu.TensorMapDescriptorType.get(
+ tensorType,
+ nvgpu.TensorMapSwizzleKind.SWIZZLE_128B,
+ nvgpu.TensorMapL2PromoKind.L2PROMO_256B,
+ nvgpu.TensorMapOOBKind.OOB_NAN,
+ nvgpu.TensorMapInterleaveKind.INTERLEAVE_NONE)
+ print(tma_desc)
+
# CHECK-LABEL: testSmoke
@constructAndPrintInModule
>From 6a2f275dc1a634ebbd167b071d2d36b56e5252a3 Mon Sep 17 00:00:00 2001
From: grypp <guray.ozen at gmail.com>
Date: Tue, 16 Apr 2024 09:20:32 +0000
Subject: [PATCH 2/3] address comments
---
mlir/lib/Bindings/Python/DialectNVGPU.cpp | 9 ++-------
mlir/python/mlir/dialects/nvgpu.py | 2 +-
mlir/test/python/dialects/nvgpu.py | 16 ++++++++++------
3 files changed, 13 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Bindings/Python/DialectNVGPU.cpp b/mlir/lib/Bindings/Python/DialectNVGPU.cpp
index d0f7b2500085d1..c6a5c2ab2f1cc7 100644
--- a/mlir/lib/Bindings/Python/DialectNVGPU.cpp
+++ b/mlir/lib/Bindings/Python/DialectNVGPU.cpp
@@ -1,4 +1,4 @@
-//===- DialectLinalg.cpp - Pybind module for Nvgpu dialect API support --===//
+//===--- DialectNvgpu.cpp - Pybind module for Nvgpu dialect API support ---===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@@ -6,13 +6,9 @@
//
//===----------------------------------------------------------------------===//
-#include "IRModule.h"
#include "mlir-c/Dialect/NVGPU.h"
#include "mlir-c/IR.h"
#include "mlir/Bindings/Python/PybindAdaptors.h"
-#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
-#include "llvm/Support/raw_ostream.h"
-#include <cstdint>
#include <pybind11/pybind11.h>
namespace py = pybind11;
@@ -32,8 +28,7 @@ static void populateDialectNvgpuSubmodule(const pybind11::module &m) {
return cls(mlirNVGPUTensorMapDescriptorTypeGet(
ctx, tensorType, swizzle, l2promo, oob, interleave));
},
- "Gets an instance of RangeType in the same context as the provided "
- "element type.",
+ "Gets an instance of TensorMapDescriptorType in the same context",
py::arg("cls"), py::arg("tensor_type"), py::arg("swizzle"),
py::arg("l2promo"), py::arg("oob"), py::arg("interleave"),
py::arg("ctx") = py::none());
diff --git a/mlir/python/mlir/dialects/nvgpu.py b/mlir/python/mlir/dialects/nvgpu.py
index e93fc956a10dd5..e19bf610ea33a6 100644
--- a/mlir/python/mlir/dialects/nvgpu.py
+++ b/mlir/python/mlir/dialects/nvgpu.py
@@ -4,4 +4,4 @@
from ._nvgpu_ops_gen import *
from ._nvgpu_enum_gen import *
-from .._mlir_libs._mlirDialectsNvgpu import *
\ No newline at end of file
+from .._mlir_libs._mlirDialectsNvgpu import *
diff --git a/mlir/test/python/dialects/nvgpu.py b/mlir/test/python/dialects/nvgpu.py
index 2b2fec587352af..b3edddae9e610c 100644
--- a/mlir/test/python/dialects/nvgpu.py
+++ b/mlir/test/python/dialects/nvgpu.py
@@ -14,17 +14,21 @@ def constructAndPrintInModule(f):
print(module)
return f
+
# CHECK-LABEL: testTypes
@constructAndPrintInModule
def testTypes():
- tensorType = MemRefType.get((128,64), F16Type.get(), memory_space=Attribute.parse("3"))
+ tensorType = MemRefType.get(
+ (128, 64), F16Type.get(), memory_space=Attribute.parse("3")
+ )
# CHECK: !nvgpu.tensormap.descriptor<tensor = memref<128x64xf16, 3>, swizzle = swizzle_128b, l2promo = l2promo_256b, oob = nan, interleave = none>
tma_desc = nvgpu.TensorMapDescriptorType.get(
- tensorType,
- nvgpu.TensorMapSwizzleKind.SWIZZLE_128B,
- nvgpu.TensorMapL2PromoKind.L2PROMO_256B,
- nvgpu.TensorMapOOBKind.OOB_NAN,
- nvgpu.TensorMapInterleaveKind.INTERLEAVE_NONE)
+ tensorType,
+ nvgpu.TensorMapSwizzleKind.SWIZZLE_128B,
+ nvgpu.TensorMapL2PromoKind.L2PROMO_256B,
+ nvgpu.TensorMapOOBKind.OOB_NAN,
+ nvgpu.TensorMapInterleaveKind.INTERLEAVE_NONE,
+ )
print(tma_desc)
>From e19f251d974aeaad08c8b89bae928f596dd43df6 Mon Sep 17 00:00:00 2001
From: grypp <guray.ozen at gmail.com>
Date: Wed, 17 Apr 2024 08:17:39 +0000
Subject: [PATCH 3/3] fix names
---
mlir/include/mlir-c/Dialect/NVGPU.h | 4 ++--
mlir/lib/Bindings/Python/DialectNVGPU.cpp | 8 ++++----
mlir/lib/CAPI/Dialect/NVGPU.cpp | 10 +++++-----
mlir/test/python/dialects/nvgpu.py | 4 ++--
4 files changed, 13 insertions(+), 13 deletions(-)
diff --git a/mlir/include/mlir-c/Dialect/NVGPU.h b/mlir/include/mlir-c/Dialect/NVGPU.h
index 143284ee32b5ae..e58015a4a3421a 100644
--- a/mlir/include/mlir-c/Dialect/NVGPU.h
+++ b/mlir/include/mlir-c/Dialect/NVGPU.h
@@ -26,8 +26,8 @@ MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(NVGPU, nvgpu);
MLIR_CAPI_EXPORTED bool mlirTypeIsANVGPUTensorMapDescriptorType(MlirType type);
MLIR_CAPI_EXPORTED MlirType mlirNVGPUTensorMapDescriptorTypeGet(
- MlirContext ctx, MlirType tensorType, int swizzle, int l2promo, int oob,
- int interleave);
+ MlirContext ctx, MlirType tensorMemrefType, int swizzle, int l2promo,
+ int oobFill, int interleave);
#ifdef __cplusplus
}
diff --git a/mlir/lib/Bindings/Python/DialectNVGPU.cpp b/mlir/lib/Bindings/Python/DialectNVGPU.cpp
index c6a5c2ab2f1cc7..341e4d55bcf219 100644
--- a/mlir/lib/Bindings/Python/DialectNVGPU.cpp
+++ b/mlir/lib/Bindings/Python/DialectNVGPU.cpp
@@ -23,14 +23,14 @@ static void populateDialectNvgpuSubmodule(const pybind11::module &m) {
nvgpuTensorMapDescriptorType.def_classmethod(
"get",
- [](py::object cls, MlirType tensorType, int swizzle, int l2promo, int oob,
- int interleave, MlirContext ctx) {
+ [](py::object cls, MlirType tensorMemrefType, int swizzle, int l2promo,
+ int oobFill, int interleave, MlirContext ctx) {
return cls(mlirNVGPUTensorMapDescriptorTypeGet(
- ctx, tensorType, swizzle, l2promo, oob, interleave));
+ ctx, tensorMemrefType, swizzle, l2promo, oobFill, interleave));
},
"Gets an instance of TensorMapDescriptorType in the same context",
py::arg("cls"), py::arg("tensor_type"), py::arg("swizzle"),
- py::arg("l2promo"), py::arg("oob"), py::arg("interleave"),
+ py::arg("l2promo"), py::arg("oob_fill"), py::arg("interleave"),
py::arg("ctx") = py::none());
}
diff --git a/mlir/lib/CAPI/Dialect/NVGPU.cpp b/mlir/lib/CAPI/Dialect/NVGPU.cpp
index 7536a525e98778..e6da529e1b6b5f 100644
--- a/mlir/lib/CAPI/Dialect/NVGPU.cpp
+++ b/mlir/lib/CAPI/Dialect/NVGPU.cpp
@@ -21,11 +21,11 @@ bool mlirTypeIsANVGPUTensorMapDescriptorType(MlirType type) {
}
MlirType mlirNVGPUTensorMapDescriptorTypeGet(MlirContext ctx,
- MlirType tensorType, int swizzle,
- int l2promo, int oob,
- int interleave) {
+ MlirType tensorMemrefType,
+ int swizzle, int l2promo,
+ int oobFill, int interleave) {
return wrap(nvgpu::TensorMapDescriptorType::get(
- unwrap(ctx), cast<MemRefType>(unwrap(tensorType)),
+ unwrap(ctx), cast<MemRefType>(unwrap(tensorMemrefType)),
TensorMapSwizzleKind(swizzle), TensorMapL2PromoKind(l2promo),
- TensorMapOOBKind(oob), TensorMapInterleaveKind(interleave)));
+ TensorMapOOBKind(oobFill), TensorMapInterleaveKind(interleave)));
}
diff --git a/mlir/test/python/dialects/nvgpu.py b/mlir/test/python/dialects/nvgpu.py
index b3edddae9e610c..6df32bdd3c2739 100644
--- a/mlir/test/python/dialects/nvgpu.py
+++ b/mlir/test/python/dialects/nvgpu.py
@@ -18,12 +18,12 @@ def constructAndPrintInModule(f):
# CHECK-LABEL: testTypes
@constructAndPrintInModule
def testTypes():
- tensorType = MemRefType.get(
+ tensorMemrefType = MemRefType.get(
(128, 64), F16Type.get(), memory_space=Attribute.parse("3")
)
# CHECK: !nvgpu.tensormap.descriptor<tensor = memref<128x64xf16, 3>, swizzle = swizzle_128b, l2promo = l2promo_256b, oob = nan, interleave = none>
tma_desc = nvgpu.TensorMapDescriptorType.get(
- tensorType,
+ tensorMemrefType,
nvgpu.TensorMapSwizzleKind.SWIZZLE_128B,
nvgpu.TensorMapL2PromoKind.L2PROMO_256B,
nvgpu.TensorMapOOBKind.OOB_NAN,
More information about the Mlir-commits
mailing list