MLIR  14.0.0git
LinalgInterfaceImpl.cpp
Go to the documentation of this file.
1 //===- LinalgInterfaceImpl.cpp - Linalg Impl. of BufferizableOpInterface --===//
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 
14 #include "mlir/IR/Dialect.h"
15 #include "mlir/IR/Operation.h"
16 
17 using namespace mlir;
18 using namespace linalg;
19 using namespace comprehensive_bufferize;
20 using namespace mlir::bufferization;
21 
22 namespace {
23 
24 // TODO: Ops in the linalg dialect can directly implement this interface.
25 
26 /// Generic conversion for any LinalgOp on tensors.
27 static LogicalResult bufferizeLinalgOp(RewriterBase &rewriter, LinalgOp op,
28  const BufferizationState &state) {
29  // Take a guard before anything else.
30  OpBuilder::InsertionGuard g(rewriter);
31  rewriter.setInsertionPoint(op);
32 
33  // Nothing to do. This op is already bufferized.
34  if (op.hasBufferSemantics())
35  return success();
36 
37  // Ensure op has only tensors. Allow mixed tensor-buffer mode on a per-need
38  // basis.
39  if (!op.hasTensorSemantics())
40  return op->emitError() << "op does not have tensor semantics";
41 
42  // New input operands for the cloned op.
43  SmallVector<Value> newInputBuffers;
44  newInputBuffers.reserve(op.getNumInputs());
45  for (OpOperand *opOperand : op.getInputOperands()) {
46  if (op.isScalar(opOperand)) {
47  newInputBuffers.push_back(opOperand->get());
48  continue;
49  }
50  // Input operands are never written to.
51  newInputBuffers.push_back(
52  *state.getBuffer(rewriter, *opOperand, /*forceInPlace=*/true));
53  }
54 
55  // New output operands for the cloned op.
56  SmallVector<Value> newOutputBuffers;
57  for (OpResult opResult : op->getOpResults()) {
58  SmallVector<OpOperand *> aliasingOpOperands =
59  state.getAliasingOpOperand(opResult);
60  assert(aliasingOpOperands.size() == 1 && "expected 1 OpOperand");
61  FailureOr<Value> resultBuffer =
62  state.getBuffer(rewriter, *aliasingOpOperands.front());
63  if (failed(resultBuffer))
64  return failure();
65  newOutputBuffers.push_back(*resultBuffer);
66  }
67 
68  // Merge input/output operands.
69  SmallVector<Value> newOperands = newInputBuffers;
70  newOperands.append(newOutputBuffers.begin(), newOutputBuffers.end());
71 
72  // Set insertion point now that potential alloc/dealloc are introduced.
73  rewriter.setInsertionPoint(op);
74  // Clone the op, but use the new operands. Move the existing block into the
75  // new op. Since the new op does not have any tensor results, it does not
76  // return anything.
77  assert(op->getNumRegions() == 1 && "expected that op has 1 region");
78  auto newOp = cast<LinalgOp>(op.cloneWithoutRegions(
79  rewriter, op.getLoc(), /*resultTypes=*/TypeRange{}, newOperands));
80  rewriter.inlineRegionBefore(op->getRegion(0), newOp->getRegion(0),
81  newOp->getRegion(0).begin());
82 
83  // Replace the results of the old op with the new output buffers.
84  replaceOpWithBufferizedValues(rewriter, op, newOutputBuffers);
85 
86  return success();
87 }
88 
89 /// Linalg OpResults usually bufferize inplace with their tied (output
90 /// OpOperands. However, if an output OpOperand is not used in the computation,
91 /// it is better to bufferize inplace with an actually used input OpOperand;
92 /// less memory will be touched that way.
93 ///
94 /// Example:
95 /// O(i, j) = A(i, j) + B(j) --> bufferizes inplace to: A(i, j) += B(j)
96 ///
97 /// O(i, j) = A(j, i) + B(j) --> cannot bufferize inplace with A because
98 /// indexing maps are not identical
99 ///
100 /// O(i, j) += A(i, j) + B(j) --> Output is used in computation.
101 /// This could bufferize inplace with A:
102 /// A(i, j) += O(i, j) + B(j)
103 /// However, we choose to bufferize inplace with O here, as there is no clear
104 /// benefit of choosing A. TODO: We may want to consider both options and make
105 /// an informed decision during analysis in the future.
106 static DenseMap<OpOperand *, OpResult> computeAliasingPairs(LinalgOp op) {
108  for (OpResult opResult : op->getOpResults()) {
109  OpOperand *tiedOperand =
110  op.getOutputTensorOperands()[opResult.getResultNumber()];
111  AffineMap outputIndexingMap = op.getTiedIndexingMap(tiedOperand);
112  bool onlyParallelIterators = op.getNumParallelLoops() == op.getNumLoops();
113  bool tiedOperandUsed = op.payloadUsesValueFromOperand(tiedOperand);
114 
115  // If the output arg is used in the computation or at least one iterator is
116  // not parallel, try to bufferize inplace with the corresponding output
117  // tensor.
118  if (tiedOperandUsed || !onlyParallelIterators) {
119  mapping[tiedOperand] = opResult;
120  continue;
121  }
122 
123  // Otherwise, try to bufferize inplace with one of the inputs.
124  OpOperand *chosenOperand = nullptr;
125  for (OpOperand *opOperand : op.getInputTensorOperands()) {
126  if (opOperand->get().getType() != opResult.getType())
127  continue;
128  if (!op.payloadUsesValueFromOperand(opOperand))
129  continue;
130  if (op.getTiedIndexingMap(opOperand) != outputIndexingMap)
131  continue;
132  // No other OpResult bufferizes aliases with this OpOperand.
133  if (mapping.count(opOperand))
134  continue;
135  assert(op.getTiedIndexingMap(opOperand).isProjectedPermutation() &&
136  "expected projected permutation");
137  chosenOperand = opOperand;
138  break;
139  }
140 
141  // No suitable input tensor found. Use output tensor.
142  // TODO: This operand could bufferize inplace with OpOperands that have the
143  // correct type, even if they are not used inside the computation.
144  if (!chosenOperand)
145  chosenOperand = tiedOperand;
146 
147  mapping[chosenOperand] = opResult;
148  }
149  return mapping;
150 }
151 
152 /// Bufferization of linalg.generic. Replace with a new linalg.generic that
153 /// operates entirely on memrefs.
154 template <typename OpTy>
155 struct LinalgOpInterface
156  : public BufferizableOpInterface::ExternalModel<LinalgOpInterface<OpTy>,
157  OpTy> {
158  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
159  const BufferizationState &state) const {
160  // Operand is read if it is used in the computation.
161  auto genericOp = cast<linalg::LinalgOp>(op);
162  return genericOp.payloadUsesValueFromOperand(&opOperand);
163  }
164 
165  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
166  const BufferizationState &state) const {
167  // Operand is written to if it has an aliasing OpResult. For more details,
168  // see `computeAliasingPairs`.
169  auto bufferizableOp = cast<BufferizableOpInterface>(op);
170  return static_cast<bool>(
171  bufferizableOp.getAliasingOpResult(opOperand, state));
172  }
173 
175  getAliasingOpOperand(Operation *op, OpResult opResult,
176  const BufferizationState &state) const {
177  auto genericOp = cast<linalg::LinalgOp>(op);
178 
179  // Aliasing OpOperand/OpResult pairs are computed by `computeAliasingPairs`.
180  DenseMap<OpOperand *, OpResult> pairs = computeAliasingPairs(genericOp);
181  for (OpOperand *opOperand : genericOp.getInputAndOutputOperands())
182  if (pairs[opOperand] == opResult)
183  return {opOperand};
184  return {};
185  }
186 
187  OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
188  const BufferizationState &state) const {
189  auto genericOp = cast<linalg::LinalgOp>(op);
190 
191  // Aliasing OpOperand/OpResult pairs are computed by `computeAliasingPairs`.
192  DenseMap<OpOperand *, OpResult> pairs = computeAliasingPairs(genericOp);
193  return pairs[&opOperand];
194  }
195 
196  BufferRelation bufferRelation(Operation *op, OpResult opResult,
197  const BufferizationState &state) const {
198  return BufferRelation::Equivalent;
199  }
200 
201  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
202  const BufferizationState &state) const {
203  return bufferizeLinalgOp(rewriter, cast<LinalgOp>(op), state);
204  }
205 };
206 
207 struct InitTensorOpInterface
208  : public BufferizableOpInterface::ExternalModel<InitTensorOpInterface,
209  linalg::InitTensorOp> {
210  bool isMemoryWrite(Operation *op, OpResult opResult,
211  const BufferizationState &state) const {
212  // InitTensorOps allocate but do not write.
213  return false;
214  }
215 
216  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
217  const BufferizationState &state) const {
218  auto initTensorOp = cast<linalg::InitTensorOp>(op);
219 
220  // The InitTensorOp may have been eliminated.
221  if (initTensorOp->getUses().empty())
222  return success();
223 
224  FailureOr<Value> alloc =
225  createAlloc(rewriter, initTensorOp->getLoc(), initTensorOp.result(),
226  state.getOptions().createDeallocs, state.getOptions());
227  if (failed(alloc))
228  return failure();
229  replaceOpWithBufferizedValues(rewriter, op, *alloc);
230  return success();
231  }
232 };
233 
234 /// Bufferization of linalg.tiled_loop. Replace with a new linalg.tiled_loop
235 /// that operates entirely on memrefs.
236 struct TiledLoopOpInterface
237  : public BufferizableOpInterface::ExternalModel<TiledLoopOpInterface,
238  linalg::TiledLoopOp> {
239  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
240  const BufferizationState &state) const {
241  auto tiledLoopOp = cast<linalg::TiledLoopOp>(op);
242 
243  // linalg.tiled_loop operands alone do not bufferize to a memory read, but
244  // one of the uses of their matching bbArgs may.
245  return state.isValueRead(tiledLoopOp.getTiedBlockArgument(opOperand));
246  }
247 
248  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
249  const BufferizationState &state) const {
250  auto bufferizableOp = cast<BufferizableOpInterface>(op);
251 
252  // Only operands with an aliasing OpResult (i.e., output operands) bufferize
253  // to a memory write.
254  return static_cast<bool>(
255  bufferizableOp.getAliasingOpResult(opOperand, state));
256  }
257 
258  OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
259  const BufferizationState &state) const {
260  auto tiledLoopOp = cast<linalg::TiledLoopOp>(op);
261 
262  // Output operands are tied to their corresponding OpResults.
263  return tiledLoopOp.getTiedOpResult(opOperand);
264  }
265 
266  BufferRelation bufferRelation(Operation *op, OpResult opResult,
267  const BufferizationState &state) const {
268  return BufferRelation::Equivalent;
269  }
270 
271  bool isWritable(Operation *op, Value value,
272  const BufferizationState &state) const {
273  // Interestingly, linalg::TiledLoopOp's bbArgs can **always** be viewed
274  // inplace from the perspective of nested ops:
275  // 1. Either the matching iter operand is not bufferized inplace and an
276  // alloc + optional copy makes the bbArg itself inplaceable.
277  // 2. Or the matching iter operand is bufferized inplace and bbArg just
278  // bufferizes to that too.
279  return true;
280  }
281 
282  bool isAllocationHoistingBarrier(Operation *op) const { return true; }
283 
284  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
285  const BufferizationState &state) const {
286  auto tiledLoopOp = cast<linalg::TiledLoopOp>(op);
287 
288  // Compute new inputs, outputs and results.
289  SmallVector<Value> newInputs, newOutputs, newResults;
290  for (unsigned i = tiledLoopOp.getNumControlOperands();
291  i < tiledLoopOp->getNumOperands(); ++i) {
292  OpOperand &operand = tiledLoopOp->getOpOperand(i);
293  Value rewrittenValue = operand.get();
294  if (rewrittenValue.getType().isa<TensorType>()) {
295  FailureOr<Value> bufferOrFailure = state.getBuffer(rewriter, operand);
296  if (failed(bufferOrFailure))
297  return failure();
298  rewrittenValue = *bufferOrFailure;
299  }
300  if (i <
301  tiledLoopOp.getNumControlOperands() + tiledLoopOp.getNumInputs()) {
302  newInputs.push_back(rewrittenValue);
303  } else {
304  newOutputs.push_back(rewrittenValue);
305  if (operand.get().getType().isa<TensorType>())
306  newResults.push_back(rewrittenValue);
307  }
308  }
309 
310  // Create new TiledLoopOp.
311  auto newTiledLoopOp = rewriter.create<TiledLoopOp>(
312  tiledLoopOp.getLoc(), tiledLoopOp.lowerBound(),
313  tiledLoopOp.upperBound(), tiledLoopOp.step(), newInputs, newOutputs,
314  tiledLoopOp.iterator_types(), tiledLoopOp.distribution_types());
315 
316  // Remove terminator.
317  if (!newTiledLoopOp.getBody()->empty())
318  rewriter.eraseOp(tiledLoopOp.getBody()->getTerminator());
319 
320  // Compute new loop body arguments.
321  SmallVector<Value> newBlockArgs, newRegionInOutArgs, oldRegionInOutArgs;
322  ValueRange newInductionVars = newTiledLoopOp.getInductionVars();
323  newBlockArgs.append(newInductionVars.begin(), newInductionVars.end());
324 
325  ValueRange newRegionInArgs = newTiledLoopOp.getRegionInputArgs();
326  ValueRange newRegionOutArgs = newTiledLoopOp.getRegionOutputArgs();
327  newRegionInOutArgs.append(newRegionInArgs.begin(), newRegionInArgs.end());
328  newRegionInOutArgs.append(newRegionOutArgs.begin(), newRegionOutArgs.end());
329 
330  ValueRange oldRegionInArgs = tiledLoopOp.getRegionInputArgs();
331  ValueRange oldRegionOutArgs = tiledLoopOp.getRegionOutputArgs();
332  oldRegionInOutArgs.append(oldRegionInArgs.begin(), oldRegionInArgs.end());
333  oldRegionInOutArgs.append(oldRegionOutArgs.begin(), oldRegionOutArgs.end());
334  assert(newRegionInArgs.size() == oldRegionInArgs.size() &&
335  "expected same number of input args");
336  assert(newRegionOutArgs.size() == oldRegionOutArgs.size() &&
337  "expected same number of output args");
338 
339  for (auto it : llvm::zip(oldRegionInOutArgs, newRegionInOutArgs)) {
340  Value oldArg = std::get<0>(it);
341  Value newArg = std::get<1>(it);
342  rewriter.setInsertionPointToStart(newTiledLoopOp.getBody());
343  if (oldArg.getType().isa<TensorType>()) {
344  newBlockArgs.push_back(rewriter.create<bufferization::ToTensorOp>(
345  oldArg.getLoc(), newArg));
346  } else {
347  newBlockArgs.push_back(newArg);
348  }
349  }
350 
351  // Move old body into new loop.
352  rewriter.mergeBlocks(tiledLoopOp.getBody(), newTiledLoopOp.getBody(),
353  newBlockArgs);
354 
355  // Replace previous terminator with a new one that does not yield anything.
356  auto oldTerminator =
357  cast<linalg::YieldOp>(newTiledLoopOp.getBody()->getTerminator());
358  rewriter.setInsertionPointToEnd(newTiledLoopOp.getBody());
359  auto newTerminator =
360  rewriter.create<linalg::YieldOp>(oldTerminator->getLoc());
361 
362  // Copy buffer of yielded tensor to output buffer. If everything bufferized
363  // inplace, this copy will fold away.
364  rewriter.setInsertionPoint(newTerminator);
365  for (auto it : llvm::zip(oldTerminator.values(), newOutputs)) {
366  Value output = std::get<1>(it);
367  Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
368  newTerminator.getLoc(), output.getType(), std::get<0>(it));
369  if (failed(createMemCpy(rewriter, newTerminator.getLoc(), toMemrefOp,
370  output, state.getOptions())))
371  return failure();
372  }
373 
374  // Erase old terminator.
375  rewriter.eraseOp(oldTerminator);
376 
377  // Replace results and delete old op.
378  replaceOpWithBufferizedValues(rewriter, op, newResults);
379 
380  return success();
381  }
382 };
383 
384 /// Bufferization of linalg.yield. Bufferized as part of linalg.tiled_loop's
385 /// bufferization.
386 struct YieldOpInterface
387  : public BufferizableOpInterface::ExternalModel<YieldOpInterface,
388  linalg::YieldOp> {
389  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
390  const BufferizationState &state) const {
391  return true;
392  }
393 
394  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
395  const BufferizationState &state) const {
396  return false;
397  }
398 
399  OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
400  const BufferizationState &state) const {
401  return OpResult();
402  }
403 
404  bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand,
405  const BufferizationState &state) const {
406  // Yield operands always bufferize inplace. Otherwise, an alloc + copy
407  // may be generated inside the block. We should not return/yield allocations
408  // when possible.
409  return true;
410  }
411 
412  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
413  const BufferizationState &state) const {
414  auto yieldOp = cast<linalg::YieldOp>(op);
415 
416  if (!yieldOp->getParentOfType<TiledLoopOp>())
417  return yieldOp->emitError(
418  "expected that linalg.yield terminates a tiled_loop");
419 
420  assert(yieldOp->getOpOperands().empty() &&
421  "expected that linalg.yield was bufferized together with"
422  " tiled_loop");
423  return success();
424  }
425 };
426 
427 /// Helper structure that iterates over all LinalgOps in `OpTys` and registers
428 /// the `BufferizableOpInterface` with each of them.
429 template <typename... OpTys>
430 struct LinalgOpInterfaceHelper;
431 
432 template <typename First, typename... Others>
433 struct LinalgOpInterfaceHelper<First, Others...> {
434  static void registerOpInterface(DialectRegistry &registry) {
435  registry.addOpInterface<First, LinalgOpInterface<First>>();
436  LinalgOpInterfaceHelper<Others...>::registerOpInterface(registry);
437  }
438 };
439 
440 template <>
441 struct LinalgOpInterfaceHelper<> {
442  static void registerOpInterface(DialectRegistry &registry) {}
443 };
444 
445 } // namespace
446 
447 /// Try to eliminate InitTensorOps inside `op`. An InitTensorOp is replaced
448 /// with the the result of `rewriteFunc` if it is anchored on a matching
449 /// OpOperand. "Anchored" means that there is a path on the reverse SSA use-def
450 /// chain, starting from the OpOperand and always following the aliasing
451 /// OpOperand, that eventually ends at a single InitTensorOp.
455  BufferizationAliasInfo &aliasInfo,
456  AnchorMatchFn anchorMatchFunc, RewriteFn rewriteFunc,
457  SmallVector<Operation *> &newOps) {
458  OpBuilder b(op->getContext());
459 
460  WalkResult status = op->walk([&](Operation *op) {
461  for (OpOperand &operand : op->getOpOperands()) {
462  // Skip operands that do not bufferize inplace.
463  if (!aliasInfo.isInPlace(operand))
464  continue;
465  // Is this a matching OpOperand?
466  if (!anchorMatchFunc(operand))
467  continue;
468  SetVector<Value> maybeInitTensor =
469  state.findValueInReverseUseDefChain(operand.get(), [&](Value val) {
470  // Continue traversal until this function returns true.
471  OpResult opResult = val.dyn_cast<OpResult>();
472  if (!opResult)
473  return true;
474  SmallVector<OpOperand *> opOperands =
475  state.getAliasingOpOperand(opResult);
476  if (!llvm::all_of(opOperands, [&](OpOperand *operand) {
477  return aliasInfo.isInPlace(*operand);
478  }))
479  return true;
480  // Only equivalent tensors are supported at the moment.
481  // TODO: Support cases such as extract_slice(init_tensor)
482  return !llvm::all_of(opOperands, [&](OpOperand *operand) {
483  return aliasInfo.areEquivalentBufferizedValues(operand->get(),
484  opResult);
485  });
486  });
487 
488  // Replace only if the reverse use-def chain ends at exactly one
489  // InitTensorOp.
490  if (maybeInitTensor.size() != 1 ||
491  !maybeInitTensor.front().getDefiningOp<InitTensorOp>())
492  return WalkResult::skip();
493  Value initTensor = maybeInitTensor.front();
494 
495  // Create a replacement for the InitTensorOp.
496  b.setInsertionPoint(initTensor.getDefiningOp());
497  Value replacement = rewriteFunc(b, initTensor.getLoc(), operand);
498  if (!replacement)
499  continue;
500 
501  // Uses of the InitTensorOp are replaced here, but the op is not deleted.
502  // InitTensorOps without uses are ignored by the bufferization.
503  initTensor.replaceAllUsesWith(replacement);
504  aliasInfo.createAliasInfoEntry(replacement);
505  aliasInfo.unionAliasSets(initTensor, replacement);
506  aliasInfo.unionEquivalenceClasses(initTensor, replacement);
507 
508  // Register replacement ops.
509  if (Operation *newOp = replacement.getDefiningOp())
510  newOps.push_back(newOp);
511  }
512 
513  // Advance to the next operation.
514  return WalkResult::advance();
515  });
516 
517  return failure(status.wasInterrupted());
518 }
519 
520 /// Try to eliminate InitTensorOps inside `op`. An InitTensorOp can be
521 /// eliminated if it is eventually inserted into another tensor (and some other
522 /// conditions are met).
523 ///
524 /// E.g.:
525 /// %0 = linalg.init_tensor
526 /// %1 = linalg.fill(%cst, %0) {inplace = [true]}
527 /// %2 = tensor.insert_slice %1 into %t[10][20][1]
528 ///
529 /// InitTensorOp elimination will try to fill %t inplace instead of filling a
530 /// new allocation %0 and inserting it into %t. This is done by replacing the
531 /// InitTensorOp with:
532 ///
533 /// %0 = tensor.extract_slice %t[10][20][1]
534 ///
535 /// The analysis looks for matching ExtractSliceOp/InsertSliceOp pairs and lets
536 /// those bufferize inplace in the absence of other conflicts.
537 ///
538 /// Starting from an InsertSliceOp, an InitTensorOp at the end of the insert
539 /// source's reverse use-def chain is eliminated if:
540 /// * The InsertSliceOp was decided to bufferize inplace.
541 /// * On the reverse use-def chain path from the InsertSliceOp to the
542 /// InitTensorOp, all ops were decided to bufferize inplace and the buffer
543 /// relation is "equivalent" (TODO: can be relaxed if needed).
544 /// * The reverse use-def chain has exactly one end, which is the InitTensorOp.
545 ///
546 /// Note that the newly inserted ExtractSliceOp may have to bufferize
547 /// out-of-place due to RaW conflicts.
550  Operation *op, BufferizationState &state,
551  BufferizationAliasInfo &aliasInfo, SmallVector<Operation *> &newOps) {
552  return eliminateInitTensors(
553  op, state, aliasInfo,
554  /*anchorMatchFunc=*/
555  [&](OpOperand &operand) {
556  auto insertSliceOp =
557  dyn_cast<tensor::InsertSliceOp>(operand.getOwner());
558  if (!insertSliceOp)
559  return false;
560  // Only inplace bufferized InsertSliceOps are eligible.
561  if (!aliasInfo.isInPlace(insertSliceOp->getOpOperand(1) /*dest*/))
562  return false;
563  return &operand == &insertSliceOp->getOpOperand(0) /*source*/;
564  },
565  /*rewriteFunc=*/
566  [](OpBuilder &b, Location loc, OpOperand &operand) {
567  auto insertOp = cast<tensor::InsertSliceOp>(operand.getOwner());
568  // Expand offsets, sizes and strides to the full rank to handle the
569  // rank-reducing case.
570  SmallVector<OpFoldResult> mixedOffsets = insertOp.getMixedOffsets();
571  SmallVector<OpFoldResult> mixedSizes = insertOp.getMixedSizes();
572  SmallVector<OpFoldResult> mixedStrides = insertOp.getMixedStrides();
573  OffsetSizeAndStrideOpInterface::expandToRank(
574  insertOp.dest(), mixedOffsets, mixedSizes, mixedStrides,
575  [&](Value target, int64_t dim) -> OpFoldResult {
576  auto shapedType = target.getType().cast<ShapedType>();
577  if (shapedType.isDynamicDim(dim))
578  return b.create<tensor::DimOp>(loc, target, dim).result();
579  return b.getIndexAttr(shapedType.getDimSize(dim));
580  });
581  auto t = tensor::ExtractSliceOp::inferRankReducedResultType(
582  insertOp.getSourceType().getRank(),
583  insertOp.dest().getType().cast<RankedTensorType>(), mixedOffsets,
584  mixedSizes, mixedStrides);
585  auto extractOp = b.create<tensor::ExtractSliceOp>(
586  loc, t, insertOp.dest(), mixedOffsets, mixedSizes, mixedStrides);
587  return extractOp.result();
588  },
589  newOps);
590 }
591 
594  registry.addOpInterface<linalg::InitTensorOp, InitTensorOpInterface>();
595  registry.addOpInterface<linalg::TiledLoopOp, TiledLoopOpInterface>();
596  registry.addOpInterface<linalg::YieldOp, YieldOpInterface>();
597 
598  // Register all Linalg structured ops. `LinalgOp` is an interface and it is
599  // not possible to attach an external interface to an existing interface.
600  // Therefore, attach the `BufferizableOpInterface` to all ops one-by-one.
601  LinalgOpInterfaceHelper<
602 #define GET_OP_LIST
603 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
604  >::registerOpInterface(registry);
605 }
Include the generated interface declarations.
OpTy create(Location location, Args &&...args)
Create an operation of specific op type at the current insertion point.
Definition: Builders.h:430
Operation is a basic unit of execution within MLIR.
Definition: Operation.h:28
This is a value defined by a result of an operation.
Definition: Value.h:423
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:329
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
This class represents a single result from folding an operation.
Definition: OpDefinition.h:214
LogicalResult eliminateInitTensors(Operation *op, bufferization::BufferizationState &state, bufferization::BufferizationAliasInfo &aliasInfo, AnchorMatchFn anchorMatchFunc, RewriteFn rewriteFunc, SmallVector< Operation *> &newOps)
Try to eliminate InitTensorOps inside op.
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value...
Definition: LogicalResult.h:72
void replaceOpWithBufferizedValues(RewriterBase &rewriter, Operation *op, ValueRange values)
Replace an op with replacement values.
bool areEquivalentBufferizedValues(Value v1, Value v2) const
Return true if v1 and v2 bufferize to equivalent buffers.
bool isValueRead(Value value) const
Return true if the given value is read by an op that bufferizes to a memory read. ...
void unionEquivalenceClasses(Value v1, Value v2)
Union the equivalence classes of v1 and v2.
virtual void inlineRegionBefore(Region &region, Region &parent, Region::iterator before)
Move the blocks that belong to "region" before the given position in another region "parent"...
The BufferizationAliasInfo class maintains a list of buffer aliases and equivalence classes to suppor...
void replaceAllUsesWith(Value newValue) const
Replace all uses of &#39;this&#39; value with the new value, updating anything in the IR that uses &#39;this&#39; to ...
Definition: Value.h:161
bool isInPlace(OpOperand &opOperand) const
Return true if a value was marked as in-place bufferized.
static constexpr const bool value
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:48
MLIRContext * getContext()
Return the context this operation is associated with.
Definition: Operation.h:99
void createAliasInfoEntry(Value v)
Add a new entry for v in the aliasInfo and equivalentInfo.
MutableArrayRef< OpOperand > getOpOperands()
Definition: Operation.h:252
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
std::enable_if< llvm::function_traits< std::decay_t< FnT > >::num_args==1, RetT >::type walk(FnT &&callback)
Walk the operation by calling the callback for each nested operation (including this one)...
Definition: Operation.h:515
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
void addOpInterface()
Add an external op interface model for an op that belongs to a dialect, both provided as template par...
Definition: Dialect.h:367
This class provides support for representing a failure result, or a valid value of type T...
Definition: LogicalResult.h:77
std::function< bool(OpOperand &)> AnchorMatchFn
A function that matches anchor OpOperands for InitTensorOp elimination.
bool createDeallocs
Specifies whether dealloc ops should be generated along with alloc ops.
U dyn_cast() const
Definition: Value.h:99
void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry)
SmallVector< OpOperand * > getAliasingOpOperand(OpResult result) const
Determine which OpOperand* will alias with result if the op is bufferized in place.
FailureOr< Value > getBuffer(RewriterBase &rewriter, OpOperand &opOperand, bool forceInPlace=false, Optional< Operation *> customCopyInsertionPoint=None) const
Return the buffer (memref) for a given OpOperand (tensor).
LogicalResult run(Operation *op, bufferization::BufferizationState &state, bufferization::BufferizationAliasInfo &aliasInfo, SmallVector< Operation *> &newOps) override
Try to eliminate InitTensorOps inside op.
This class provides an abstraction over the various different ranges of value types.
Definition: TypeRange.h:38
static WalkResult advance()
Definition: Visitors.h:51
IRValueT get() const
Return the current value being used by this operand.
Definition: UseDefLists.h:133
A multi-dimensional affine map Affine map&#39;s are immutable like Type&#39;s, and they are uniqued...
Definition: AffineMap.h:38
A utility result that is used to signal how to proceed with an ongoing walk:
Definition: Visitors.h:34
Tensor types represent multi-dimensional arrays, and have two variants: RankedTensorType and Unranked...
Definition: BuiltinTypes.h:73
Location getLoc() const
Return the location of this value.
Definition: Value.cpp:26
LogicalResult createMemCpy(OpBuilder &b, Location loc, Value from, Value to, const BufferizationOptions &options)
Creates a memcpy between two given buffers.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:84
static WalkResult skip()
Definition: Visitors.h:52
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Definition: Builders.h:362
BufferizationState provides a variety of helper functions for dealing with tensor values and memref b...
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:279
Type getType() const
Return the type of this value.
Definition: Value.h:117
std::function< Value(OpBuilder &, Location, OpOperand &)> RewriteFn
A function that rewrites matched anchors.
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
Definition: Dialect.h:282
Operation * getOwner() const
Return the owner of this operand.
Definition: UseDefLists.h:37
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:20
void unionAliasSets(Value v1, Value v2)
Union the alias sets of v1 and v2.
This class represents an operand of an operation.
Definition: Value.h:249
const BufferizationOptions & getOptions() const
Return a reference to the BufferizationOptions.
BufferRelation
Specify fine-grain relationship between buffers to enable more analysis.
SetVector< Value > findValueInReverseUseDefChain(Value value, llvm::function_ref< bool(Value)> condition) const
Starting from value, follow the use-def chain in reverse, always selecting the aliasing OpOperands...
void setInsertionPointToEnd(Block *block)
Sets the insertion point to the end of the specified block.
Definition: Builders.h:367
bool isa() const
Definition: Types.h:234
virtual void mergeBlocks(Block *source, Block *dest, ValueRange argValues=llvm::None)
Merge the operations of block &#39;source&#39; into the end of block &#39;dest&#39;.
This class helps build Operations.
Definition: Builders.h:177
This class provides an abstraction over the different types of ranges over Values.
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:95
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:688
FailureOr< Value > createAlloc(OpBuilder &b, Location loc, MemRefType type, ArrayRef< Value > dynShape, const BufferizationOptions &options)
Creates a memref allocation.