MLIR  21.0.0git
CodegenEnv.cpp
Go to the documentation of this file.
1 //===- CodegenEnv.cpp - Code generation environment class ----------------===//
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 "CodegenEnv.h"
10 
15 
16 #include <optional>
17 
18 using namespace mlir;
19 using namespace mlir::sparse_tensor;
20 
21 //===----------------------------------------------------------------------===//
22 // Code generation environment helper functions
23 //===----------------------------------------------------------------------===//
24 
25 /// Returns true if tensor materializes uninitialized into the computation.
26 static bool isMaterializing(Value val) {
27  return val.getDefiningOp<tensor::EmptyOp>() ||
28  val.getDefiningOp<bufferization::AllocTensorOp>();
29 }
30 
31 /// Sorts the dependent loops such that it is ordered in the same sequence in
32 /// which loops will be generated.
33 static void sortDependentLoops(std::vector<LoopCoeffPair> &target) {
34  llvm::sort(target, [](const LoopCoeffPair &l, const LoopCoeffPair &r) {
35  assert(std::addressof(l) == std::addressof(r) || l != r);
36  return l.first < r.first;
37  });
38 }
39 //===----------------------------------------------------------------------===//
40 // Code generation environment constructor and general methods
41 //===----------------------------------------------------------------------===//
42 
43 CodegenEnv::CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts,
44  unsigned numTensors, unsigned numLoops, unsigned maxRank)
45  : linalgOp(linop), sparseOptions(opts),
46  latticeMerger(numTensors, numLoops, maxRank), loopEmitter(),
47  sparseOut(nullptr), outerParNest(-1u), insChain(), expValues(),
48  expFilled(), expAdded(), expCount(), redVal(), redExp(detail::kInvalidId),
49  redCustom(detail::kInvalidId), redValidLexInsert() {}
50 
51 LogicalResult CodegenEnv::initTensorExp() {
52  // Builds the tensor expression for the Linalg operation in SSA form.
53  std::optional<ExprId> optExp = latticeMerger.buildTensorExpFromLinalg(op());
54  if (!optExp || !isAdmissibleTensorExp(*optExp))
55  return failure();
56 
57  tensorExp = *optExp;
58  return success();
59 }
60 
62  assert(insChain == nullptr && "must only start emitting once");
63  if (sparseOut) {
64  insChain = sparseOut->get();
65  latticeMerger.setHasSparseOut(true);
66  }
67 
68  // Sort the related loop array such that they are in the same order as they
69  // appears on the topoOrder.
70  // TODO: since we only handle affine addition for slice based codegen, and
71  // addition is assoicative, the order how we evaluate the expression does
72  // not matter. However, to support multiplication, the order of the loop
73  // index should match the evaluation order to the affine expression AST.
74 
75  // Initialize loop emitter.
76  SmallVector<Value> tensors; // input tensors passed to loop emitter
77  for (OpOperand &t : linalgOp->getOpOperands()) {
78  tensors.push_back(t.get());
79  const TensorId tid = makeTensorId(t.getOperandNumber());
80  const Level lvlRank = linalgOp.getMatchingIndexingMap(&t).getNumResults();
81  const auto enc = getSparseTensorEncoding(t.get().getType());
82  (void)enc;
83  assert(!enc || lvlRank == enc.getLvlRank());
84  for (Level lvl = 0; lvl < lvlRank; lvl++)
85  sortDependentLoops(latticeMerger.getDependentLoops(tid, lvl));
86  }
87  loopEmitter.initialize(
88  tensors,
89  StringAttr::get(linalgOp.getContext(),
90  linalg::GenericOp::getOperationName()),
91  /*hasOutput=*/true,
92  /*isSparseOut=*/sparseOut != nullptr, /*numLoops=*/getLoopNum(),
93  // TODO: compute the map and pass it to loop emitter directly instead of
94  // passing in a callback.
95  /*dependentLvlGetter=*/
96  [this](TensorId t, Level lvl) -> std::vector<LoopCoeffPair> {
97  return merger().getDependentLoops(t, lvl);
98  },
99  emitStrategy);
100 }
101 
102 std::optional<Operation *> CodegenEnv::genLoopBoundary(
103  function_ref<std::optional<Operation *>(MutableArrayRef<Value> parameters)>
104  callback) {
105  SmallVector<Value> params;
106  if (isReduc()) {
107  params.push_back(redVal);
108  if (isValidLexInsert())
109  params.push_back(redValidLexInsert);
110  } else {
111  assert(!isValidLexInsert());
112  }
113  if (isExpand())
114  params.push_back(expCount);
115  if (insChain != nullptr)
116  params.push_back(insChain);
117  auto r = callback(params); // may update parameters
118  unsigned i = 0;
119  if (isReduc()) {
120  updateReduc(params[i++]);
121  if (isValidLexInsert())
122  updateValidLexInsert(params[i++]);
123  }
124  if (isExpand())
125  updateExpandCount(params[i++]);
126  if (insChain != nullptr)
127  updateInsertionChain(params[i]);
128  return r;
129 }
130 
131 //===----------------------------------------------------------------------===//
132 // Code generation environment verify functions.
133 //===----------------------------------------------------------------------===//
134 
136  // We reject any expression that makes a reduction from `-outTensor`, as those
137  // expressions create a dependency between the current iteration (i) and the
138  // previous iteration (i-1). It would require iterating over the whole
139  // coordinate space, which prevent exploiting sparsity for faster code.
140  for (utils::IteratorType it : linalgOp.getIteratorTypesArray()) {
141  if (it == utils::IteratorType::reduction) {
142  if (latticeMerger.hasNegateOnOut(exp))
143  return false;
144  break;
145  }
146  }
147 
148  OpOperand *lhs = linalgOp.getDpsInitOperand(0);
149  const TensorId tensor = makeTensorId(lhs->getOperandNumber());
150  // An non-annotated output tensor is assumed dense, and becomes a random
151  // access n-dim memref. Admissible since insertions cannot occur.
152  if (getSparseTensorType(lhs->get()).isAllDense())
153  return true;
154 
155  // A tensor expression with a sparse output tensor that changes its values
156  // but not its nonzero structure, an operation called "simply dynamic" in
157  // [Bik96,Ch9], is also admissible without special env.
158  if (latticeMerger.isSingleCondition(tensor, exp))
159  return true;
160 
161  // Accept "truly dynamic" if the output tensor materializes uninitialized
162  // into the computation and insertions occur in lexicographic index order.
163  sparseOut = lhs;
164 
165  // Find the outermost parallel nest to determine whether compress/expand is
166  // needed.
167  outerParNest = 0;
168  const auto iteratorTypes = linalgOp.getIteratorTypesArray();
169  for (unsigned i = 0, e = getLoopNum(); i < e; i++) {
170  if (linalg::isReductionIterator(iteratorTypes[i]))
171  break; // terminate at first reduction
172  outerParNest++;
173  }
174 
175  // Inadmissible kernel should have already been rejected by the previous
176  // path during loop scheduling.
177  assert(static_cast<int64_t>(outerParNest) >=
178  linalgOp.getRank(linalgOp.getDpsInitOperand(0)) - 1);
179  return isMaterializing(lhs->get());
180 }
181 
182 //===----------------------------------------------------------------------===//
183 // Code generation environment topological sort methods
184 //===----------------------------------------------------------------------===//
185 
187  return loopEmitter.getLoopIV(i);
188 }
189 
190 //===----------------------------------------------------------------------===//
191 // Code generation environment sparse tensor output and expansion methods
192 //===----------------------------------------------------------------------===//
193 
195  assert(sparseOut != nullptr && insChain != nullptr);
196  insChain = chain;
197 }
198 
199 bool CodegenEnv::atExpandLevel(OpOperand *o, unsigned rank, LoopId n) const {
200  return sparseOut == o && outerParNest == static_cast<LoopId>(rank - 1) &&
201  outerParNest == n;
202 }
203 
204 void CodegenEnv::startExpand(Value values, Value filled, Value added,
205  Value count) {
206  assert(sparseOut != nullptr && expValues == nullptr);
207  expValues = values;
208  expFilled = filled;
209  expAdded = added;
210  expCount = count;
211 }
212 
214  assert(sparseOut != nullptr && expValues != nullptr);
215  expCount = count;
216 }
217 
219  assert(sparseOut != nullptr && expValues != nullptr);
220  expValues = expFilled = expAdded = expCount = Value();
221 }
222 
223 //===----------------------------------------------------------------------===//
224 // Code generation environment reduction methods
225 //===----------------------------------------------------------------------===//
226 
228  assert(!isReduc() && exp != detail::kInvalidId && val);
229  redExp = exp;
230  redVal = val;
231  latticeMerger.setExprValue(exp, val);
232 }
233 
235  assert(isReduc() && val);
236  redVal = val;
237  latticeMerger.clearExprValue(redExp);
238  latticeMerger.setExprValue(redExp, val);
239 }
240 
242  assert(isReduc());
243  Value val = redVal;
244  redVal = val;
245  latticeMerger.clearExprValue(redExp);
246  redExp = detail::kInvalidId;
247  return val;
248 }
249 
251  assert(!isValidLexInsert() && isReduc() && val);
252  redValidLexInsert = val;
253 }
254 
256  assert(redValidLexInsert && isReduc() && val);
257  redValidLexInsert = val;
258 }
259 
261  assert(isValidLexInsert() && !isReduc());
262  redValidLexInsert = Value();
263 }
264 
266  assert(!isCustomReduc() && exp != detail::kInvalidId);
267  redCustom = exp;
268 }
269 
271  assert(isCustomReduc());
272  return dyn_cast<sparse_tensor::ReduceOp>(exp(redCustom).op).getIdentity();
273 }
274 
276  assert(isCustomReduc());
277  redCustom = detail::kInvalidId;
278 }
static bool isMaterializing(Value val)
Returns true if tensor materializes uninitialized into the computation.
Definition: CodegenEnv.cpp:26
static void sortDependentLoops(std::vector< LoopCoeffPair > &target)
Sorts the dependent loops such that it is ordered in the same sequence in which loops will be generat...
Definition: CodegenEnv.cpp:33
IRValueT get() const
Return the current value being used by this operand.
Definition: UseDefLists.h:160
This class represents an operand of an operation.
Definition: Value.h:257
unsigned getOperandNumber()
Return which operand this is in the OpOperand list of the Operation.
Definition: Value.cpp:228
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:20
void startReduc(ExprId exp, Value val)
Definition: CodegenEnv.cpp:227
void updateValidLexInsert(Value val)
Definition: CodegenEnv.cpp:255
std::optional< Operation * > genLoopBoundary(function_ref< std::optional< Operation * >(MutableArrayRef< Value > parameters)> callback)
Generates loop boundary statements (entering/exiting loops).
Definition: CodegenEnv.cpp:102
bool isAdmissibleTensorExp(ExprId e)
Whether the tensor expression is admissible for codegen.
Definition: CodegenEnv.cpp:135
bool atExpandLevel(OpOperand *o, unsigned rank, LoopId n) const
Definition: CodegenEnv.cpp:199
CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts, unsigned numTensors, unsigned numLoops, unsigned maxRank)
Constructs a code generation environment which can be passed around during sparsification for bookkee...
Definition: CodegenEnv.cpp:43
constexpr TensorId makeTensorId(unsigned t) const
Definition: CodegenEnv.h:72
void startExpand(Value values, Value filled, Value added, Value count)
Definition: CodegenEnv.cpp:204
unsigned getLoopNum() const
Definition: CodegenEnv.h:89
void updateInsertionChain(Value chain)
Definition: CodegenEnv.cpp:194
void startCustomReduc(ExprId exp)
Definition: CodegenEnv.cpp:265
linalg::GenericOp op() const
Definition: CodegenEnv.h:50
Value getLoopVar(LoopId i) const
Returns the induction-variable for the given loop.
Definition: CodegenEnv.cpp:186
void startEmit(SparseEmitStrategy emitStrategy)
Definition: CodegenEnv.cpp:61
const TensorExp & exp(ExprId e) const
Definition: CodegenEnv.h:81
void updateExpandCount(Value count)
Definition: CodegenEnv.cpp:213
void startValidLexInsert(Value val)
Definition: CodegenEnv.cpp:250
void initialize(ValueRange tensors, StringAttr loopTag=nullptr, bool hasOutput=false, bool isSparseOut=false, unsigned numLoops=0, DependentLvlGetter getter=nullptr, SparseEmitStrategy emitStrategy=SparseEmitStrategy::kFunctional)
Takes an array of input tensors, which the generated loops will iterate over.
Value getLoopIV(LoopId n) const
Gets loop induction variable for the given loop.
Definition: LoopEmitter.h:175
void setHasSparseOut(bool s)
Sets whether the output tensor is sparse or not.
Definition: Merger.h:469
bool isSingleCondition(TensorId t, ExprId e) const
Returns true if given tensor iterates only in the given tensor expression.
Definition: Merger.cpp:577
bool hasNegateOnOut(ExprId e) const
Returns true if the expression contains a negation on output tensor.
Definition: Merger.cpp:544
std::optional< ExprId > buildTensorExpFromLinalg(linalg::GenericOp op)
Builds a tensor expression from the given Linalg operation.
Definition: Merger.cpp:1193
void clearExprValue(ExprId e)
Clears the value associated with the expression.
Definition: Merger.h:567
std::vector< LoopCoeffPair > & getDependentLoops(TensorId t, Level lvl)
Returns the list of loop indices which appear in the non-trivial index expression on t_l,...
Definition: Merger.h:488
void setExprValue(ExprId e, Value v)
Sets the expression to have the associated value.
Definition: Merger.h:559
bool isAllDense() const
Returns true for tensors where every level is dense.
bool isReductionIterator(utils::IteratorType iteratorType)
Check if iterator type has "reduction" semantics.
Definition: Utils.cpp:242
static constexpr unsigned kInvalidId
A constant serving as the canonically invalid identifier, regardless of the identifier type.
Definition: Merger.h:30
uint64_t Level
The type of level identifiers and level-ranks.
Definition: SparseTensor.h:42
unsigned LoopId
Loop identifiers.
Definition: Merger.h:38
SparseTensorEncodingAttr getSparseTensorEncoding(Type type)
Convenience method to get a sparse encoding attribute from a type.
SparseTensorType getSparseTensorType(Value val)
Convenience methods to obtain a SparseTensorType from a Value.
unsigned ExprId
TensorExp identifiers.
Definition: Merger.h:48
std::pair< LoopId, unsigned > LoopCoeffPair
A pair of loop id and its coefficients.
Definition: Merger.h:64
unsigned TensorId
Tensor identifiers, chosen to be the BlockArgument::getArgNumber of the value passed to Merger::build...
Definition: Merger.h:35
Include the generated interface declarations.
SparseEmitStrategy
Defines a scope for reinterpret map pass.
Definition: Passes.h:52
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
Options for the Sparsification pass.
Definition: Passes.h:93