MLIR  18.0.0git
SparseTensorStorageLayout.h
Go to the documentation of this file.
1 //===- SparseTensorStorageLayout.h ------------------------------*- C++ -*-===//
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 // This header file defines utilities for the sparse memory layout.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #ifndef MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSORSTORAGELAYOUT_H_
14 #define MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSORSTORAGELAYOUT_H_
15 
18 
19 namespace mlir {
20 namespace sparse_tensor {
21 
22 ///===----------------------------------------------------------------------===//
23 /// The sparse tensor storage scheme for a tensor is organized as a single
24 /// compound type with the following fields. Note that every memref with `?`
25 /// size actually behaves as a "vector", i.e. the stored size is the capacity
26 /// and the used size resides in the storage_specifier struct.
27 ///
28 /// struct {
29 /// ; per-level l:
30 /// ; if dense:
31 /// <nothing>
32 /// ; if compressed:
33 /// memref<? x pos> positions ; positions for level l
34 /// memref<? x crd> coordinates ; coordinates for level l
35 /// ; if loose-compressed:
36 /// memref<? x pos> positions ; lo/hi position pairs for level l
37 /// memref<? x crd> coordinates ; coordinates for level l
38 /// ; if singleton/2-out-of-4:
39 /// memref<? x crd> coordinates ; coordinates for level l
40 ///
41 /// memref<? x eltType> values ; values
42 ///
43 /// struct sparse_tensor.storage_specifier {
44 /// array<rank x int> lvlSizes ; sizes/cardinalities for each level
45 /// array<n x int> memSizes; ; sizes/lengths for each data memref
46 /// }
47 /// };
48 ///
49 /// In addition, for a "trailing COO region", defined as a compressed level
50 /// followed by one or more singleton levels, the default SOA storage that
51 /// is inherent to the TACO format is optimized into an AOS storage where
52 /// all coordinates of a stored element appear consecutively. In such cases,
53 /// a special operation (sparse_tensor.coordinates_buffer) must be used to
54 /// access the AOS coordinates array. In the code below, the method
55 /// `getCOOStart` is used to find the start of the "trailing COO region".
56 ///
57 /// If the sparse tensor is a slice (produced by `tensor.extract_slice`
58 /// operation), instead of allocating a new sparse tensor for it, it reuses the
59 /// same sets of MemRefs but attaching a additional set of slicing-metadata for
60 /// per-dimension slice offset and stride.
61 ///
62 /// Examples.
63 ///
64 /// #CSR storage of 2-dim matrix yields
65 /// memref<?xindex> ; positions-1
66 /// memref<?xindex> ; coordinates-1
67 /// memref<?xf64> ; values
68 /// struct<(array<2 x i64>, array<3 x i64>)>) ; lvl0, lvl1, 3xsizes
69 ///
70 /// #COO storage of 2-dim matrix yields
71 /// memref<?xindex>, ; positions-0, essentially [0,sz]
72 /// memref<?xindex> ; AOS coordinates storage
73 /// memref<?xf64> ; values
74 /// struct<(array<2 x i64>, array<3 x i64>)>) ; lvl0, lvl1, 3xsizes
75 ///
76 /// Slice on #COO storage of 2-dim matrix yields
77 /// ;; Inherited from the original sparse tensors
78 /// memref<?xindex>, ; positions-0, essentially [0,sz]
79 /// memref<?xindex> ; AOS coordinates storage
80 /// memref<?xf64> ; values
81 /// struct<(array<2 x i64>, array<3 x i64>, ; lvl0, lvl1, 3xsizes
82 /// ;; Extra slicing-metadata
83 /// array<2 x i64>, array<2 x i64>)>) ; dim offset, dim stride.
84 ///
85 ///===----------------------------------------------------------------------===//
86 
87 enum class SparseTensorFieldKind : uint32_t {
88  StorageSpec = 0,
89  PosMemRef = static_cast<uint32_t>(StorageSpecifierKind::PosMemSize),
90  CrdMemRef = static_cast<uint32_t>(StorageSpecifierKind::CrdMemSize),
91  ValMemRef = static_cast<uint32_t>(StorageSpecifierKind::ValMemSize)
92 };
93 
94 inline StorageSpecifierKind toSpecifierKind(SparseTensorFieldKind kind) {
95  assert(kind != SparseTensorFieldKind::StorageSpec);
96  return static_cast<StorageSpecifierKind>(kind);
97 }
98 
99 inline SparseTensorFieldKind toFieldKind(StorageSpecifierKind kind) {
100  assert(kind != StorageSpecifierKind::LvlSize);
101  return static_cast<SparseTensorFieldKind>(kind);
102 }
103 
104 /// The type of field indices. This alias is to help code be more
105 /// self-documenting; unfortunately it is not type-checked, so it only
106 /// provides documentation rather than doing anything to prevent mixups.
107 using FieldIndex = unsigned;
108 
109 /// Provides methods to access fields of a sparse tensor with the given
110 /// encoding.
112 public:
113  explicit StorageLayout(const SparseTensorType &stt)
114  : StorageLayout(stt.getEncoding()) {}
115  explicit StorageLayout(SparseTensorEncodingAttr enc) : enc(enc) {
116  assert(enc);
117  }
118 
119  /// For each field that will be allocated for the given sparse tensor
120  /// encoding, calls the callback with the corresponding field index,
121  /// field kind, level, and level-type (the last two are only for level
122  /// memrefs). The field index always starts with zero and increments
123  /// by one between each callback invocation. Ideally, all other methods
124  /// should rely on this function to query a sparse tensor fields instead
125  /// of relying on ad-hoc index computation.
126  void foreachField(
127  llvm::function_ref<bool(
128  FieldIndex /*fieldIdx*/, SparseTensorFieldKind /*fieldKind*/,
129  Level /*lvl (if applicable)*/, LevelType /*LT (if applicable)*/)>)
130  const;
131 
132  /// Gets the field index for required field.
134  std::optional<Level> lvl) const {
135  return getFieldIndexAndStride(kind, lvl).first;
136  }
137 
138  /// Gets the total number of fields for the given sparse tensor encoding.
139  unsigned getNumFields() const;
140 
141  /// Gets the total number of data fields (coordinate arrays, position
142  /// arrays, and a value array) for the given sparse tensor encoding.
143  unsigned getNumDataFields() const;
144 
145  std::pair<FieldIndex, unsigned>
147  std::optional<Level> lvl) const;
148 
149 private:
150  const SparseTensorEncodingAttr enc;
151 };
152 
153 //
154 // Wrapper functions to invoke StorageLayout-related method.
155 //
156 
157 inline unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc) {
158  return StorageLayout(enc).getNumFields();
159 }
160 
161 inline unsigned getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc) {
162  return StorageLayout(enc).getNumDataFields();
163 }
164 
166  SparseTensorEncodingAttr enc,
168  LevelType)>
169  callback) {
170  return StorageLayout(enc).foreachField(callback);
171 }
172 
174  SparseTensorType,
176  LevelType)>);
177 
178 } // namespace sparse_tensor
179 } // namespace mlir
180 
181 #endif // MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSORSTORAGELAYOUT_H_
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
A wrapper around RankedTensorType, which has three goals:
Provides methods to access fields of a sparse tensor with the given encoding.
unsigned getNumDataFields() const
Gets the total number of data fields (coordinate arrays, position arrays, and a value array) for the ...
unsigned getNumFields() const
Gets the total number of fields for the given sparse tensor encoding.
StorageLayout(SparseTensorEncodingAttr enc)
FieldIndex getMemRefFieldIndex(SparseTensorFieldKind kind, std::optional< Level > lvl) const
Gets the field index for required field.
void foreachField(llvm::function_ref< bool(FieldIndex, SparseTensorFieldKind, Level, LevelType)>) const
For each field that will be allocated for the given sparse tensor encoding, calls the callback with t...
StorageLayout(const SparseTensorType &stt)
std::pair< FieldIndex, unsigned > getFieldIndexAndStride(SparseTensorFieldKind kind, std::optional< Level > lvl) const
void foreachFieldAndTypeInSparseTensor(SparseTensorType, llvm::function_ref< bool(Type, FieldIndex, SparseTensorFieldKind, Level, LevelType)>)
unsigned FieldIndex
The type of field indices.
uint64_t Level
The type of level identifiers and level-ranks.
Definition: SparseTensor.h:38
LevelType
This enum defines all the sparse representations supportable by the SparseTensor dialect.
Definition: Enums.h:168
SparseTensorFieldKind toFieldKind(StorageSpecifierKind kind)
unsigned getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc)
unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc)
SparseTensorFieldKind
===-------------------------------------------------------------------—===// The sparse tensor storag...
StorageSpecifierKind toSpecifierKind(SparseTensorFieldKind kind)
void foreachFieldInSparseTensor(SparseTensorEncodingAttr enc, llvm::function_ref< bool(FieldIndex, SparseTensorFieldKind, Level, LevelType)> callback)
Include the generated interface declarations.