149 lines · cpp
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-exception6//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/Nanobind.h"16#include "mlir/Bindings/Python/NanobindAdaptors.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 [](const nb::object &cls,42 std::vector<MlirSparseTensorLevelType> lvlTypes,43 std::optional<MlirAffineMap> dimToLvl,44 std::optional<MlirAffineMap> lvlToDim, int posWidth, int crdWidth,45 std::optional<MlirAttribute> explicitVal,46 std::optional<MlirAttribute> implicitVal, 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, explicitVal ? *explicitVal : MlirAttribute{nullptr},52 implicitVal ? *implicitVal : MlirAttribute{nullptr}));53 },54 nb::arg("cls"), nb::arg("lvl_types"), nb::arg("dim_to_lvl").none(),55 nb::arg("lvl_to_dim").none(), nb::arg("pos_width"),56 nb::arg("crd_width"), nb::arg("explicit_val") = nb::none(),57 nb::arg("implicit_val") = nb::none(), nb::arg("context") = nb::none(),58 "Gets a sparse_tensor.encoding from parameters.")59 .def_classmethod(60 "build_level_type",61 [](const 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