MLIR  19.0.0git
DialectSparseTensor.cpp
Go to the documentation of this file.
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"
11 #include "mlir-c/IR.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 
38  mlir_attribute_subclass(m, "EncodingAttr",
41  "get",
42  [](py::object cls, 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) {
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  py::arg("cls"), py::arg("lvl_types"), py::arg("dim_to_lvl"),
55  py::arg("lvl_to_dim"), py::arg("pos_width"), py::arg("crd_width"),
56  py::arg("explicit_val") = py::none(),
57  py::arg("implicit_val") = py::none(), py::arg("context") = py::none(),
58  "Gets a sparse_tensor.encoding from parameters.")
59  .def_classmethod(
60  "build_level_type",
61  [](py::object cls, MlirSparseTensorLevelFormat lvlFmt,
62  const std::vector<MlirSparseTensorLevelPropertyNondefault>
63  &properties,
64  unsigned n, unsigned m) {
66  lvlFmt, properties.data(), properties.size(), n, m);
67  },
68  py::arg("cls"), py::arg("lvl_fmt"),
69  py::arg("properties") =
70  std::vector<MlirSparseTensorLevelPropertyNondefault>(),
71  py::arg("n") = 0, py::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",
101  .def_property_readonly("crd_width",
104  "explicit_val",
105  [](MlirAttribute self) -> std::optional<MlirAttribute> {
106  MlirAttribute ret =
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 =
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);
126  mlirSparseTensorEncodingAttrGetLvlType(self, lvlRank - 1));
127  })
128  .def_property_readonly(
129  "structured_m",
130  [](MlirAttribute self) -> unsigned {
131  const int lvlRank = mlirSparseTensorEncodingGetLvlRank(self);
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 PYBIND11_MODULE(_mlirDialectsSparseTensor, m) {
146  m.doc() = "MLIR SparseTensor dialect.";
148 }
PYBIND11_MODULE(_mlirDialectsSparseTensor, m)
static void populateDialectSparseTensorSubmodule(const py::module &m)
Creates a custom subclass of mlir.ir.Attribute, implementing a casting constructor and type checking ...
pure_subclass & def_classmethod(const char *name, Func &&f, const Extra &...extra)
pure_subclass & def_property_readonly(const char *name, Func &&f, const Extra &...extra)
static bool mlirAffineMapIsNull(MlirAffineMap affineMap)
Checks whether an affine map is null.
Definition: AffineMap.h:47
MLIR_CAPI_EXPORTED MlirAffineMap mlirSparseTensorEncodingAttrGetDimToLvl(MlirAttribute attr)
Returns the dimension-to-level mapping of the sparse_tensor.encoding attribute.
MLIR_CAPI_EXPORTED MlirAttribute mlirSparseTensorEncodingAttrGetImplicitVal(MlirAttribute attr)
Returns the implicit value of the sparse_tensor.encoding attribute.
MLIR_CAPI_EXPORTED intptr_t mlirSparseTensorEncodingGetLvlRank(MlirAttribute attr)
Returns the level-rank of the sparse_tensor.encoding attribute.
MLIR_CAPI_EXPORTED MlirAttribute mlirSparseTensorEncodingAttrGet(MlirContext ctx, intptr_t lvlRank, MlirSparseTensorLevelType const *lvlTypes, MlirAffineMap dimToLvl, MlirAffineMap lvlTodim, int posWidth, int crdWidth, MlirAttribute explicitVal, MlirAttribute implicitVal)
Creates a sparse_tensor.encoding attribute with the given parameters.
MLIR_CAPI_EXPORTED enum MlirSparseTensorLevelFormat mlirSparseTensorEncodingAttrGetLvlFmt(MlirAttribute attr, intptr_t lvl)
Returns a specified level-format of the sparse_tensor.encoding attribute.
MLIR_CAPI_EXPORTED int mlirSparseTensorEncodingAttrGetCrdWidth(MlirAttribute attr)
Returns the coordinate bitwidth of the sparse_tensor.encoding attribute.
MLIR_CAPI_EXPORTED MlirAffineMap mlirSparseTensorEncodingAttrGetLvlToDim(MlirAttribute attr)
Returns the level-to-dimension mapping of the sparse_tensor.encoding attribute.
MLIR_CAPI_EXPORTED MlirAttribute mlirSparseTensorEncodingAttrGetExplicitVal(MlirAttribute attr)
Returns the explicit value of the sparse_tensor.encoding attribute.
MlirSparseTensorLevelFormat
Definition: SparseTensor.h:30
@ MLIR_SPARSE_TENSOR_LEVEL_N_OUT_OF_M
Definition: SparseTensor.h:36
@ MLIR_SPARSE_TENSOR_LEVEL_COMPRESSED
Definition: SparseTensor.h:33
@ MLIR_SPARSE_TENSOR_LEVEL_DENSE
Definition: SparseTensor.h:31
@ MLIR_SPARSE_TENSOR_LEVEL_LOOSE_COMPRESSED
Definition: SparseTensor.h:35
@ MLIR_SPARSE_TENSOR_LEVEL_SINGLETON
Definition: SparseTensor.h:34
MLIR_CAPI_EXPORTED unsigned mlirSparseTensorEncodingAttrGetStructuredN(MlirSparseTensorLevelType lvlType)
MLIR_CAPI_EXPORTED MlirSparseTensorLevelType mlirSparseTensorEncodingAttrBuildLvlType(enum MlirSparseTensorLevelFormat lvlFmt, const enum MlirSparseTensorLevelPropertyNondefault *properties, unsigned propSize, unsigned n, unsigned m)
MLIR_CAPI_EXPORTED unsigned mlirSparseTensorEncodingAttrGetStructuredM(MlirSparseTensorLevelType lvlType)
MLIR_CAPI_EXPORTED bool mlirAttributeIsASparseTensorEncodingAttr(MlirAttribute attr)
Checks whether the given attribute is a sparse_tensor.encoding attribute.
MLIR_CAPI_EXPORTED MlirSparseTensorLevelType mlirSparseTensorEncodingAttrGetLvlType(MlirAttribute attr, intptr_t lvl)
Returns a specified level-type of the sparse_tensor.encoding attribute.
MLIR_CAPI_EXPORTED int mlirSparseTensorEncodingAttrGetPosWidth(MlirAttribute attr)
Returns the position bitwidth of the sparse_tensor.encoding attribute.
@ MLIR_SPARSE_PROPERTY_NON_UNIQUE
Definition: SparseTensor.h:40
@ MLIR_SPARSE_PROPERTY_NON_ORDERED
Definition: SparseTensor.h:41
static bool mlirAttributeIsNull(MlirAttribute attr)
Checks whether an attribute is null.
Definition: IR.h:1026
Include the generated interface declarations.
Definition: CallGraph.h:229
Include the generated interface declarations.