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 18 namespace nb = nanobind; 19 using namespace llvm; 20 using namespace mlir; 21 using namespace mlir::python::nanobind_adaptors; 22 23 static 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 145 NB_MODULE(_mlirDialectsSparseTensor, m) { 146 m.doc() = "MLIR SparseTensor dialect."; 147 populateDialectSparseTensorSubmodule(m); 148 } 149