1//===- DialectSparseTensor.cpp - 'sparse_tensor' dialect submodule --------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#include "mlir-c/AffineMap.h"
10#include "mlir-c/Dialect/SparseTensor.h"
11#include "mlir-c/IR.h"
12#include "mlir/Bindings/Python/PybindAdaptors.h"
13#include <optional>
14#include <pybind11/cast.h>
15#include <pybind11/detail/common.h>
16#include <pybind11/pybind11.h>
17#include <pybind11/pytypes.h>
18#include <vector>
19
20namespace py = pybind11;
21using namespace llvm;
22using namespace mlir;
23using namespace mlir::python::adaptors;
24
25static void populateDialectSparseTensorSubmodule(const py::module &m) {
26 py::enum_<MlirSparseTensorLevelFormat>(m, "LevelFormat", py::module_local())
27 .value("dense", MLIR_SPARSE_TENSOR_LEVEL_DENSE)
28 .value("n_out_of_m", MLIR_SPARSE_TENSOR_LEVEL_N_OUT_OF_M)
29 .value("compressed", MLIR_SPARSE_TENSOR_LEVEL_COMPRESSED)
30 .value("singleton", MLIR_SPARSE_TENSOR_LEVEL_SINGLETON)
31 .value("loose_compressed", MLIR_SPARSE_TENSOR_LEVEL_LOOSE_COMPRESSED);
32
33 py::enum_<MlirSparseTensorLevelPropertyNondefault>(m, "LevelProperty",
34 py::module_local())
35 .value("non_ordered", MLIR_SPARSE_PROPERTY_NON_ORDERED)
36 .value("non_unique", MLIR_SPARSE_PROPERTY_NON_UNIQUE);
37
38 mlir_attribute_subclass(m, "EncodingAttr",
39 mlirAttributeIsASparseTensorEncodingAttr)
40 .def_classmethod(
41 "get",
42 [](py::object cls, std::vector<MlirSparseTensorLevelType> lvlTypes,
43 std::optional<MlirAffineMap> dimToLvl,
44 std::optional<MlirAffineMap> lvlToDim, int posWidth, int crdWidth,
45 MlirContext context) {
46 return cls(mlirSparseTensorEncodingAttrGet(
47 context, lvlTypes.size(), lvlTypes.data(),
48 dimToLvl ? *dimToLvl : MlirAffineMap{nullptr},
49 lvlToDim ? *lvlToDim : MlirAffineMap{nullptr}, posWidth,
50 crdWidth));
51 },
52 py::arg("cls"), py::arg("lvl_types"), py::arg("dim_to_lvl"),
53 py::arg("lvl_to_dim"), py::arg("pos_width"), py::arg("crd_width"),
54 py::arg("context") = py::none(),
55 "Gets a sparse_tensor.encoding from parameters.")
56 .def_classmethod(
57 "build_level_type",
58 [](py::object cls, MlirSparseTensorLevelFormat lvlFmt,
59 const std::vector<MlirSparseTensorLevelPropertyNondefault>
60 &properties,
61 unsigned n, unsigned m) {
62 return mlirSparseTensorEncodingAttrBuildLvlType(
63 lvlFmt, properties.data(), properties.size(), n, m);
64 },
65 py::arg("cls"), py::arg("lvl_fmt"),
66 py::arg("properties") =
67 std::vector<MlirSparseTensorLevelPropertyNondefault>(),
68 py::arg("n") = 0, py::arg("m") = 0,
69 "Builds a sparse_tensor.encoding.level_type from parameters.")
70 .def_property_readonly(
71 "lvl_types",
72 [](MlirAttribute self) {
73 const int lvlRank = mlirSparseTensorEncodingGetLvlRank(self);
74 std::vector<MlirSparseTensorLevelType> ret;
75 ret.reserve(lvlRank);
76 for (int l = 0; l < lvlRank; ++l)
77 ret.push_back(mlirSparseTensorEncodingAttrGetLvlType(self, l));
78 return ret;
79 })
80 .def_property_readonly(
81 "dim_to_lvl",
82 [](MlirAttribute self) -> std::optional<MlirAffineMap> {
83 MlirAffineMap ret = mlirSparseTensorEncodingAttrGetDimToLvl(self);
84 if (mlirAffineMapIsNull(ret))
85 return {};
86 return ret;
87 })
88 .def_property_readonly(
89 "lvl_to_dim",
90 [](MlirAttribute self) -> std::optional<MlirAffineMap> {
91 MlirAffineMap ret = mlirSparseTensorEncodingAttrGetLvlToDim(self);
92 if (mlirAffineMapIsNull(ret))
93 return {};
94 return ret;
95 })
96 .def_property_readonly("pos_width",
97 mlirSparseTensorEncodingAttrGetPosWidth)
98 .def_property_readonly("crd_width",
99 mlirSparseTensorEncodingAttrGetCrdWidth)
100 .def_property_readonly(
101 "structured_n",
102 [](MlirAttribute self) -> unsigned {
103 const int lvlRank = mlirSparseTensorEncodingGetLvlRank(self);
104 return mlirSparseTensorEncodingAttrGetStructuredN(
105 mlirSparseTensorEncodingAttrGetLvlType(self, lvlRank - 1));
106 })
107 .def_property_readonly(
108 "structured_m",
109 [](MlirAttribute self) -> unsigned {
110 const int lvlRank = mlirSparseTensorEncodingGetLvlRank(self);
111 return mlirSparseTensorEncodingAttrGetStructuredM(
112 mlirSparseTensorEncodingAttrGetLvlType(self, lvlRank - 1));
113 })
114 .def_property_readonly("lvl_formats_enum", [](MlirAttribute self) {
115 const int lvlRank = mlirSparseTensorEncodingGetLvlRank(self);
116 std::vector<MlirSparseTensorLevelFormat> ret;
117 ret.reserve(lvlRank);
118 for (int l = 0; l < lvlRank; l++)
119 ret.push_back(mlirSparseTensorEncodingAttrGetLvlFmt(self, l));
120 return ret;
121 });
122}
123
124PYBIND11_MODULE(_mlirDialectsSparseTensor, m) {
125 m.doc() = "MLIR SparseTensor dialect.";
126 populateDialectSparseTensorSubmodule(m);
127}
128

source code of mlir/lib/Bindings/Python/DialectSparseTensor.cpp