[Mlir-commits] [mlir] [mlir][py] Add NVGPU's `TensorMapDescriptorType` in py bindings (PR #88855)

Maksim Levental llvmlistbot at llvm.org
Tue Apr 16 10:57:35 PDT 2024


================
@@ -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 tensorType, int swizzle, int l2promo, int oob,
+         int interleave, MlirContext ctx) {
+        return cls(mlirNVGPUTensorMapDescriptorTypeGet(
+            ctx, tensorType, swizzle, l2promo, oob, 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("ctx") = py::none());
----------------
makslevental wrote:

If you wanted to be ambitious you could define a couple properties so this type isn't completely opaque, like this https://github.com/llvm/llvm-project/blob/681eacc1b670fd7137d8677fef6fc76c6e37dca9/mlir/lib/Bindings/Python/DialectTransform.cpp#L88, but if there's not a current need for introspection we can leave it till someone asks/needs.

https://github.com/llvm/llvm-project/pull/88855


More information about the Mlir-commits mailing list