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