[Mlir-commits] [mlir] 4f88c23 - [mlir][py] Add NVGPU's `TensorMapDescriptorType` in py bindings (#88855)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Apr 17 06:59:22 PDT 2024
Author: Guray Ozen
Date: 2024-04-17T15:59:18+02:00
New Revision: 4f88c2311130791cf69da34b743b1b3ba7584a7b
URL: https://github.com/llvm/llvm-project/commit/4f88c2311130791cf69da34b743b1b3ba7584a7b
DIFF: https://github.com/llvm/llvm-project/commit/4f88c2311130791cf69da34b743b1b3ba7584a7b.diff
LOG: [mlir][py] Add NVGPU's `TensorMapDescriptorType` in py bindings (#88855)
This PR adds NVGPU dialects' TensorMapDescriptorType in the py bindings.
This is a follow-up issue from [this
PR](https://github.com/llvm/llvm-project/pull/87153#discussion_r1546193095)
Added:
mlir/lib/Bindings/Python/DialectNVGPU.cpp
Modified:
mlir/include/mlir-c/Dialect/NVGPU.h
mlir/lib/CAPI/Dialect/NVGPU.cpp
mlir/python/CMakeLists.txt
mlir/python/mlir/dialects/nvgpu.py
mlir/test/python/dialects/nvgpu.py
Removed:
################################################################################
diff --git a/mlir/include/mlir-c/Dialect/NVGPU.h b/mlir/include/mlir-c/Dialect/NVGPU.h
index 580d566794c09f..e58015a4a3421a 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 tensorMemrefType, int swizzle, int l2promo,
+ int oobFill, 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..341e4d55bcf219
--- /dev/null
+++ b/mlir/lib/Bindings/Python/DialectNVGPU.cpp
@@ -0,0 +1,41 @@
+//===--- 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.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/NVGPU.h"
+#include "mlir-c/IR.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
+#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 tensorMemrefType, int swizzle, int l2promo,
+ int oobFill, int interleave, MlirContext ctx) {
+ return cls(mlirNVGPUTensorMapDescriptorTypeGet(
+ 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_fill"), 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..e6da529e1b6b5f 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 tensorMemrefType,
+ int swizzle, int l2promo,
+ int oobFill, int interleave) {
+ return wrap(nvgpu::TensorMapDescriptorType::get(
+ unwrap(ctx), cast<MemRefType>(unwrap(tensorMemrefType)),
+ TensorMapSwizzleKind(swizzle), TensorMapL2PromoKind(l2promo),
+ TensorMapOOBKind(oobFill), 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..e19bf610ea33a6 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 *
diff --git a/mlir/test/python/dialects/nvgpu.py b/mlir/test/python/dialects/nvgpu.py
index 3158388f0e6869..6df32bdd3c2739 100644
--- a/mlir/test/python/dialects/nvgpu.py
+++ b/mlir/test/python/dialects/nvgpu.py
@@ -15,6 +15,23 @@ def constructAndPrintInModule(f):
return f
+# CHECK-LABEL: testTypes
+ at constructAndPrintInModule
+def testTypes():
+ 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(
+ tensorMemrefType,
+ nvgpu.TensorMapSwizzleKind.SWIZZLE_128B,
+ nvgpu.TensorMapL2PromoKind.L2PROMO_256B,
+ nvgpu.TensorMapOOBKind.OOB_NAN,
+ nvgpu.TensorMapInterleaveKind.INTERLEAVE_NONE,
+ )
+ print(tma_desc)
+
+
# CHECK-LABEL: testSmoke
@constructAndPrintInModule
def testSmoke():
More information about the Mlir-commits
mailing list