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/Dialect/SparseTensor.h" 10 #include "mlir-c/IR.h" 11 #include "mlir/Bindings/Python/PybindAdaptors.h" 12 #include <optional> 13 14 namespace py = pybind11; 15 using namespace llvm; 16 using namespace mlir; 17 using namespace mlir::python::adaptors; 18 19 static void populateDialectSparseTensorSubmodule(const py::module &m) { 20 py::enum_<MlirSparseTensorDimLevelType>(m, "DimLevelType", py::module_local()) 21 .value("dense", MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE) 22 .value("compressed24", MLIR_SPARSE_TENSOR_DIM_LEVEL_TWO_OUT_OF_FOUR) 23 .value("compressed", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED) 24 .value("compressed_nu", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU) 25 .value("compressed_no", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NO) 26 .value("compressed_nu_no", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU_NO) 27 .value("singleton", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON) 28 .value("singleton_nu", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU) 29 .value("singleton_no", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NO) 30 .value("singleton_nu_no", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU_NO) 31 .value("loose_compressed", MLIR_SPARSE_TENSOR_DIM_LEVEL_LOOSE_COMPRESSED) 32 .value("loose_compressed_nu", 33 MLIR_SPARSE_TENSOR_DIM_LEVEL_LOOSE_COMPRESSED_NU) 34 .value("loose_compressed_no", 35 MLIR_SPARSE_TENSOR_DIM_LEVEL_LOOSE_COMPRESSED_NO) 36 .value("loose_compressed_nu_no", 37 MLIR_SPARSE_TENSOR_DIM_LEVEL_LOOSE_COMPRESSED_NU_NO); 38 39 mlir_attribute_subclass(m, "EncodingAttr", 40 mlirAttributeIsASparseTensorEncodingAttr) 41 .def_classmethod( 42 "get", 43 [](py::object cls, std::vector<MlirSparseTensorDimLevelType> lvlTypes, 44 std::optional<MlirAffineMap> dimToLvl, 45 std::optional<MlirAffineMap> lvlToDim, int posWidth, int crdWidth, 46 MlirContext context) { 47 return cls(mlirSparseTensorEncodingAttrGet( 48 context, lvlTypes.size(), lvlTypes.data(), 49 dimToLvl ? *dimToLvl : MlirAffineMap{nullptr}, 50 lvlToDim ? *lvlToDim : MlirAffineMap{nullptr}, posWidth, 51 crdWidth)); 52 }, 53 py::arg("cls"), py::arg("lvl_types"), py::arg("dim_to_lvl"), 54 py::arg("lvl_to_dim"), py::arg("pos_width"), py::arg("crd_width"), 55 py::arg("context") = py::none(), 56 "Gets a sparse_tensor.encoding from parameters.") 57 .def_property_readonly( 58 "lvl_types", 59 [](MlirAttribute self) { 60 const int lvlRank = mlirSparseTensorEncodingGetLvlRank(self); 61 std::vector<MlirSparseTensorDimLevelType> ret; 62 ret.reserve(lvlRank); 63 for (int l = 0; l < lvlRank; ++l) 64 ret.push_back(mlirSparseTensorEncodingAttrGetLvlType(self, l)); 65 return ret; 66 }) 67 .def_property_readonly( 68 "dim_to_lvl", 69 [](MlirAttribute self) -> std::optional<MlirAffineMap> { 70 MlirAffineMap ret = mlirSparseTensorEncodingAttrGetDimToLvl(self); 71 if (mlirAffineMapIsNull(ret)) 72 return {}; 73 return ret; 74 }) 75 .def_property_readonly( 76 "lvl_to_dim", 77 [](MlirAttribute self) -> std::optional<MlirAffineMap> { 78 MlirAffineMap ret = mlirSparseTensorEncodingAttrGetLvlToDim(self); 79 if (mlirAffineMapIsNull(ret)) 80 return {}; 81 return ret; 82 }) 83 .def_property_readonly("pos_width", 84 mlirSparseTensorEncodingAttrGetPosWidth) 85 .def_property_readonly("crd_width", 86 mlirSparseTensorEncodingAttrGetCrdWidth); 87 } 88 89 PYBIND11_MODULE(_mlirDialectsSparseTensor, m) { 90 m.doc() = "MLIR SparseTensor dialect."; 91 populateDialectSparseTensorSubmodule(m); 92 } 93