MLIR  16.0.0git
CodegenUtils.h
Go to the documentation of this file.
1 //===- CodegenUtils.h - Utilities for generating MLIR -----------*- 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 generating MLIR.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_
14 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_
15 
23 #include "mlir/IR/Builders.h"
24 
25 namespace mlir {
26 
27 class Location;
28 class Type;
29 class Value;
30 
31 namespace sparse_tensor {
32 
33 /// Shorthand aliases for the `emitCInterface` argument to `getFunc()`,
34 /// `createFuncCall()`, and `replaceOpWithFuncCall()`.
35 enum class EmitCInterface : bool { Off = false, On = true };
36 
37 //===----------------------------------------------------------------------===//
38 // ExecutionEngine/SparseTensorUtils helper functions.
39 //===----------------------------------------------------------------------===//
40 
41 /// Converts an overhead storage bitwidth to its internal type-encoding.
42 OverheadType overheadTypeEncoding(unsigned width);
43 
44 /// Converts an overhead storage type to its internal type-encoding.
46 
47 /// Converts the internal type-encoding for overhead storage to an mlir::Type.
49 
50 /// Returns the OverheadType for pointer overhead storage.
51 OverheadType pointerOverheadTypeEncoding(const SparseTensorEncodingAttr &enc);
52 
53 /// Returns the OverheadType for index overhead storage.
54 OverheadType indexOverheadTypeEncoding(const SparseTensorEncodingAttr &enc);
55 
56 /// Returns the mlir::Type for pointer overhead storage.
58  const SparseTensorEncodingAttr &enc);
59 
60 /// Returns the mlir::Type for index overhead storage.
62  const SparseTensorEncodingAttr &enc);
63 
64 /// Convert OverheadType to its function-name suffix.
66 
67 /// Converts an overhead storage type to its function-name suffix.
68 StringRef overheadTypeFunctionSuffix(Type overheadTp);
69 
70 /// Converts a primary storage type to its internal type-encoding.
72 
73 /// Convert PrimaryType to its function-name suffix.
75 
76 /// Converts a primary storage type to its function-name suffix.
77 StringRef primaryTypeFunctionSuffix(Type elemTp);
78 
79 //===----------------------------------------------------------------------===//
80 // Misc code generators and utilities.
81 //===----------------------------------------------------------------------===//
82 
83 /// Generates a 1-valued attribute of the given type. This supports
84 /// all the same types as `getZeroAttr`; however, unlike `getZeroAttr`,
85 /// for unsupported types we raise `llvm_unreachable` rather than
86 /// returning a null attribute.
87 Attribute getOneAttr(Builder &builder, Type tp);
88 
89 /// Generates the comparison `v != 0` where `v` is of numeric type.
90 /// For floating types, we use the "unordered" comparator (i.e., returns
91 /// true if `v` is NaN).
92 Value genIsNonzero(OpBuilder &builder, Location loc, Value v);
93 
94 /// Computes the shape of destination tensor of a reshape operator. This is only
95 /// used when operands have dynamic shape. The shape of the destination is
96 /// stored into dstShape.
97 void genReshapeDstShape(Location loc, PatternRewriter &rewriter,
98  SmallVectorImpl<Value> &dstShape,
99  ArrayRef<Value> srcShape,
100  ArrayRef<int64_t> staticDstShape,
101  ArrayRef<ReassociationIndices> reassociation);
102 
103 /// Translate indices during a reshaping operation.
104 void translateIndicesArray(OpBuilder &builder, Location loc,
105  ArrayRef<ReassociationIndices> reassociation,
106  ValueRange srcIndices, ArrayRef<Value> srcShape,
107  ArrayRef<Value> dstShape,
108  SmallVectorImpl<Value> &dstIndices);
109 
110 /// Returns a function reference (first hit also inserts into module). Sets
111 /// the "_emit_c_interface" on the function declaration when requested,
112 /// so that LLVM lowering generates a wrapper function that takes care
113 /// of ABI complications with passing in and returning MemRefs to C functions.
114 FlatSymbolRefAttr getFunc(ModuleOp module, StringRef name, TypeRange resultType,
115  ValueRange operands, EmitCInterface emitCInterface);
116 
117 /// Creates a `CallOp` to the function reference returned by `getFunc()` in
118 /// the builder's module.
119 func::CallOp createFuncCall(OpBuilder &builder, Location loc, StringRef name,
120  TypeRange resultType, ValueRange operands,
121  EmitCInterface emitCInterface);
122 
123 /// Returns the equivalent of `void*` for opaque arguments to the
124 /// execution engine.
126 
127 /// Generates an uninitialized temporary buffer of the given size and
128 /// type, but returns it as type `memref<? x $tp>` (rather than as type
129 /// `memref<$sz x $tp>`).
130 Value genAlloca(OpBuilder &builder, Location loc, Value sz, Type tp);
131 
132 /// Generates an uninitialized temporary buffer of the given size and
133 /// type, but returns it as type `memref<? x $tp>` (rather than as type
134 /// `memref<$sz x $tp>`).
135 Value genAlloca(OpBuilder &builder, Location loc, unsigned sz, Type tp);
136 
137 /// Generates an uninitialized temporary buffer with room for one value
138 /// of the given type, and returns the `memref<$tp>`.
139 Value genAllocaScalar(OpBuilder &builder, Location loc, Type tp);
140 
141 /// Generates code to allocate a buffer of the given type, and zero
142 /// initialize it. If the buffer type has any dynamic sizes, then the
143 /// `sizes` parameter should be as filled by sizesFromPtr(); that way
144 /// we can reuse the genDimSizeCall() results generated by sizesFromPtr().
146  RankedTensorType tensorTp, ValueRange sizes);
147 
148 /// Generates code to deallocate a dense buffer.
149 void deallocDenseTensor(OpBuilder &builder, Location loc, Value buffer);
150 
151 /// Generates the code to read the value from tensor[ivs]. The generated code
152 /// looks like the following and the insertion point after this routine is
153 /// inside the if-then branch behind the assignment to ind.
154 /// if (tensor[ivs] != 0)
155 /// insert_point
156 Value genValueForDense(OpBuilder &builder, Location loc, Value tensor,
157  ValueRange ivs);
158 
159 /// Generates the loop structure to iterate over a dense tensor or a sparse
160 /// tensor constant to support the lowering of dense-to-sparse convert operator.
161 //
162 // The loop to iterate a dense tensor:
163 // for i1 in dim1
164 // ..
165 // for ik in dimk
166 // val = a[i1,..,ik]
167 // if val != 0
168 // loop-body
169 //
170 // The loop to iterate a sparse tensor constant:
171 // for i in range(NNZ)
172 // val = values[i]
173 // [i1,..,ik] = indices[i]
174 // loop-body
176  OpBuilder &builder, Location loc, Value src, unsigned rank,
177  function_ref<void(OpBuilder &, Location, Value, ValueRange)> bodyBuilder);
178 
179 /// Populates given sizes array from dense tensor or sparse tensor constant.
180 void sizesFromSrc(OpBuilder &builder, SmallVectorImpl<Value> &sizes,
181  Location loc, Value src);
182 
183 /// Scans to top of generated loop.
185 
186 /// Iterate over a sparse constant, generates constantOp for value and indices.
187 /// E.g.,
188 /// sparse<[ [0], [28], [31] ],
189 /// [ (-5.13, 2.0), (3.0, 4.0), (5.0, 6.0) ] >
190 /// =>
191 /// %c1 = arith.constant 0
192 /// %v1 = complex.constant (5.13, 2.0)
193 /// callback({%c1}, %v1)
194 ///
195 /// %c2 = arith.constant 28
196 /// %v2 = complex.constant (3.0, 4.0)
197 /// callback({%c2}, %v2)
198 ///
199 /// %c3 = arith.constant 31
200 /// %v3 = complex.constant (5.0, 6.0)
201 /// callback({%c3}, %v3)
203  Location loc, RewriterBase &rewriter, SparseElementsAttr attr,
204  function_ref<void(ArrayRef<Value>, Value)> callback);
205 
206 //===----------------------------------------------------------------------===//
207 // Inlined constant generators.
208 //
209 // All these functions are just wrappers to improve code legibility;
210 // therefore, we mark them as `inline` to avoid introducing any additional
211 // overhead due to the legibility.
212 //
213 // TODO: Ideally these should move upstream, so that we don't
214 // develop a design island. However, doing so will involve
215 // substantial design work. For related prior discussion, see
216 // <https://llvm.discourse.group/t/evolving-builder-apis-based-on-lessons-learned-from-edsc/879>
217 //===----------------------------------------------------------------------===//
218 
219 /// Generates a 0-valued constant of the given type. In addition to
220 /// the scalar types (`ComplexType`, ``FloatType`, `IndexType`,
221 /// `IntegerType`), this also works for `RankedTensorType` and `VectorType`
222 /// (for which it generates a constant `DenseElementsAttr` of zeros).
223 inline Value constantZero(OpBuilder &builder, Location loc, Type tp) {
224  if (auto ctp = tp.dyn_cast<ComplexType>()) {
225  auto zeroe = builder.getZeroAttr(ctp.getElementType());
226  auto zeroa = builder.getArrayAttr({zeroe, zeroe});
227  return builder.create<complex::ConstantOp>(loc, tp, zeroa);
228  }
229  return builder.create<arith::ConstantOp>(loc, tp, builder.getZeroAttr(tp));
230 }
231 
232 /// Generates a 1-valued constant of the given type. This supports all
233 /// the same types as `constantZero`.
234 inline Value constantOne(OpBuilder &builder, Location loc, Type tp) {
235  if (auto ctp = tp.dyn_cast<ComplexType>()) {
236  auto zeroe = builder.getZeroAttr(ctp.getElementType());
237  auto onee = getOneAttr(builder, ctp.getElementType());
238  auto zeroa = builder.getArrayAttr({onee, zeroe});
239  return builder.create<complex::ConstantOp>(loc, tp, zeroa);
240  }
241  return builder.create<arith::ConstantOp>(loc, tp, getOneAttr(builder, tp));
242 }
243 
244 /// Generates a constant of `index` type.
245 inline Value constantIndex(OpBuilder &builder, Location loc, int64_t i) {
246  return builder.create<arith::ConstantIndexOp>(loc, i);
247 }
248 
249 /// Generates a constant of `i32` type.
250 inline Value constantI32(OpBuilder &builder, Location loc, int32_t i) {
251  return builder.create<arith::ConstantIntOp>(loc, i, 32);
252 }
253 
254 /// Generates a constant of `i16` type.
255 inline Value constantI16(OpBuilder &builder, Location loc, int16_t i) {
256  return builder.create<arith::ConstantIntOp>(loc, i, 16);
257 }
258 
259 /// Generates a constant of `i8` type.
260 inline Value constantI8(OpBuilder &builder, Location loc, int8_t i) {
261  return builder.create<arith::ConstantIntOp>(loc, i, 8);
262 }
263 
264 /// Generates a constant of `i1` type.
265 inline Value constantI1(OpBuilder &builder, Location loc, bool b) {
266  return builder.create<arith::ConstantIntOp>(loc, b, 1);
267 }
268 
269 /// Generates a constant of the given `Action`.
270 inline Value constantAction(OpBuilder &builder, Location loc, Action action) {
271  return constantI32(builder, loc, static_cast<uint32_t>(action));
272 }
273 
274 /// Generates a constant of the internal type-encoding for overhead storage.
276  unsigned width) {
277  return constantI32(builder, loc,
278  static_cast<uint32_t>(overheadTypeEncoding(width)));
279 }
280 
281 /// Generates a constant of the internal type-encoding for pointer
282 /// overhead storage.
284  const SparseTensorEncodingAttr &enc) {
285  return constantOverheadTypeEncoding(builder, loc, enc.getPointerBitWidth());
286 }
287 
288 /// Generates a constant of the internal type-encoding for index overhead
289 /// storage.
291  const SparseTensorEncodingAttr &enc) {
292  return constantOverheadTypeEncoding(builder, loc, enc.getIndexBitWidth());
293 }
294 
295 /// Generates a constant of the internal type-encoding for primary storage.
297  Type elemTp) {
298  return constantI32(builder, loc,
299  static_cast<uint32_t>(primaryTypeEncoding(elemTp)));
300 }
301 
302 /// Generates a constant of the internal dimension level type encoding.
304  DimLevelType dlt) {
305  return constantI8(builder, loc, static_cast<uint8_t>(dlt));
306 }
307 
308 inline bool isZeroRankedTensorOrScalar(Type type) {
309  auto rtp = type.dyn_cast<RankedTensorType>();
310  return !rtp || rtp.getRank() == 0;
311 }
312 
313 //===----------------------------------------------------------------------===//
314 // SparseTensorLoopEmiter class, manages sparse tensors and helps to generate
315 // loop structure to (co)-iterate sparse tensors.
316 //
317 // An example usage:
318 // To generate the following loops over T1<?x?> and T2<?x?>
319 //
320 // for i in TENSOR_1_0 {
321 // for j : TENSOR_2_0 {
322 // for k : TENSOR_1_1 {}
323 // for k : TENSOR_2_1 {}
324 // }
325 // }
326 //
327 // One can use
328 //
329 // SparseTensorLoopEmiter loopEmiter({T1, T1});
330 // loopEmiter.initializeLoopEmit();
331 // loopEmiter.enterLoopOverTensorAtDim(T1, 0);
332 // loopEmiter.enterLoopOverTensorAtDim(T2, 0);
333 // loopEmiter.enterLoopOverTensorAtDim(T1, 1);
334 // loopEmiter.exitCurrentLoop();
335 // loopEmiter.enterLoopOverTensorAtDim(T2, 1);
336 // loopEmiter.exitCurrentLoop(); // exit k
337 // loopEmiter.exitCurrentLoop(); // exit j
338 // loopEmiter.exitCurrentLoop(); // exit i
339 //===----------------------------------------------------------------------===//
340 
342 public:
343  /// Optional callback function to setup dense output tensors when
344  /// initializing the loop emitter (e.g., to fill a dense output with zeros).
346  Value memref, Value tensor)>;
347 
348  /// Constructor: take an array of tensors inputs, on which the generated loops
349  /// will iterate on. The index of the tensor in the array is also the
350  /// tensor id (tid) used in related functions.
351  /// If isSparseOut is set, loop emitter assume that the sparse output tensor
352  /// is empty, and will always generate loops on it based on the dim sizes.
353  /// An optional array could be provided (by sparsification) to indicate the
354  /// loop id sequence that will be generated. It is used to establish the
355  /// mapping between affineDimExpr to the corresponding loop index in the loop
356  /// stack that are maintained by the loop emitter.
357  explicit SparseTensorLoopEmitter(ValueRange tensors,
358  StringAttr loopTag = nullptr,
359  bool hasOutput = false,
360  bool isSparseOut = false,
361  ArrayRef<unsigned> topSort = {});
362 
363  /// Starts a loop emitting session by generating all the buffers needed to
364  /// iterate tensors.
365  void initializeLoopEmit(OpBuilder &builder, Location loc,
366  OutputUpdater updater = nullptr);
367 
368  /// Generates a list of operations to compute the affine expression.
369  Value genAffine(OpBuilder &builder, AffineExpr a, Location loc);
370 
371  /// Enters a new loop sequence, the loops within the same sequence starts from
372  /// the break points of previous loop instead of starting over from 0.
373  /// e.g.,
374  /// {
375  /// // loop sequence start.
376  /// p0 = while(xxx)
377  /// ...
378  /// break p0
379  ///
380  /// // Starts loop from p0
381  /// for (i = p0; i < end; i++)
382  /// ...
383  /// // loop sequence end.
384  /// }
385  void enterNewLoopSeq(OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
386  ArrayRef<size_t> dims);
387 
388  // exit the current loop sequence, this will reset universal index to 0.
390  assert(loopSeqStack.size() == loopStack.size() + 1);
391  loopSeqStack.pop_back();
392  }
393 
394  // TODO: Gets rid of `dim` in the argument list? Track the dimension we
395  // are currently at internally. Then it would be enterNextDimForTensor.
396  // Still need a way to specify the dim for non annoated dense tensor though,
397  // as it can be accessed out of order.
398  /// Emits loop over tensor_tid_dim, it assumes that loops between
399  /// tensor_tid_[0, dim - 1] have already been generated.
400  /// The function will also perform in-place update on the `reduc` vector to
401  /// return the reduction variable used inside the generated loop.
403  size_t tid, size_t dim,
404  MutableArrayRef<Value> reduc = {},
405  bool isParallel = false,
406  ArrayRef<size_t> extraTids = {},
407  ArrayRef<size_t> extraDims = {});
408 
409  Operation *enterFilterLoopOverTensorAtDim(OpBuilder &builder, Location loc,
410  size_t tid, size_t dim,
411  AffineExpr affine,
412  MutableArrayRef<Value> reduc = {});
413 
414  void genDenseAffineAddressAtCurLevel(OpBuilder &builder, Location loc,
415  size_t tid, size_t dim,
416  AffineExpr affine);
417 
418  /// Emits a co-iteration loop over a set of tensors.
420  OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
421  ArrayRef<size_t> dims, bool needsUniv, MutableArrayRef<Value> reduc = {},
422  ArrayRef<size_t> extraTids = {}, ArrayRef<size_t> extraDims = {});
423 
424  void exitCurrentLoop(RewriterBase &rewriter, Location loc,
425  MutableArrayRef<Value> reduc = {});
426 
427  /// Returns the array of coordinate for all the loop generated till now.
429  for (auto &l : loopStack)
430  coords.push_back(l.iv);
431  }
432 
433  /// Gets loop induction variable at the given level.
434  unsigned getCurrentDepth() const { return loopStack.size(); }
435 
436  /// Gets loop induction variable at the given level.
437  Value getLoopIV(size_t level) const {
438  if (level < loopStack.size())
439  return loopStack[level].iv;
440  return nullptr;
441  }
442 
443  ///
444  /// Getters.
445  ///
446  const std::vector<std::vector<Value>> &getPidxs() const { return pidxs; };
447  const std::vector<std::vector<Value>> &getCoord() const { return coord; };
448  const std::vector<std::vector<Value>> &getHighs() const { return highs; };
449  const std::vector<std::vector<Value>> &getPtrBuffer() const {
450  return ptrBuffer;
451  };
452  const std::vector<std::vector<Value>> &getIdxBuffer() const {
453  return idxBuffer;
454  };
455  const std::vector<Value> &getValBuffer() const { return valBuffer; };
456 
457  constexpr static llvm::StringLiteral getLoopEmitterLoopAttrName() {
458  return llvm::StringLiteral("Emitted from");
459  }
460 
461 private:
462  struct LoopLevelInfo {
463  LoopLevelInfo(ArrayRef<size_t> tids, ArrayRef<size_t> dims, Operation *loop,
464  Value iv, StringAttr loopTag)
465  : tids(tids), dims(dims), loop(loop), iv(iv) {
466  // Attached a special tag to loop emitter generated loop.
467  if (loopTag)
469  loopTag);
470  }
471  // TODO: maybe use a vector<pair> for tid and dim?
472  // The set of tensors that the loop is operating on
473  const llvm::SmallVector<size_t> tids;
474  // The corresponding dims for the tensors
475  const llvm::SmallVector<size_t> dims;
476  const Operation *loop; // the loop operation
477  const Value iv; // the induction variable for the loop
478  };
479 
480  /// Linearizes address for dense dimension (i.e., p = (i * d0) + j).
481  Value genAddress(OpBuilder &builder, Location loc, size_t tid, size_t dim,
482  Value iv) {
483  Value p = dim == 0 ? constantIndex(builder, loc, 0) : pidxs[tid][dim - 1];
484  Value mul = builder.create<arith::MulIOp>(loc, highs[tid][dim], p);
485  Value add = builder.create<arith::AddIOp>(loc, mul, iv);
486  return add;
487  }
488 
489  bool isOutputTensor(size_t tid) {
490  return hasOutput && tid == tensors.size() - 1;
491  }
492 
493  bool isSparseOutput(size_t tid) { return isOutputTensor(tid) && isSparseOut; }
494 
495  /// Setups [lo, hi] for iterating tensor[dim], it assumes that tensor[0
496  /// ...dims-1] has already been setup.
497  void prepareLoopOverTensorAtDim(OpBuilder &builder, Location loc, size_t tid,
498  size_t dim);
499 
500  /// Emits extra locals, since the locals might not be in simplified lattices
501  /// point used to generate the loops, but are still required to generates
502  /// expressions.
503  void emitExtraLocalsForTensorsAtDenseDims(OpBuilder &builder, Location loc,
504  ArrayRef<size_t> tids,
505  ArrayRef<size_t> dims);
506 
507  /// Exits a for loop, returns the reduction results, e.g.,
508  /// For sequential for loops:
509  /// %ret = for () {
510  /// ...
511  /// %val = addi %args, %c
512  /// yield %val
513  /// }
514  /// For parallel loops, the following generated code by users:
515  /// %ret = parallel () init(%args) {
516  /// ...
517  /// %val = op %args, %c
518  /// }
519  /// will be transformed into
520  /// %ret = parallel () init(%args) {
521  /// ...
522  /// scf.reduce(%c) bb0(%0, %1){
523  /// %val = op %0, %1
524  /// scf.reduce.return %val
525  /// }
526  /// }
527  /// NOTE: only one instruction will be moved into reduce block, transformation
528  /// will fail if multiple instructions are used to compute the reduction
529  /// value.
530  /// Return %ret to user, while %val is provided by users (`reduc`).
531  void exitForLoop(RewriterBase &rewriter, Location loc,
532  MutableArrayRef<Value> reduc);
533 
534  /// Exits a while loop, returns the reduction results.
535  void exitCoIterationLoop(OpBuilder &builder, Location loc,
536  MutableArrayRef<Value> reduc);
537 
538  /// A optional string attribute that should be attached to the loop generated
539  /// by loop emitter, it might help following passes to identify loops that
540  /// operates on sparse tensors more easily.
541  StringAttr loopTag;
542  /// Whether the loop emitter needs to treat the last tensor as the output
543  /// tensor.
544  bool hasOutput;
545  bool isSparseOut;
546  /// Input and (optional) output tensors.
547  std::vector<Value> tensors;
548  /// The dim type array for each tensor.
549  std::vector<std::vector<DimLevelType>> dimTypes;
550  /// Sparse iteration information (by tensor and dim). These arrays
551  /// are updated to remain current within the current loop.
552  std::vector<std::vector<Value>> pidxs;
553  std::vector<std::vector<Value>> coord;
554  std::vector<std::vector<Value>> highs;
555  std::vector<std::vector<Value>> ptrBuffer; // to_pointers
556  std::vector<std::vector<Value>> idxBuffer; // to_indices
557  std::vector<Value> valBuffer; // to_value
558 
559  // Loop Stack, stores the information of all the nested loops that are alive.
560  std::vector<LoopLevelInfo> loopStack;
561 
562  // Loop Sequence Stack, stores the unversial index for the current loop
563  // sequence.
564  std::vector<Value> loopSeqStack;
565 
566  // Maps AffineDimExpr to the index of the loop in loopStack.
567  // TODO: We should probably use a callback function here to make it more
568  // general.
569  std::vector<unsigned> sparsiferLoopLvlMap;
570 
571  // TODO: not yet used, it should track the current level for each tensor
572  // to help eliminate `dim` paramters from above APIs.
573  // std::vector<size_t> curLv;
574 };
575 
576 } // namespace sparse_tensor
577 } // namespace mlir
578 
579 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_
Base type for affine expression.
Definition: AffineExpr.h:68
Attributes are known-constant values of operations.
Definition: Attributes.h:25
This class is a general helper class for creating context-global objects like types,...
Definition: Builders.h:49
Attribute getZeroAttr(Type type)
Definition: Builders.cpp:306
ArrayAttr getArrayAttr(ArrayRef< Attribute > value)
Definition: Builders.cpp:247
A symbol reference with a reference path containing a single element.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:64
This class helps build Operations.
Definition: Builders.h:198
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:422
Operation is a basic unit of execution within MLIR.
Definition: Operation.h:31
void setAttr(StringAttr name, Attribute value)
If the an attribute exists with the specified name, change it to the new value.
Definition: Operation.h:395
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:605
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:398
This class provides an abstraction over the various different ranges of value types.
Definition: TypeRange.h:36
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
U dyn_cast() const
Definition: Types.h:270
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:349
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:85
Specialization of arith.constant op that returns an integer of index type.
Definition: Arith.h:89
Specialization of arith.constant op that returns an integer value.
Definition: Arith.h:52
const std::vector< std::vector< Value > > & getPidxs() const
Getters.
Definition: CodegenUtils.h:446
Operation * enterCoIterationOverTensorsAtDims(OpBuilder &builder, Location loc, ArrayRef< size_t > tids, ArrayRef< size_t > dims, bool needsUniv, MutableArrayRef< Value > reduc={}, ArrayRef< size_t > extraTids={}, ArrayRef< size_t > extraDims={})
Emits a co-iteration loop over a set of tensors.
const std::vector< std::vector< Value > > & getIdxBuffer() const
Definition: CodegenUtils.h:452
Operation * enterFilterLoopOverTensorAtDim(OpBuilder &builder, Location loc, size_t tid, size_t dim, AffineExpr affine, MutableArrayRef< Value > reduc={})
void genDenseAffineAddressAtCurLevel(OpBuilder &builder, Location loc, size_t tid, size_t dim, AffineExpr affine)
Value getLoopIV(size_t level) const
Gets loop induction variable at the given level.
Definition: CodegenUtils.h:437
Value genAffine(OpBuilder &builder, AffineExpr a, Location loc)
Generates a list of operations to compute the affine expression.
void getCoordinateArray(SmallVectorImpl< Value > &coords) const
Returns the array of coordinate for all the loop generated till now.
Definition: CodegenUtils.h:428
void initializeLoopEmit(OpBuilder &builder, Location loc, OutputUpdater updater=nullptr)
Starts a loop emitting session by generating all the buffers needed to iterate tensors.
Operation * enterLoopOverTensorAtDim(OpBuilder &builder, Location loc, size_t tid, size_t dim, MutableArrayRef< Value > reduc={}, bool isParallel=false, ArrayRef< size_t > extraTids={}, ArrayRef< size_t > extraDims={})
Emits loop over tensor_tid_dim, it assumes that loops between tensor_tid_[0, dim - 1] have already be...
unsigned getCurrentDepth() const
Gets loop induction variable at the given level.
Definition: CodegenUtils.h:434
void enterNewLoopSeq(OpBuilder &builder, Location loc, ArrayRef< size_t > tids, ArrayRef< size_t > dims)
Enters a new loop sequence, the loops within the same sequence starts from the break points of previo...
function_ref< Value(OpBuilder &builder, Location loc, Value memref, Value tensor)> OutputUpdater
Optional callback function to setup dense output tensors when initializing the loop emitter (e....
Definition: CodegenUtils.h:346
SparseTensorLoopEmitter(ValueRange tensors, StringAttr loopTag=nullptr, bool hasOutput=false, bool isSparseOut=false, ArrayRef< unsigned > topSort={})
Constructor: take an array of tensors inputs, on which the generated loops will iterate on.
const std::vector< std::vector< Value > > & getPtrBuffer() const
Definition: CodegenUtils.h:449
const std::vector< Value > & getValBuffer() const
Definition: CodegenUtils.h:455
constexpr static llvm::StringLiteral getLoopEmitterLoopAttrName()
Definition: CodegenUtils.h:457
void exitCurrentLoop(RewriterBase &rewriter, Location loc, MutableArrayRef< Value > reduc={})
const std::vector< std::vector< Value > > & getHighs() const
Definition: CodegenUtils.h:448
const std::vector< std::vector< Value > > & getCoord() const
Definition: CodegenUtils.h:447
@ Type
An inlay hint that for a type annotation.
Attribute getOneAttr(Builder &builder, Type tp)
Generates a 1-valued attribute of the given type.
void translateIndicesArray(OpBuilder &builder, Location loc, ArrayRef< ReassociationIndices > reassociation, ValueRange srcIndices, ArrayRef< Value > srcShape, ArrayRef< Value > dstShape, SmallVectorImpl< Value > &dstIndices)
Translate indices during a reshaping operation.
Value constantIndexTypeEncoding(OpBuilder &builder, Location loc, const SparseTensorEncodingAttr &enc)
Generates a constant of the internal type-encoding for index overhead storage.
Definition: CodegenUtils.h:290
FlatSymbolRefAttr getFunc(ModuleOp module, StringRef name, TypeRange resultType, ValueRange operands, EmitCInterface emitCInterface)
Returns a function reference (first hit also inserts into module).
Value genAllocaScalar(OpBuilder &builder, Location loc, Type tp)
Generates an uninitialized temporary buffer with room for one value of the given type,...
Value constantIndex(OpBuilder &builder, Location loc, int64_t i)
Generates a constant of index type.
Definition: CodegenUtils.h:245
Value constantZero(OpBuilder &builder, Location loc, Type tp)
Generates a 0-valued constant of the given type.
Definition: CodegenUtils.h:223
Value constantOne(OpBuilder &builder, Location loc, Type tp)
Generates a 1-valued constant of the given type.
Definition: CodegenUtils.h:234
OverheadType
Encoding of overhead types (both pointer overhead and indices overhead), for "overloading" @newSparse...
Definition: Enums.h:49
Value constantI8(OpBuilder &builder, Location loc, int8_t i)
Generates a constant of i8 type.
Definition: CodegenUtils.h:260
Action
The actions performed by @newSparseTensor.
Definition: Enums.h:123
Type getPointerOverheadType(Builder &builder, const SparseTensorEncodingAttr &enc)
Returns the mlir::Type for pointer overhead storage.
Value genValueForDense(OpBuilder &builder, Location loc, Value tensor, ValueRange ivs)
Generates the code to read the value from tensor[ivs].
OverheadType overheadTypeEncoding(unsigned width)
Converts an overhead storage bitwidth to its internal type-encoding.
Value constantI1(OpBuilder &builder, Location loc, bool b)
Generates a constant of i1 type.
Definition: CodegenUtils.h:265
Value constantAction(OpBuilder &builder, Location loc, Action action)
Generates a constant of the given Action.
Definition: CodegenUtils.h:270
StringRef overheadTypeFunctionSuffix(OverheadType ot)
Convert OverheadType to its function-name suffix.
PrimaryType
Encoding of the elemental type, for "overloading" @newSparseTensor.
Definition: Enums.h:80
PrimaryType primaryTypeEncoding(Type elemTp)
Converts a primary storage type to its internal type-encoding.
Operation * getTop(Operation *op)
Scans to top of generated loop.
Type getIndexOverheadType(Builder &builder, const SparseTensorEncodingAttr &enc)
Returns the mlir::Type for index overhead storage.
void genDenseTensorOrSparseConstantIterLoop(OpBuilder &builder, Location loc, Value src, unsigned rank, function_ref< void(OpBuilder &, Location, Value, ValueRange)> bodyBuilder)
Generates the loop structure to iterate over a dense tensor or a sparse tensor constant to support th...
Value genIsNonzero(OpBuilder &builder, Location loc, Value v)
Generates the comparison v != 0 where v is of numeric type.
Type getOpaquePointerType(OpBuilder &builder)
Returns the equivalent of void* for opaque arguments to the execution engine.
DimLevelType
This enum defines all the sparse representations supportable by the SparseTensor dialect.
Definition: Enums.h:147
bool isZeroRankedTensorOrScalar(Type type)
Definition: CodegenUtils.h:308
void deallocDenseTensor(OpBuilder &builder, Location loc, Value buffer)
Generates code to deallocate a dense buffer.
Value genAlloca(OpBuilder &builder, Location loc, Value sz, Type tp)
Generates an uninitialized temporary buffer of the given size and type, but returns it as type memref...
Value constantOverheadTypeEncoding(OpBuilder &builder, Location loc, unsigned width)
Generates a constant of the internal type-encoding for overhead storage.
Definition: CodegenUtils.h:275
OverheadType pointerOverheadTypeEncoding(const SparseTensorEncodingAttr &enc)
Returns the OverheadType for pointer overhead storage.
void foreachInSparseConstant(Location loc, RewriterBase &rewriter, SparseElementsAttr attr, function_ref< void(ArrayRef< Value >, Value)> callback)
Iterate over a sparse constant, generates constantOp for value and indices.
Value constantPointerTypeEncoding(OpBuilder &builder, Location loc, const SparseTensorEncodingAttr &enc)
Generates a constant of the internal type-encoding for pointer overhead storage.
Definition: CodegenUtils.h:283
Value constantI16(OpBuilder &builder, Location loc, int16_t i)
Generates a constant of i16 type.
Definition: CodegenUtils.h:255
func::CallOp createFuncCall(OpBuilder &builder, Location loc, StringRef name, TypeRange resultType, ValueRange operands, EmitCInterface emitCInterface)
Creates a CallOp to the function reference returned by getFunc() in the builder's module.
OverheadType indexOverheadTypeEncoding(const SparseTensorEncodingAttr &enc)
Returns the OverheadType for index overhead storage.
StringRef primaryTypeFunctionSuffix(PrimaryType pt)
Convert PrimaryType to its function-name suffix.
void genReshapeDstShape(Location loc, PatternRewriter &rewriter, SmallVectorImpl< Value > &dstShape, ArrayRef< Value > srcShape, ArrayRef< int64_t > staticDstShape, ArrayRef< ReassociationIndices > reassociation)
Computes the shape of destination tensor of a reshape operator.
Value constantI32(OpBuilder &builder, Location loc, int32_t i)
Generates a constant of i32 type.
Definition: CodegenUtils.h:250
Value constantPrimaryTypeEncoding(OpBuilder &builder, Location loc, Type elemTp)
Generates a constant of the internal type-encoding for primary storage.
Definition: CodegenUtils.h:296
void sizesFromSrc(OpBuilder &builder, SmallVectorImpl< Value > &sizes, Location loc, Value src)
Populates given sizes array from dense tensor or sparse tensor constant.
Type getOverheadType(Builder &builder, OverheadType ot)
Converts the internal type-encoding for overhead storage to an mlir::Type.
Value constantDimLevelTypeEncoding(OpBuilder &builder, Location loc, DimLevelType dlt)
Generates a constant of the internal dimension level type encoding.
Definition: CodegenUtils.h:303
EmitCInterface
Shorthand aliases for the emitCInterface argument to getFunc(), createFuncCall(), and replaceOpWithFu...
Definition: CodegenUtils.h:35
Value allocDenseTensor(OpBuilder &builder, Location loc, RankedTensorType tensorTp, ValueRange sizes)
Generates code to allocate a buffer of the given type, and zero initialize it.
Include the generated interface declarations.