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 13 namespace py = pybind11; 14 using namespace llvm; 15 using namespace mlir; 16 using namespace mlir::python::adaptors; 17 18 static void populateDialectSparseTensorSubmodule(const py::module &m) { 19 py::enum_<MlirSparseTensorDimLevelType>(m, "DimLevelType", py::module_local()) 20 .value("dense", MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE) 21 .value("compressed", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED) 22 .value("compressed-nu", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU) 23 .value("compressed-no", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NO) 24 .value("compressed-nu-no", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU_NO) 25 .value("singleton", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON) 26 .value("singleton-nu", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU) 27 .value("singleton-no", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NO) 28 .value("singleton-nu-no", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU_NO); 29 30 mlir_attribute_subclass(m, "EncodingAttr", 31 mlirAttributeIsASparseTensorEncodingAttr) 32 .def_classmethod( 33 "get", 34 [](py::object cls, 35 std::vector<MlirSparseTensorDimLevelType> dimLevelTypes, 36 llvm::Optional<MlirAffineMap> dimOrdering, 37 llvm::Optional<MlirAffineMap> higherOrdering, int pointerBitWidth, 38 int indexBitWidth, MlirContext context) { 39 return cls(mlirSparseTensorEncodingAttrGet( 40 context, dimLevelTypes.size(), dimLevelTypes.data(), 41 dimOrdering ? *dimOrdering : MlirAffineMap{nullptr}, 42 higherOrdering ? *higherOrdering : MlirAffineMap{nullptr}, 43 pointerBitWidth, indexBitWidth)); 44 }, 45 py::arg("cls"), py::arg("dim_level_types"), py::arg("dim_ordering"), 46 py::arg("higher_ordering"), py::arg("pointer_bit_width"), 47 py::arg("index_bit_width"), py::arg("context") = py::none(), 48 "Gets a sparse_tensor.encoding from parameters.") 49 .def_property_readonly( 50 "dim_level_types", 51 [](MlirAttribute self) { 52 std::vector<MlirSparseTensorDimLevelType> ret; 53 for (int i = 0, 54 e = mlirSparseTensorEncodingGetNumDimLevelTypes(self); 55 i < e; ++i) 56 ret.push_back( 57 mlirSparseTensorEncodingAttrGetDimLevelType(self, i)); 58 return ret; 59 }) 60 .def_property_readonly( 61 "dim_ordering", 62 [](MlirAttribute self) -> llvm::Optional<MlirAffineMap> { 63 MlirAffineMap ret = 64 mlirSparseTensorEncodingAttrGetDimOrdering(self); 65 if (mlirAffineMapIsNull(ret)) 66 return {}; 67 return ret; 68 }) 69 .def_property_readonly( 70 "higher_ordering", 71 [](MlirAttribute self) -> llvm::Optional<MlirAffineMap> { 72 MlirAffineMap ret = 73 mlirSparseTensorEncodingAttrGetHigherOrdering(self); 74 if (mlirAffineMapIsNull(ret)) 75 return {}; 76 return ret; 77 }) 78 .def_property_readonly( 79 "pointer_bit_width", 80 [](MlirAttribute self) { 81 return mlirSparseTensorEncodingAttrGetPointerBitWidth(self); 82 }) 83 .def_property_readonly("index_bit_width", [](MlirAttribute self) { 84 return mlirSparseTensorEncodingAttrGetIndexBitWidth(self); 85 }); 86 } 87 88 PYBIND11_MODULE(_mlirDialectsSparseTensor, m) { 89 m.doc() = "MLIR SparseTensor dialect."; 90 populateDialectSparseTensorSubmodule(m); 91 } 92