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 <optional>
10#include <vector>
11
12#include "mlir-c/AffineMap.h"
13#include "mlir-c/Dialect/SparseTensor.h"
14#include "mlir-c/IR.h"
15#include "mlir/Bindings/Python/NanobindAdaptors.h"
16#include "mlir/Bindings/Python/Nanobind.h"
17
18namespace nb = nanobind;
19using namespace llvm;
20using namespace mlir;
21using namespace mlir::python::nanobind_adaptors;
22
23static void populateDialectSparseTensorSubmodule(const nb::module_ &m) {
24 nb::enum_<MlirSparseTensorLevelFormat>(m, "LevelFormat", nb::is_arithmetic(),
25 nb::is_flag())
26 .value("dense", MLIR_SPARSE_TENSOR_LEVEL_DENSE)
27 .value("n_out_of_m", MLIR_SPARSE_TENSOR_LEVEL_N_OUT_OF_M)
28 .value("compressed", MLIR_SPARSE_TENSOR_LEVEL_COMPRESSED)
29 .value("singleton", MLIR_SPARSE_TENSOR_LEVEL_SINGLETON)
30 .value("loose_compressed", MLIR_SPARSE_TENSOR_LEVEL_LOOSE_COMPRESSED);
31
32 nb::enum_<MlirSparseTensorLevelPropertyNondefault>(m, "LevelProperty")
33 .value("non_ordered", MLIR_SPARSE_PROPERTY_NON_ORDERED)
34 .value("non_unique", MLIR_SPARSE_PROPERTY_NON_UNIQUE)
35 .value("soa", MLIR_SPARSE_PROPERTY_SOA);
36
37 mlir_attribute_subclass(m, "EncodingAttr",
38 mlirAttributeIsASparseTensorEncodingAttr)
39 .def_classmethod(
40 "get",
41 [](nb::object cls, std::vector<MlirSparseTensorLevelType> lvlTypes,
42 std::optional<MlirAffineMap> dimToLvl,
43 std::optional<MlirAffineMap> lvlToDim, int posWidth, int crdWidth,
44 std::optional<MlirAttribute> explicitVal,
45 std::optional<MlirAttribute> implicitVal, 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, explicitVal ? *explicitVal : MlirAttribute{nullptr},
51 implicitVal ? *implicitVal : MlirAttribute{nullptr}));
52 },
53 nb::arg("cls"), nb::arg("lvl_types"), nb::arg("dim_to_lvl").none(),
54 nb::arg("lvl_to_dim").none(), nb::arg("pos_width"),
55 nb::arg("crd_width"), nb::arg("explicit_val").none() = nb::none(),
56 nb::arg("implicit_val").none() = nb::none(),
57 nb::arg("context").none() = nb::none(),
58 "Gets a sparse_tensor.encoding from parameters.")
59 .def_classmethod(
60 "build_level_type",
61 [](nb::object cls, MlirSparseTensorLevelFormat lvlFmt,
62 const std::vector<MlirSparseTensorLevelPropertyNondefault>
63 &properties,
64 unsigned n, unsigned m) {
65 return mlirSparseTensorEncodingAttrBuildLvlType(
66 lvlFmt, properties.data(), properties.size(), n, m);
67 },
68 nb::arg("cls"), nb::arg("lvl_fmt"),
69 nb::arg("properties") =
70 std::vector<MlirSparseTensorLevelPropertyNondefault>(),
71 nb::arg("n") = 0, nb::arg("m") = 0,
72 "Builds a sparse_tensor.encoding.level_type from parameters.")
73 .def_property_readonly(
74 "lvl_types",
75 [](MlirAttribute self) {
76 const int lvlRank = mlirSparseTensorEncodingGetLvlRank(self);
77 std::vector<MlirSparseTensorLevelType> ret;
78 ret.reserve(lvlRank);
79 for (int l = 0; l < lvlRank; ++l)
80 ret.push_back(mlirSparseTensorEncodingAttrGetLvlType(self, l));
81 return ret;
82 })
83 .def_property_readonly(
84 "dim_to_lvl",
85 [](MlirAttribute self) -> std::optional<MlirAffineMap> {
86 MlirAffineMap ret = mlirSparseTensorEncodingAttrGetDimToLvl(self);
87 if (mlirAffineMapIsNull(ret))
88 return {};
89 return ret;
90 })
91 .def_property_readonly(
92 "lvl_to_dim",
93 [](MlirAttribute self) -> std::optional<MlirAffineMap> {
94 MlirAffineMap ret = mlirSparseTensorEncodingAttrGetLvlToDim(self);
95 if (mlirAffineMapIsNull(ret))
96 return {};
97 return ret;
98 })
99 .def_property_readonly("pos_width",
100 mlirSparseTensorEncodingAttrGetPosWidth)
101 .def_property_readonly("crd_width",
102 mlirSparseTensorEncodingAttrGetCrdWidth)
103 .def_property_readonly(
104 "explicit_val",
105 [](MlirAttribute self) -> std::optional<MlirAttribute> {
106 MlirAttribute ret =
107 mlirSparseTensorEncodingAttrGetExplicitVal(self);
108 if (mlirAttributeIsNull(ret))
109 return {};
110 return ret;
111 })
112 .def_property_readonly(
113 "implicit_val",
114 [](MlirAttribute self) -> std::optional<MlirAttribute> {
115 MlirAttribute ret =
116 mlirSparseTensorEncodingAttrGetImplicitVal(self);
117 if (mlirAttributeIsNull(ret))
118 return {};
119 return ret;
120 })
121 .def_property_readonly(
122 "structured_n",
123 [](MlirAttribute self) -> unsigned {
124 const int lvlRank = mlirSparseTensorEncodingGetLvlRank(self);
125 return mlirSparseTensorEncodingAttrGetStructuredN(
126 mlirSparseTensorEncodingAttrGetLvlType(self, lvlRank - 1));
127 })
128 .def_property_readonly(
129 "structured_m",
130 [](MlirAttribute self) -> unsigned {
131 const int lvlRank = mlirSparseTensorEncodingGetLvlRank(self);
132 return mlirSparseTensorEncodingAttrGetStructuredM(
133 mlirSparseTensorEncodingAttrGetLvlType(self, lvlRank - 1));
134 })
135 .def_property_readonly("lvl_formats_enum", [](MlirAttribute self) {
136 const int lvlRank = mlirSparseTensorEncodingGetLvlRank(self);
137 std::vector<MlirSparseTensorLevelFormat> ret;
138 ret.reserve(lvlRank);
139 for (int l = 0; l < lvlRank; l++)
140 ret.push_back(mlirSparseTensorEncodingAttrGetLvlFmt(self, l));
141 return ret;
142 });
143}
144
145NB_MODULE(_mlirDialectsSparseTensor, m) {
146 m.doc() = "MLIR SparseTensor dialect.";
147 populateDialectSparseTensorSubmodule(m);
148}
149

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