MLIR  18.0.0git
CodegenUtils.cpp
Go to the documentation of this file.
1 //===- CodegenUtils.cpp - Utilities for generating MLIR -------------------===//
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 "CodegenUtils.h"
10 #include "SparseTensorDescriptor.h"
11 
18 #include "mlir/IR/Matchers.h"
19 #include "mlir/IR/Types.h"
20 #include "mlir/IR/Value.h"
21 #include <optional>
22 
23 using namespace mlir;
24 using namespace mlir::sparse_tensor;
25 
26 /// If the tensor is a sparse constant, generates and returns the pair of
27 /// the constants for the coordinates and the values.
28 static std::optional<std::pair<Value, Value>>
30  if (auto constOp = tensor.getDefiningOp<arith::ConstantOp>()) {
31  if (auto a = dyn_cast<SparseElementsAttr>(constOp.getValue())) {
32  auto coordinates = builder.create<arith::ConstantOp>(loc, a.getIndices());
33  auto values = builder.create<arith::ConstantOp>(loc, a.getValues());
34  return std::make_pair(coordinates, values);
35  }
36  }
37  return {};
38 }
39 
40 /// Reads `coordinates[k][0..rank-1]` and `value[k]`, appending the
41 /// former onto `cvs` and returning the latter.
42 // FIXME: Change the `rank` argument to `Dimension dimRank` or `Level lvlRank`,
43 // to clarify its intended meaning.
45  Value coordinates, Value values,
47  unsigned rank) {
48  for (unsigned d = 0; d < rank; d++) {
49  Value dim = constantIndex(builder, loc, d);
50  Value crd =
51  builder.create<tensor::ExtractOp>(loc, coordinates, ValueRange{k, dim});
52  crd = builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), crd);
53  // builder.create<memref::StoreOp>(loc, crd, cvs, dim);
54  cvs.push_back(crd);
55  }
56  return builder.create<tensor::ExtractOp>(loc, values, k);
57 }
58 
59 /// Generates code to read the value from `tensor[ivs]`, and open
60 /// a conditional for whether the value is non-zero. The generated code
61 /// looks like the following and the insertion point after this routine
62 /// is inside the then-branch.
63 /// if (tensor[ivs] != 0)
64 /// insert_point
66  Value tensor,
68  ValueRange ivs) {
69  Value val = genValueForDense(builder, loc, tensor, ivs);
70  cvs.append(ivs.begin(), ivs.end());
71  return val;
72 }
73 
74 //===----------------------------------------------------------------------===//
75 // ExecutionEngine/SparseTensorUtils helper functions.
76 //===----------------------------------------------------------------------===//
77 
79  switch (width) {
80  case 64:
81  return OverheadType::kU64;
82  case 32:
83  return OverheadType::kU32;
84  case 16:
85  return OverheadType::kU16;
86  case 8:
87  return OverheadType::kU8;
88  case 0:
89  return OverheadType::kIndex;
90  }
91  llvm_unreachable("Unsupported overhead bitwidth");
92 }
93 
95  if (tp.isIndex())
96  return OverheadType::kIndex;
97  if (auto intTp = dyn_cast<IntegerType>(tp))
98  return overheadTypeEncoding(intTp.getWidth());
99  llvm_unreachable("Unknown overhead type");
100 }
101 
102 // TODO: should offer an overload of this that takes a `MLIRContext*`
103 // instead of the builder, similar to `detail::getIntegerOrIndexType`.
105  switch (ot) {
107  return builder.getIndexType();
108  case OverheadType::kU64:
109  return builder.getIntegerType(64);
110  case OverheadType::kU32:
111  return builder.getIntegerType(32);
112  case OverheadType::kU16:
113  return builder.getIntegerType(16);
114  case OverheadType::kU8:
115  return builder.getIntegerType(8);
116  }
117  llvm_unreachable("Unknown OverheadType");
118 }
119 
121 mlir::sparse_tensor::posTypeEncoding(SparseTensorEncodingAttr enc) {
122  return overheadTypeEncoding(enc.getPosWidth());
123 }
124 
126 mlir::sparse_tensor::crdTypeEncoding(SparseTensorEncodingAttr enc) {
127  return overheadTypeEncoding(enc.getCrdWidth());
128 }
129 
130 // TODO: we ought to add some `static_assert` tests to ensure that the
131 // `STEA::get{Pos,Crd}Type` methods agree with `getOverheadType(builder,
132 // {pos,crd}OverheadTypeEncoding(enc))`
133 
134 // TODO: Adjust the naming convention for the constructors of
135 // `OverheadType` so we can use the `MLIR_SPARSETENSOR_FOREVERY_O` x-macro
136 // here instead of `MLIR_SPARSETENSOR_FOREVERY_FIXED_O`; to further reduce
137 // the possibility of typo bugs or things getting out of sync.
139  switch (ot) {
141  return "0";
142 #define CASE(ONAME, O) \
143  case OverheadType::kU##ONAME: \
144  return #ONAME;
146 #undef CASE
147  }
148  llvm_unreachable("Unknown OverheadType");
149 }
150 
153 }
154 
156  if (elemTp.isF64())
157  return PrimaryType::kF64;
158  if (elemTp.isF32())
159  return PrimaryType::kF32;
160  if (elemTp.isF16())
161  return PrimaryType::kF16;
162  if (elemTp.isBF16())
163  return PrimaryType::kBF16;
164  if (elemTp.isInteger(64))
165  return PrimaryType::kI64;
166  if (elemTp.isInteger(32))
167  return PrimaryType::kI32;
168  if (elemTp.isInteger(16))
169  return PrimaryType::kI16;
170  if (elemTp.isInteger(8))
171  return PrimaryType::kI8;
172  if (auto complexTp = dyn_cast<ComplexType>(elemTp)) {
173  auto complexEltTp = complexTp.getElementType();
174  if (complexEltTp.isF64())
175  return PrimaryType::kC64;
176  if (complexEltTp.isF32())
177  return PrimaryType::kC32;
178  }
179  llvm_unreachable("Unknown primary type");
180 }
181 
183  switch (pt) {
184 #define CASE(VNAME, V) \
185  case PrimaryType::k##VNAME: \
186  return #VNAME;
188 #undef CASE
189  }
190  llvm_unreachable("Unknown PrimaryType");
191 }
192 
195 }
196 
197 //===----------------------------------------------------------------------===//
198 // Misc code generators.
199 //===----------------------------------------------------------------------===//
200 
202  Type dstTp) {
203  const Type srcTp = value.getType();
204  if (srcTp == dstTp)
205  return value;
206 
207  // int <=> index
208  if (isa<IndexType>(srcTp) || isa<IndexType>(dstTp))
209  return builder.create<arith::IndexCastOp>(loc, dstTp, value);
210 
211  const auto srcIntTp = dyn_cast_or_null<IntegerType>(srcTp);
212  const bool isUnsignedCast = srcIntTp ? srcIntTp.isUnsigned() : false;
213  return mlir::convertScalarToDtype(builder, loc, value, dstTp, isUnsignedCast);
214 }
215 
217  Value s) {
218  Value load = builder.create<memref::LoadOp>(loc, mem, s);
219  if (!isa<IndexType>(load.getType())) {
220  if (load.getType().getIntOrFloatBitWidth() < 64)
221  load = builder.create<arith::ExtUIOp>(loc, builder.getI64Type(), load);
222  load =
223  builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), load);
224  }
225  return load;
226 }
227 
228 mlir::TypedAttr mlir::sparse_tensor::getOneAttr(Builder &builder, Type tp) {
229  if (isa<FloatType>(tp))
230  return builder.getFloatAttr(tp, 1.0);
231  if (isa<IndexType>(tp))
232  return builder.getIndexAttr(1);
233  if (auto intTp = dyn_cast<IntegerType>(tp))
234  return builder.getIntegerAttr(tp, APInt(intTp.getWidth(), 1));
235  if (isa<RankedTensorType, VectorType>(tp)) {
236  auto shapedTp = cast<ShapedType>(tp);
237  if (auto one = getOneAttr(builder, shapedTp.getElementType()))
238  return DenseElementsAttr::get(shapedTp, one);
239  }
240  llvm_unreachable("Unsupported attribute type");
241 }
242 
244  Value v) {
245  Type tp = v.getType();
246  Value zero = constantZero(builder, loc, tp);
247  if (isa<FloatType>(tp))
248  return builder.create<arith::CmpFOp>(loc, arith::CmpFPredicate::UNE, v,
249  zero);
250  if (tp.isIntOrIndex())
251  return builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ne, v,
252  zero);
253  if (dyn_cast<ComplexType>(tp))
254  return builder.create<complex::NotEqualOp>(loc, v, zero);
255  llvm_unreachable("Non-numeric type");
256 }
257 
259  OpBuilder &builder, Location loc, SmallVectorImpl<Value> &dstShape,
260  ArrayRef<Value> srcShape, ArrayRef<StaticSize> staticDstShape,
261  ArrayRef<ReassociationIndices> reassociation) {
262  // Collapse shape.
263  if (reassociation.size() < srcShape.size()) {
264  unsigned start = 0;
265  for (const auto &map : llvm::enumerate(reassociation)) {
266  auto dstDim = constantIndex(builder, loc, 1);
267  for (unsigned i = start; i < start + map.value().size(); i++) {
268  dstDim = builder.create<arith::MulIOp>(loc, dstDim, srcShape[i]);
269  }
270  dstShape.push_back(dstDim);
271  start = start + map.value().size();
272  }
273  assert(start == srcShape.size());
274  return;
275  }
276 
277  // Expand shape.
278  assert(reassociation.size() == srcShape.size());
279  unsigned start = 0;
280  // Expand the i-th dimension in srcShape.
281  for (unsigned i = 0, size = srcShape.size(); i < size; i++) {
282  const auto &map = reassociation[i];
283  auto srcDim = srcShape[i];
284  // Iterate through dimensions expanded from the i-th dimension.
285  for (unsigned j = start; j < start + map.size(); j++) {
286  // There can be only one dynamic sized dimension among dimensions
287  // expanded from the i-th dimension in srcShape.
288  // For example, if srcDim = 8, then the expanded shape could be <2x?x2>,
289  // but not <2x?x?>.
290  if (staticDstShape[j] == ShapedType::kDynamic) {
291  // The expanded dimension has dynamic size. We compute the dimension
292  // by dividing srcDim by the product of the static dimensions.
293  StaticSize product = 1;
294  for (unsigned k = start; k < start + map.size(); k++) {
295  if (staticDstShape[k] != ShapedType::kDynamic) {
296  product *= staticDstShape[k];
297  }
298  }
299  // Compute the dynamic dimension size.
300  Value productVal = constantIndex(builder, loc, product);
301  Value dynamicSize =
302  builder.create<arith::DivUIOp>(loc, srcDim, productVal);
303  dstShape.push_back(dynamicSize);
304  } else {
305  // The expanded dimension is statically known.
306  dstShape.push_back(constantIndex(builder, loc, staticDstShape[j]));
307  }
308  }
309  start = start + map.size();
310  }
311  assert(start == staticDstShape.size());
312 }
313 
315  OpBuilder &builder, Location loc,
316  ArrayRef<ReassociationIndices> reassociation, // NOLINT
317  ValueRange srcSizes, ValueRange srcCvs, // NOLINT
318  ValueRange dstSizes, SmallVectorImpl<Value> &dstCvs) {
319  const unsigned srcRank = srcSizes.size();
320  const unsigned dstRank = dstSizes.size();
321  assert(srcRank == srcCvs.size() && "Source rank mismatch");
322  const bool isCollapse = srcRank > dstRank;
323  const ValueRange sizes = isCollapse ? srcSizes : dstSizes;
324  // Iterate over reassociation map.
325  unsigned i = 0;
326  unsigned start = 0;
327  for (const auto &map : llvm::enumerate(reassociation)) {
328  // Prepare strides information in dimension slice.
329  Value linear = constantIndex(builder, loc, 1);
330  for (unsigned j = start, end = start + map.value().size(); j < end; j++) {
331  linear = builder.create<arith::MulIOp>(loc, linear, sizes[j]);
332  }
333  // Start expansion.
334  Value val;
335  if (!isCollapse)
336  val = srcCvs[i];
337  // Iterate over dimension slice.
338  for (unsigned j = start, end = start + map.value().size(); j < end; j++) {
339  linear = builder.create<arith::DivUIOp>(loc, linear, sizes[j]);
340  if (isCollapse) {
341  const Value mul = builder.create<arith::MulIOp>(loc, srcCvs[j], linear);
342  val = val ? builder.create<arith::AddIOp>(loc, val, mul) : mul;
343  } else {
344  const Value old = val;
345  val = builder.create<arith::DivUIOp>(loc, val, linear);
346  assert(dstCvs.size() == j);
347  dstCvs.push_back(val);
348  val = builder.create<arith::RemUIOp>(loc, old, linear);
349  }
350  }
351  // Finalize collapse.
352  if (isCollapse) {
353  assert(dstCvs.size() == i);
354  dstCvs.push_back(val);
355  }
356  start += map.value().size();
357  i++;
358  }
359  assert(dstCvs.size() == dstRank);
360 }
361 
362 FlatSymbolRefAttr mlir::sparse_tensor::getFunc(ModuleOp module, StringRef name,
363  TypeRange resultType,
364  ValueRange operands,
365  EmitCInterface emitCInterface) {
366  MLIRContext *context = module.getContext();
367  auto result = SymbolRefAttr::get(context, name);
368  auto func = module.lookupSymbol<func::FuncOp>(result.getAttr());
369  if (!func) {
370  OpBuilder moduleBuilder(module.getBodyRegion());
371  func = moduleBuilder.create<func::FuncOp>(
372  module.getLoc(), name,
373  FunctionType::get(context, operands.getTypes(), resultType));
374  func.setPrivate();
375  if (static_cast<bool>(emitCInterface))
376  func->setAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(),
377  UnitAttr::get(context));
378  }
379  return result;
380 }
381 
383  OpBuilder &builder, Location loc, StringRef name, TypeRange resultType,
384  ValueRange operands, EmitCInterface emitCInterface) {
385  auto module = builder.getBlock()->getParentOp()->getParentOfType<ModuleOp>();
386  FlatSymbolRefAttr fn =
387  getFunc(module, name, resultType, operands, emitCInterface);
388  return builder.create<func::CallOp>(loc, resultType, fn, operands);
389 }
390 
393 }
394 
396  return getOpaquePointerType(builder.getContext());
397 }
398 
400  unsigned sz, Type tp, bool staticShape) {
401  if (staticShape) {
402  auto memTp = MemRefType::get({sz}, tp);
403  return builder.create<memref::AllocaOp>(loc, memTp);
404  }
405  return genAlloca(builder, loc, constantIndex(builder, loc, sz), tp);
406 }
407 
409  Type tp) {
410  auto memTp = MemRefType::get({ShapedType::kDynamic}, tp);
411  return builder.create<memref::AllocaOp>(loc, memTp, ValueRange{sz});
412 }
413 
415  Type tp) {
416  return builder.create<memref::AllocaOp>(loc, MemRefType::get({}, tp));
417 }
418 
420  ValueRange values) {
421  const unsigned sz = values.size();
422  assert(sz >= 1);
423  Value buffer = genAlloca(builder, loc, sz, values[0].getType());
424  for (unsigned i = 0; i < sz; i++) {
425  Value idx = constantIndex(builder, loc, i);
426  builder.create<memref::StoreOp>(loc, values[i], buffer, idx);
427  }
428  return buffer;
429 }
430 
432  RankedTensorType tensorTp,
433  ValueRange sizes) {
434  Type elemTp = tensorTp.getElementType();
435  auto shape = tensorTp.getShape();
436  auto memTp = MemRefType::get(shape, elemTp);
437  SmallVector<Value> dynamicSizes;
438  for (unsigned i = 0, rank = tensorTp.getRank(); i < rank; i++) {
439  if (shape[i] == ShapedType::kDynamic)
440  dynamicSizes.push_back(sizes[i]);
441  }
442  Value mem = builder.create<memref::AllocOp>(loc, memTp, dynamicSizes);
443  Value zero = constantZero(builder, loc, elemTp);
444  builder.create<linalg::FillOp>(loc, ValueRange{zero}, ValueRange{mem});
445  return mem;
446 }
447 
449  Value buffer) {
450  builder.create<memref::DeallocOp>(loc, buffer);
451 }
452 
454  Value tensor, ValueRange ivs) {
455  Value val = builder.create<tensor::ExtractOp>(loc, tensor, ivs);
456  Value cond = genIsNonzero(builder, loc, val);
457  scf::IfOp ifOp = builder.create<scf::IfOp>(loc, cond, /*else*/ false);
458  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
459  return val;
460 }
461 
462 // FIXME:
463 // 1. Dense tensors loop should be generated by loop emitter.
464 // 2. Support reduction variables to propagate SSA chains properly.
465 // 3. Change the `rank` argument to `Dimension dimRank` or `Level lvlRank`,
466 // to clarify its meaning.
468  OpBuilder &builder, Location loc, Value src, unsigned rank,
469  function_ref<void(OpBuilder &, Location, Value, ValueRange)> bodyBuilder) {
470  // `cvs` is actually the flattened coordinates array for all elements,
471  // not just for one element (since we do not `SmallVector::clear` after
472  // each iteration of the body of the loopnest.
473  SmallVector<Value> cvs;
477  const Value zero = constantIndex(builder, loc, 0);
478  const Value one = constantIndex(builder, loc, 1);
479  const auto splitSrc = genSplitSparseConstant(builder, loc, src);
480  if (splitSrc.has_value()) {
481  const Value srcCoordinates = splitSrc->first;
482  const Value srcValues = splitSrc->second;
483  lo.push_back(zero);
484  hi.push_back(linalg::createOrFoldDimOp(builder, loc, srcValues, 0));
485  st.push_back(one);
486  scf::buildLoopNest(builder, loc, lo, hi, st, {},
487  [&](OpBuilder &builder, Location loc, ValueRange ivs,
488  ValueRange /*args*/) -> scf::ValueVector {
490  builder, loc, srcCoordinates, srcValues, cvs,
491  ivs[0], rank);
492  bodyBuilder(builder, loc, val, cvs);
493  return {};
494  });
495  } else {
496  for (unsigned i = 0; i < rank; i++) {
497  lo.push_back(zero);
498  hi.push_back(linalg::createOrFoldDimOp(builder, loc, src, i));
499  st.push_back(one);
500  }
501  scf::buildLoopNest(builder, loc, lo, hi, st, {},
502  [&](OpBuilder &builder, Location loc, ValueRange ivs,
503  ValueRange /*args*/) -> scf::ValueVector {
504  Value val = genCoordsAndValueForDense(builder, loc,
505  src, cvs, ivs);
506  bodyBuilder(builder, loc, val, cvs);
507  return {};
508  });
509  }
510 }
511 
513  SmallVectorImpl<Value> &sizes,
514  Location loc, Value src) {
515  const Dimension dimRank = getSparseTensorType(src).getDimRank();
516  for (Dimension d = 0; d < dimRank; d++)
517  sizes.push_back(linalg::createOrFoldDimOp(builder, loc, src, d));
518 }
519 
521  for (; isa<scf::ForOp>(op->getParentOp()) ||
522  isa<scf::WhileOp>(op->getParentOp()) ||
523  isa<scf::ParallelOp>(op->getParentOp()) ||
524  isa<scf::IfOp>(op->getParentOp());
525  op = op->getParentOp())
526  ;
527  return op;
528 }
529 
531  OpBuilder &builder, Location loc, SparseElementsAttr attr, AffineMap order,
532  function_ref<void(ArrayRef<Value>, Value)> callback) {
533  const Dimension dimRank = getSparseTensorType(attr).getDimRank();
534  const auto coordinates = attr.getIndices().getValues<IntegerAttr>();
535  const auto values = attr.getValues().getValues<Attribute>();
536 
537  // This is like the `Element<V>` class in the runtime library, but for
538  // MLIR attributes. In the future we may want to move this out into
539  // a proper class definition to help improve code legibility (e.g.,
540  // `first` -> `coords`, `second` -> `value`) as well as being able
541  // to factor out analogues of `ElementLT<V>` for the sort below, etc.
542  using ElementAttr = std::pair<SmallVector<IntegerAttr>, Attribute>;
543 
544  // Construct the COO from the SparseElementsAttr.
546  for (size_t i = 0, nse = values.size(); i < nse; i++) {
547  elems.emplace_back();
548  elems.back().second = values[i];
549  auto &coords = elems.back().first;
550  coords.reserve(dimRank);
551  for (Dimension d = 0; d < dimRank; d++)
552  coords.push_back(coordinates[i * dimRank + d]);
553  }
554 
555  // Sorts the sparse element attribute based on coordinates.
556  std::sort(elems.begin(), elems.end(),
557  [order, dimRank](const ElementAttr &lhs, const ElementAttr &rhs) {
558  const auto &lhsCoords = lhs.first;
559  const auto &rhsCoords = rhs.first;
560  for (Dimension d = 0; d < dimRank; d++) {
561  // FIXME: This only makes sense for permutations.
562  // And since we don't check that `order` is a permutation,
563  // it can also cause OOB errors when we use `l`.
564  const Level l = order ? order.getDimPosition(d) : d;
565  if (lhsCoords[l].getInt() == rhsCoords[l].getInt())
566  continue;
567  return lhsCoords[l].getInt() < rhsCoords[l].getInt();
568  }
569  if (std::addressof(lhs) == std::addressof(rhs))
570  return false;
571  llvm_unreachable("no equal coordinate in sparse element attr");
572  });
573 
574  SmallVector<Value> cvs;
575  cvs.reserve(dimRank);
576  for (size_t i = 0, nse = values.size(); i < nse; i++) {
577  // Remap coordinates.
578  cvs.clear();
579  for (Dimension d = 0; d < dimRank; d++) {
580  auto crd = elems[i].first[d].getInt();
581  cvs.push_back(builder.create<arith::ConstantIndexOp>(loc, crd));
582  }
583  // Remap value.
584  Value val;
585  if (isa<ComplexType>(attr.getElementType())) {
586  auto valAttr = cast<ArrayAttr>(elems[i].second);
587  val = builder.create<complex::ConstantOp>(loc, attr.getElementType(),
588  valAttr);
589  } else {
590  auto valAttr = cast<TypedAttr>(elems[i].second);
591  val = builder.create<arith::ConstantOp>(loc, valAttr);
592  }
593  assert(val);
594  callback(cvs, val);
595  }
596 }
597 
599  size_t size, Value mem,
600  size_t offsetIdx, Value offsetVal) {
601 #ifndef NDEBUG
602  const auto memTp = cast<MemRefType>(mem.getType());
603  assert(memTp.getRank() == 1);
604  const DynSize memSh = memTp.getDimSize(0);
605  assert(ShapedType::isDynamic(memSh) || memSh >= static_cast<DynSize>(size));
606  assert(offsetIdx == 0 || offsetIdx < size);
607 #endif // NDEBUG
609  vs.reserve(size);
610  for (unsigned i = 0; i < size; i++) {
611  Value v = builder.create<memref::LoadOp>(loc, mem,
612  constantIndex(builder, loc, i));
613  if (i == offsetIdx && offsetVal)
614  v = builder.create<arith::AddIOp>(loc, v, offsetVal);
615  vs.push_back(v);
616  }
617  return vs;
618 }
619 
621  ValueRange vs, size_t offsetIdx, Value offsetVal) {
622 #ifndef NDEBUG
623  const size_t vsize = vs.size();
624  const auto memTp = cast<MemRefType>(mem.getType());
625  assert(memTp.getRank() == 1);
626  const DynSize memSh = memTp.getDimSize(0);
627  assert(ShapedType::isDynamic(memSh) || memSh >= static_cast<DynSize>(vsize));
628  assert(offsetIdx == 0 || offsetIdx < vsize);
629 #endif // NDEBUG
630  for (const auto &v : llvm::enumerate(vs)) {
631  const Value w =
632  (offsetIdx == v.index() && offsetVal)
633  ? builder.create<arith::AddIOp>(loc, v.value(), offsetVal)
634  : v.value();
635  builder.create<memref::StoreOp>(loc, w, mem,
636  constantIndex(builder, loc, v.index()));
637  }
638 }
639 
641  SparseTensorEncodingAttr enc,
642  ValueRange dimSizes,
643  Value valuesBuffer,
644  Value lvlCoords) {
645  // Reuse the `lvlCoords` buffer to store the level-sizes.
646  const Level lvlRank = enc.getLvlRank();
647  SmallVector<Value> lvlSizes;
648  lvlSizes.reserve(lvlRank);
649  for (Level l = 0; l < lvlRank; l++)
650  // FIXME: `toOrigDim` is deprecated.
651  lvlSizes.push_back(dimSizes[toOrigDim(enc, l)]);
652  storeAll(builder, loc, lvlCoords, lvlSizes);
653  // The memref ReshapeOp requires the sizes buffer to have a static
654  // shape.
655  const auto iTp = builder.getIndexType();
656  const SmallVector<DynSize, 1> lvlSizesShape{static_cast<DynSize>(lvlRank)};
657  const auto lvlSizesTp = MemRefType::get(lvlSizesShape, iTp);
658  lvlCoords = builder.create<memref::CastOp>(loc, lvlSizesTp, lvlCoords);
659  // Finally, create the ReshapeOp.
660  const SmallVector<DynSize> resShape(lvlRank, ShapedType::kDynamic);
661  const Type elemTp = getMemRefType(valuesBuffer).getElementType();
662  const auto resTp = MemRefType::get(resShape, elemTp);
663  return builder.create<memref::ReshapeOp>(loc, resTp, valuesBuffer, lvlCoords);
664 }
665 
668  auto tTp = llvm::cast<TensorType>(tensor.getType());
669  auto mTp = MemRefType::get(tTp.getShape(), tTp.getElementType());
670  return builder.create<bufferization::ToMemrefOp>(loc, mTp, tensor)
671  .getResult();
672 }
673 
675  Value tensor, Level lvl) {
676  const auto srcTp = getSparseTensorType(tensor);
677  const Type posTp = srcTp.getPosType();
678  const Type memTp = get1DMemRefType(posTp, /*withLayout=*/false);
679  return builder.create<ToPositionsOp>(loc, memTp, tensor,
680  builder.getIndexAttr(lvl));
681 }
682 
684  Value tensor, Level lvl, Level cooStart) {
685  const auto srcTp = getSparseTensorType(tensor);
686  const Type crdTp = srcTp.getCrdType();
687  const Type memTp = get1DMemRefType(crdTp, /*withLayout=*/lvl >= cooStart);
688  return builder.create<ToCoordinatesOp>(loc, memTp, tensor,
689  builder.getIndexAttr(lvl));
690 }
691 
693  Value tensor) {
694  const auto srcTp = getSparseTensorType(tensor);
695  const Type crdTp = srcTp.getCrdType();
696  const Type memTp = get1DMemRefType(crdTp, /*withLayout=*/false);
697  return builder.create<ToCoordinatesBufferOp>(loc, memTp, tensor);
698 }
699 
701  Value tensor) {
702  RankedTensorType srcTp = getRankedTensorType(tensor);
703  Type valTp = get1DMemRefType(srcTp.getElementType(),
704  /*withLayout=*/false);
705  return builder.create<ToValuesOp>(loc, valTp, tensor);
706 }
707 
709  Value tensor) {
710  return getDescriptorFromTensorTuple(tensor).getValMemSize(builder, loc);
711 }
712 
714  Value tensor, Dimension dim) {
715  auto enc = getSparseTensorEncoding(tensor.getType());
716  assert(enc && enc.isSlice());
717  std::optional<unsigned> offset = enc.getStaticDimSliceOffset(dim);
718  if (offset.has_value())
719  return constantIndex(builder, loc, *offset);
720  return builder.create<ToSliceOffsetOp>(loc, tensor, APInt(64, dim));
721 }
722 
724  Value tensor, Dimension dim) {
725  auto enc = getSparseTensorEncoding(tensor.getType());
726  assert(enc && enc.isSlice());
727  std::optional<unsigned> stride = enc.getStaticDimSliceStride(dim);
728  if (stride.has_value())
729  return constantIndex(builder, loc, *stride);
730  return builder.create<ToSliceStrideOp>(loc, tensor, APInt(64, dim));
731 }
#define CASE(ONAME, O)
static Value genCoordsAndValueForDense(OpBuilder &builder, Location loc, Value tensor, SmallVectorImpl< Value > &cvs, ValueRange ivs)
Generates code to read the value from tensor[ivs], and open a conditional for whether the value is no...
static Value genCoordsAndValueForSparse(OpBuilder &builder, Location loc, Value coordinates, Value values, SmallVectorImpl< Value > &cvs, Value k, unsigned rank)
Reads coordinates[k][0..rank-1] and value[k], appending the former onto cvs and returning the latter.
static std::optional< std::pair< Value, Value > > genSplitSparseConstant(OpBuilder &builder, Location loc, Value tensor)
If the tensor is a sparse constant, generates and returns the pair of the constants for the coordinat...
#define MLIR_SPARSETENSOR_FOREVERY_FIXED_O(DO)
Definition: Enums.h:61
#define MLIR_SPARSETENSOR_FOREVERY_V(DO)
Definition: Enums.h:94
static int64_t product(ArrayRef< int64_t > vals)
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:44
Attributes are known-constant values of operations.
Definition: Attributes.h:25
Type getElementType() const
Returns the element type of this memref type.
Operation * getParentOp()
Returns the closest surrounding operation that contains this block.
Definition: Block.cpp:30
This class is a general helper class for creating context-global objects like types,...
Definition: Builders.h:50
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:124
IntegerAttr getIntegerAttr(Type type, int64_t value)
Definition: Builders.cpp:238
FloatAttr getFloatAttr(Type type, double value)
Definition: Builders.cpp:261
IntegerType getI64Type()
Definition: Builders.cpp:85
IntegerType getIntegerType(unsigned width)
Definition: Builders.cpp:87
MLIRContext * getContext() const
Definition: Builders.h:55
IndexType getIndexType()
Definition: Builders.cpp:71
static DenseElementsAttr get(ShapedType type, ArrayRef< Attribute > values)
Constructs a dense elements attribute from an array of element values.
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:63
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
This class helps build Operations.
Definition: Builders.h:206
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Definition: Builders.h:416
Block * getBlock() const
Returns the current block of the builder.
Definition: Builders.h:433
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:446
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
Operation * getParentOp()
Returns the closest surrounding operation that contains this operation or nullptr if this is a top-le...
Definition: Operation.h:234
OpTy getParentOfType()
Return the closest surrounding parent operation that is of type 'OpTy'.
Definition: Operation.h:238
void setAttr(StringAttr name, Attribute value)
If the an attribute exists with the specified name, change it to the new value.
Definition: Operation.h:560
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
bool isF64() const
Definition: Types.cpp:52
bool isInteger(unsigned width) const
Return true if this is an integer type with the specified width.
Definition: Types.cpp:59
bool isIndex() const
Definition: Types.cpp:56
bool isF32() const
Definition: Types.cpp:51
bool isIntOrIndex() const
Return true if this is an integer (of any signedness) or an index type.
Definition: Types.cpp:113
bool isF16() const
Definition: Types.cpp:49
unsigned getIntOrFloatBitWidth() const
Return the bit width of an integer or a float type, assert failure on other types.
Definition: Types.cpp:123
bool isBF16() const
Definition: Types.cpp:48
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:372
type_range getTypes() const
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:93
Type getType() const
Return the type of this value.
Definition: Value.h:122
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:20
Specialization of arith.constant op that returns an integer of index type.
Definition: Arith.h:90
Value getValMemSize(OpBuilder &builder, Location loc) const
Dimension getDimRank() const
Returns the dimension-rank.
BaseMemRefType getMemRefType(Value value, const BufferizationOptions &options, MemRefLayoutAttrInterface layout={}, Attribute memorySpace=nullptr)
Return a MemRefType to which the type of the given value can be bufferized.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:285
Value createOrFoldDimOp(OpBuilder &b, Location loc, Value val, int64_t dim)
Create one memref::DimOp or tensor::DimOp depending on the type of val.
Definition: LinalgOps.cpp:88
LoopNest buildLoopNest(OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs, ValueRange steps, ValueRange iterArgs, function_ref< ValueVector(OpBuilder &, Location, ValueRange, ValueRange)> bodyBuilder=nullptr)
Creates a perfect nest of "for" loops, i.e.
Definition: SCF.cpp:667
SmallVector< Value > ValueVector
An owning vector of values, handy to return from functions.
Definition: SCF.h:70
TypedAttr getOneAttr(Builder &builder, Type tp)
Generates a 1-valued attribute of the given type.
Value genToCoordinatesBuffer(OpBuilder &builder, Location loc, Value tensor)
Infers the result type and generates ToCoordinatesBufferOp.
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:350
void foreachInSparseConstant(OpBuilder &builder, Location loc, SparseElementsAttr attr, AffineMap order, function_ref< void(ArrayRef< Value >, Value)> callback)
Iterate over a sparse constant, generates constantOp for value and coordinates.
Value constantZero(OpBuilder &builder, Location loc, Type tp)
Generates a 0-valued constant of the given type.
Definition: CodegenUtils.h:328
Value allocaBuffer(OpBuilder &builder, Location loc, ValueRange values)
Generates a temporary buffer, initializes it with the given contents, and returns it as type memref<?...
Value reshapeValuesToLevels(OpBuilder &builder, Location loc, SparseTensorEncodingAttr enc, ValueRange dimSizes, Value valuesBuffer, Value lvlCoords)
Reshapes the linear values buffer for an annotated all dense sparse tensor to match the shape of the ...
OverheadType posTypeEncoding(SparseTensorEncodingAttr enc)
Returns the OverheadType for position overhead storage.
OverheadType
Encoding of overhead types (both position overhead and coordinate overhead), for "overloading" @newSp...
Definition: Enums.h:49
uint64_t Dimension
The type of dimension identifiers, and dimension-ranks.
Definition: SparseTensor.h:40
OverheadType crdTypeEncoding(SparseTensorEncodingAttr enc)
Returns the OverheadType for coordinate overhead storage.
SparseTensorType getSparseTensorType(T t)
Convenience method to abbreviate wrapping getRankedTensorType.
Value genToValues(OpBuilder &builder, Location loc, Value tensor)
Infers the result type and generates ToValuesOp.
Value genValueForDense(OpBuilder &builder, Location loc, Value tensor, ValueRange ivs)
Generates code to read the value from tensor[ivs].
uint64_t Level
The type of level identifiers, and level-ranks.
Definition: SparseTensor.h:46
TypedValue< BaseMemRefType > genToMemref(OpBuilder &builder, Location loc, Value tensor)
Value genToPositions(OpBuilder &builder, Location loc, Value tensor, Level lvl)
Infers the result type and generates ToPositionsOp.
OverheadType overheadTypeEncoding(unsigned width)
Converts an overhead storage bitwidth to its internal type-encoding.
MemRefType get1DMemRefType(Type etp, bool withLayout)
Generates a 1D MemRefType with a dynamic size.
Definition: CodegenUtils.h:258
StringRef overheadTypeFunctionSuffix(OverheadType ot)
Convert OverheadType to its function-name suffix.
PrimaryType
Encoding of the elemental type, for "overloading" @newSparseTensor.
Definition: Enums.h:80
RankedTensorType getRankedTensorType(T &&t)
Convenience method to abbreviate casting getType().
Definition: SparseTensor.h:96
int64_t DynSize
The type for individual components of a compile-time shape.
Definition: SparseTensor.h:52
PrimaryType primaryTypeEncoding(Type elemTp)
Converts a primary storage type to its internal type-encoding.
DEPRECATED Dimension toOrigDim(SparseTensorEncodingAttr enc, Level l)
[deprecated] Convenience method to translate the given level to the corresponding dimension.
Operation * getTop(Operation *op)
Scans to top of generated loop.
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 createOrFoldSliceStrideOp(OpBuilder &builder, Location loc, Value tensor, Dimension dim)
Generates code to retrieve the slice slice for the sparse tensor slice, return a constant if the offs...
Type getOpaquePointerType(MLIRContext *ctx)
Returns the equivalent of void* for opaque arguments to the execution engine.
SparseTensorEncodingAttr getSparseTensorEncoding(Type type)
Convenience method to get a sparse encoding attribute from a type.
Value genIsNonzero(OpBuilder &builder, Location loc, Value v)
Generates the comparison v != 0 where v is of numeric type.
Value genIndexLoad(OpBuilder &builder, Location loc, Value mem, Value s)
Generates a pointer/index load from the sparse storage scheme.
SparseTensorDescriptor getDescriptorFromTensorTuple(Value tensor)
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...
SmallVector< Value > loadAll(OpBuilder &builder, Location loc, size_t size, Value mem, size_t offsetIdx=0, Value offsetVal=Value())
Loads size-many values from the memref, which must have rank-1 and size greater-or-equal to size.
Value genToCoordinates(OpBuilder &builder, Location loc, Value tensor, Level lvl, Level cooStart)
Infers the result type and generates ToCoordinatesOp.
int64_t StaticSize
The type for individual components of a compile-time shape which are known not to be ShapedType::kDyn...
Definition: SparseTensor.h:56
void reshapeCvs(OpBuilder &builder, Location loc, ArrayRef< ReassociationIndices > reassociation, ValueRange srcSizes, ValueRange srcCvs, ValueRange dstSizes, SmallVectorImpl< Value > &dstCvs)
Reshape coordinates during a reshaping operation.
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.
Value genCast(OpBuilder &builder, Location loc, Value value, Type dstTy)
Add type casting between arith and index types when needed.
StringRef primaryTypeFunctionSuffix(PrimaryType pt)
Convert PrimaryType to its function-name suffix.
Value createOrFoldSliceOffsetOp(OpBuilder &builder, Location loc, Value tensor, Dimension dim)
Generates code to retrieve the slice offset for the sparse tensor slice, return a constant if the off...
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 genValMemSize(OpBuilder &builder, Location loc, Value tensor)
Generates code to retrieve the values size for the sparse tensor.
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.
void storeAll(OpBuilder &builder, Location loc, Value mem, ValueRange vs, size_t offsetIdx=0, Value offsetVal=Value())
Stores all the values of vs into the memref mem, which must have rank-1 and size greater-or-equal to ...
void genReshapeDstShape(OpBuilder &builder, Location loc, SmallVectorImpl< Value > &dstShape, ArrayRef< Value > srcShape, ArrayRef< StaticSize > staticDstShape, ArrayRef< ReassociationIndices > reassociation)
Computes the shape of destination tensor of a reshape operator.
This header declares functions that assist transformations in the MemRef dialect.
Value convertScalarToDtype(OpBuilder &b, Location loc, Value operand, Type toType, bool isUnsignedCast)
Converts a scalar value operand to type toType.
Definition: Utils.cpp:168
std::conditional_t< std::is_same_v< Ty, mlir::Type >, mlir::Value, detail::TypedValue< Ty > > TypedValue
If Ty is mlir::Type this will select Value instead of having a wrapper around it.
Definition: Value.h:489
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
Eliminates variable at the specified position using Fourier-Motzkin variable elimination.