MLIR  21.0.0git
ConvertToDestinationStyle.cpp
Go to the documentation of this file.
1 //===- ConvertToDestinationStyle.cpp - Convert non-DPS to DPS ops ---------===//
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 file contains patterns to convert non-DPS ops to DPS ops. New
10 // tensor.empty ops are inserted as a destination. Such tensor.empty can be
11 // eliminated with "empty tensor elimination", allowing them to bufferize
12 // without an allocation (assuming there are no further conflicts).
13 //
14 //===----------------------------------------------------------------------===//
15 //
24 #include "mlir/IR/Matchers.h"
25 #include "mlir/IR/PatternMatch.h"
26 #include "llvm/ADT/STLExtras.h"
27 #include "llvm/Support/Debug.h"
28 
29 using namespace mlir;
30 using namespace mlir::tensor;
31 
32 // Implements backtracking to traverse indices of the output buffer while
33 // iterating over op.elements().
34 static Value createInserts(RewriterBase &rewriter, Location loc, int dim,
35  Value destination, ArrayRef<int64_t> shape,
36  ArrayRef<Value> constants,
37  OperandRange::iterator &elementIt,
38  SmallVectorImpl<Value> &indices) {
39  if (dim == static_cast<int>(shape.size()) - 1) {
40  for (int i = 0; i < shape.back(); ++i) {
41  indices.back() = constants[i];
42  destination = rewriter.create<tensor::InsertOp>(loc, *elementIt,
43  destination, indices);
44  ++elementIt;
45  }
46  return destination;
47  }
48  for (int i = 0; i < shape[dim]; ++i) {
49  indices[dim] = constants[i];
50  destination = createInserts(rewriter, loc, dim + 1, destination, shape,
51  constants, elementIt, indices);
52  }
53  return destination;
54 }
55 
56 /// Create a memcpy from the given source tensor to the given destination
57 /// memref. The copy op type can be specified in the `options`.
58 static void createMemcpy(OpBuilder &b, Location loc, Value tensorSource,
59  Value memrefDest,
61  auto tensorType = dyn_cast<RankedTensorType>(tensorSource.getType());
62  assert(tensorType && "expected ranked tensor");
63  assert(isa<MemRefType>(memrefDest.getType()) && "expected ranked memref");
64 
65  switch (options.memcpyOp) {
68  // Note: This is the preferred way of memcpy'ing because no layout map
69  // and/or memory space must be specified for the source.
70  auto materializeOp = b.create<bufferization::MaterializeInDestinationOp>(
71  loc, tensorSource, memrefDest);
72  materializeOp.setWritable(true);
73  } break;
75  // TODO: Support custom memory space on source.
76  // We do not know the layout map of the source yet, so use a fully dynamic
77  // layout for best compatibility.
78  Value toBuffer = b.create<bufferization::ToBufferOp>(
80  tensorSource, /*readOnly=*/true);
81  b.create<memref::CopyOp>(loc, toBuffer, memrefDest);
82  } break;
84  // TODO: Support custom memory space on source.
85  // We do not know the layout map of the source yet, so use a fully dynamic
86  // layout for best compatibility.
87  Value toBuffer = b.create<bufferization::ToBufferOp>(
89  tensorSource, /*readOnly=*/true);
90  b.create<linalg::CopyOp>(loc, toBuffer, memrefDest);
91  } break;
92  };
93 }
94 
96  Location loc, PadOp padOp,
97  Value dest) {
98  OpBuilder::InsertionGuard g(rewriter);
99  RankedTensorType resultType = padOp.getResultType();
100 
101  // Examine the yielded value to decide if a linalg.generic is neede or a
102  // linalg.fill is sufficient.
103  Value yieldedValue =
104  cast<tensor::YieldOp>(padOp.getBody()->getTerminator()).getValue();
105  Attribute constYieldedValue;
106  // Is the yielded value a bbArg defined outside of the PadOp?
107  bool outsideBbArg =
108  isa<BlockArgument>(yieldedValue) &&
109  cast<BlockArgument>(yieldedValue).getOwner()->getParentOp() !=
110  padOp.getOperation();
111  // Is the yielded value an OpResult defined outside of the PadOp?
112  bool outsideOpResult =
113  isa<OpResult>(yieldedValue) &&
114  yieldedValue.getDefiningOp()->getParentOp() != padOp.getOperation();
115  bool invariantYieldedValue = outsideBbArg || outsideOpResult;
116  if (matchPattern(yieldedValue, m_Constant(&constYieldedValue))) {
117  // Padding with a constant: Create linalg.fill.
118  Dialect *arithDialect =
119  rewriter.getContext()->getLoadedDialect<arith::ArithDialect>();
120  Value fillValue =
121  arithDialect
122  ->materializeConstant(rewriter, constYieldedValue,
123  yieldedValue.getType(), yieldedValue.getLoc())
124  ->getResult(0);
125  auto fillOp = rewriter.create<linalg::FillOp>(loc, ValueRange(fillValue),
126  ValueRange(dest));
127  return fillOp;
128  }
129 
130  if (invariantYieldedValue) {
131  // Padding with an invariant value.
132  auto fillOp = rewriter.create<linalg::FillOp>(loc, ValueRange(yieldedValue),
133  ValueRange(dest));
134  return fillOp;
135  }
136 
137  // Create linalg.generic.
138  SmallVector<utils::IteratorType> iteratorTypes(resultType.getRank(),
139  utils::IteratorType::parallel);
140  SmallVector<AffineMap> indexingMaps(
141  1, rewriter.getMultiDimIdentityMap(resultType.getRank()));
142  auto genericOp = rewriter.create<linalg::GenericOp>(
143  loc, resultType, /*inputs=*/ValueRange(),
144  /*outputs=*/ValueRange{dest}, /*indexingMaps=*/
145  indexingMaps, iteratorTypes);
146  Block *body = rewriter.createBlock(&genericOp->getRegion(0), {},
147  resultType.getElementType(), loc);
148  rewriter.setInsertionPointToStart(body);
149  SmallVector<Value> bbArgReplacements;
150  for (int64_t i = 0; i < resultType.getRank(); ++i)
151  bbArgReplacements.push_back(rewriter.create<linalg::IndexOp>(loc, i));
152  rewriter.mergeBlocks(padOp.getBody(), body, bbArgReplacements);
153 
154  // Update terminator.
155  auto yieldOp = cast<tensor::YieldOp>(body->getTerminator());
156  rewriter.replaceOpWithNewOp<linalg::YieldOp>(yieldOp, yieldOp.getValue());
157  return genericOp;
158 }
159 
161  Value value) {
162  auto tensorType = cast<RankedTensorType>(value.getType());
163  if (tensorType.hasStaticShape())
164  return {};
165 
166  // Try to reify dynamic sizes.
167  ReifiedRankedShapedTypeDims reifiedShape;
168  if (isa<OpResult>(value) &&
169  succeeded(reifyResultShapes(b, value.getDefiningOp(), reifiedShape))) {
170  SmallVector<Value> dynSizes;
171  for (int64_t i = 0; i < tensorType.getRank(); ++i) {
172  if (tensorType.isDynamicDim(i))
173  dynSizes.push_back(cast<Value>(
174  reifiedShape[cast<OpResult>(value).getResultNumber()][i]));
175  }
176  return dynSizes;
177  }
178 
179  // Create tensor.dim ops.
180  SmallVector<Value> dynSizes;
181  for (int64_t i = 0; i < tensorType.getRank(); ++i) {
182  if (tensorType.isDynamicDim(i))
183  dynSizes.push_back(
184  b.create<DimOp>(value.getLoc(), value,
185  b.create<arith::ConstantIndexOp>(value.getLoc(), i)));
186  }
187  return dynSizes;
188 }
189 
190 static Value
193  Attribute memorySpace = {}) {
194  OpBuilder::InsertionGuard g(rewriter);
195  auto tensorType = cast<RankedTensorType>(value.getType());
196 
197  // Create buffer allocation.
198  auto memrefType =
200  tensorType, memorySpace));
201  SmallVector<Value> dynamicSizes = reifyOrComputeDynamicSizes(rewriter, value);
202 
203  Value alloc;
204  if (options.allocOp ==
206  alloc = rewriter.create<memref::AllocOp>(loc, memrefType, dynamicSizes);
207  if (options.emitDealloc) {
208  // Place deallocation at the end of the block.
209  rewriter.setInsertionPoint(rewriter.getInsertionBlock()->getTerminator());
210  rewriter.create<memref::DeallocOp>(loc, alloc);
211  }
212  } else if (options.allocOp ==
214  alloc = rewriter.create<memref::AllocaOp>(loc, memrefType, dynamicSizes);
215  // No dealloc is needed.
216  }
217 
218  return alloc;
219 }
220 
223  PadOp padOp, Attribute memorySpace, Operation *insertionPoint) {
224  // tensor.pad does not have a destination operand.
225  assert(!options.bufferizeDestinationOnly && "invalid options");
226 
227  OpBuilder::InsertionGuard g(rewriter);
228  rewriter.setInsertionPoint(insertionPoint ? insertionPoint : padOp);
229  Location loc = padOp.getLoc();
230 
231  // Create buffer allocation.
232  Value alloc = createAllocationForTensor(rewriter, loc, padOp.getResult(),
233  options, memorySpace);
234  rewriter.setInsertionPoint(padOp);
235 
236  if (!padOp.hasZeroLowPad() || !padOp.hasZeroHighPad()) {
237  // Create linalg.fill or linalg.generic. Not needed if there is no padding.
238  Operation *fillOp =
239  movePaddingToFillOrGenericOp(rewriter, loc, padOp, alloc);
240  rewriter.setInsertionPointAfter(fillOp);
241  }
242 
243  // Create memcpy.
245  getMixedSizes(rewriter, loc, padOp.getSource());
246  SmallVector<OpFoldResult> strides(padOp.getResultType().getRank(),
247  rewriter.getIndexAttr(1));
248  Value subview = rewriter.create<memref::SubViewOp>(
249  loc, alloc, /*offsets=*/padOp.getMixedLowPad(), sizes, strides);
250  createMemcpy(rewriter, loc, padOp.getSource(), subview, options);
251 
252  // Create bufferization.to_tensor with "restrict" and "writable". The returned
253  // tensor is a new buffer allocation, so it does not alias with any buffer.
254  Value toTensorOp = rewriter.create<bufferization::ToTensorOp>(
255  loc, padOp.getResult().getType(), alloc, /*restrict=*/true,
256  /*writable=*/true);
257  rewriter.replaceOp(padOp, toTensorOp);
258  return alloc;
259 }
260 
263  vector::MaskOp maskOp, Attribute memorySpace, Operation *insertionPoint) {
264  assert(llvm::range_size(maskOp.getMaskBlock()->without_terminator()) == 1 &&
265  "expected single masked op");
266  OpBuilder::InsertionGuard g(rewriter);
267 
268  // Should the bufferization options and state be function arguments?
269  bufferization::BufferizationOptions bufferizationOptions;
270  bufferization::BufferizationState bufferizationState;
271 
272  Operation *yieldOp = maskOp.getMaskRegion().front().getTerminator();
273  assert(isa<vector::YieldOp>(yieldOp) && "expected yield op terminator");
274 
275  // Bufferize maskable op. By default, place the buffer allocation right before
276  // the mask op.
278  rewriter, options, maskOp.getMaskableOp(), memorySpace,
279  /*insertionPoint=*/insertionPoint ? insertionPoint : maskOp);
280 
281  if (options.bufferizeDestinationOnly)
282  return alloc;
283 
284  // Bufferize terminator.
285  rewriter.setInsertionPoint(yieldOp);
286  if (failed(cast<bufferization::BufferizableOpInterface>(yieldOp).bufferize(
287  rewriter, bufferizationOptions, bufferizationState)))
288  return nullptr;
289 
290  // Erase dead to_tensor ops inside of the mask op. This is necessary because
291  // there only be one op (apart from the terminator) inside the mask op.
292  // TODO: Remove dead to_tensor ops more aggressively during bufferization.
293  SmallVector<Operation *> toTensorOps;
294  maskOp.walk([&](bufferization::ToTensorOp toTensorOp) {
295  if (toTensorOp->getUses().empty())
296  toTensorOps.push_back(toTensorOp.getOperation());
297  });
298  for (Operation *op : toTensorOps)
299  rewriter.eraseOp(op);
300 
301  // Bufferize mask op.
302  SmallVector<OpOperand *> resultUses;
303  for (Value result : maskOp.getResults())
304  if (isa<TensorType>(result.getType()))
305  for (OpOperand &use : result.getUses())
306  resultUses.push_back(&use);
307  rewriter.setInsertionPoint(maskOp);
308  if (failed(
309  cast<bufferization::BufferizableOpInterface>(maskOp.getOperation())
310  .bufferize(rewriter, bufferizationOptions, bufferizationState)))
311  return nullptr;
312 
313  // Set "restrict" attribute, indicating that no other tensor aliases with
314  // this tensor. That is because we just allocated a new buffer for the tensor.
315  for (OpOperand *resultUse : resultUses) {
316  auto toTensorOp =
317  resultUse->get().getDefiningOp<bufferization::ToTensorOp>();
318  assert(toTensorOp && "expected to_tensor op");
319  rewriter.modifyOpInPlace(toTensorOp, [&]() {
320  toTensorOp.setRestrict(true);
321  toTensorOp.setWritable(true);
322  });
323  }
324 
325  return alloc;
326 }
327 
330  bufferization::AllocTensorOp allocTensorOp, Attribute memorySpace,
331  Operation *insertionPoint) {
332  Location loc = allocTensorOp.getLoc();
333  OpBuilder::InsertionGuard g(rewriter);
334  rewriter.setInsertionPoint(insertionPoint ? insertionPoint : allocTensorOp);
335  bufferization::BufferizationOptions bufferizationOptions;
336 
337  // Create buffer allocation.
339  rewriter, loc, allocTensorOp.getResult(), options, memorySpace);
340 
341  // Create bufferization.to_tensor with "restrict" and "writable". The returned
342  // tensor is a new buffer allocation, so it does not alias with any buffer.
343  Value toTensorOp = rewriter.create<bufferization::ToTensorOp>(
344  loc, allocTensorOp.getResult().getType(), alloc, /*restrict=*/true,
345  /*writable=*/true);
346  rewriter.replaceOp(allocTensorOp, toTensorOp);
347  return alloc;
348 }
349 
350 /// Lower tensor.from_elements to a sequence of chained tensor.insert.
352  RewriterBase &rewriter, tensor::FromElementsOp fromElementsOp) {
353  Location loc = fromElementsOp.getLoc();
354  RankedTensorType tensorType =
355  cast<RankedTensorType>(fromElementsOp.getType());
356  auto shape = tensorType.getShape();
357 
358  // Create tensor.empty.
359  auto emptyOp = rewriter.create<EmptyOp>(loc, tensorType, ValueRange());
360 
361  // Case: tensor<elem_type>.
362  if (shape.empty()) {
363  Operation *res = rewriter.replaceOpWithNewOp<tensor::InsertOp>(
364  fromElementsOp, fromElementsOp.getElements().front(),
365  emptyOp.getResult(), ValueRange());
366  return res;
367  }
368 
369  // Create constants for the range of possible indices [0, max{shape_i}).
370  auto maxDim = *llvm::max_element(shape);
371  SmallVector<Value, 2> constants;
372  constants.reserve(maxDim);
373  for (int i = 0; i < maxDim; ++i)
374  constants.push_back(rewriter.create<arith::ConstantIndexOp>(loc, i));
375 
376  // Traverse all elements and create tensor.insert ops.
377  auto elementIt = fromElementsOp.getElements().begin();
378  SmallVector<Value, 2> indices(tensorType.getRank(), constants[0]);
379  Value result = createInserts(rewriter, loc, /*dim=*/0, emptyOp.getResult(),
380  shape, constants, elementIt, indices);
381 
382  // Replace tensor.from_elements.
383  rewriter.replaceOp(fromElementsOp, result);
384  return result.getDefiningOp();
385 }
386 
387 /// Lower tensor.generate to linalg.generic.
388 FailureOr<Operation *>
390  tensor::GenerateOp generateOp) {
391  // Only ops with exactly one block are supported.
392  if (!generateOp.getBody().hasOneBlock())
393  return failure();
394 
395  Location loc = generateOp.getLoc();
396  RankedTensorType tensorType = cast<RankedTensorType>(generateOp.getType());
397 
398  // Create tensor.empty.
399  auto emptyOp =
400  rewriter.create<EmptyOp>(loc, tensorType, generateOp.getDynamicExtents());
401 
402  // Create linalg.generic.
403  SmallVector<utils::IteratorType> iteratorTypes(tensorType.getRank(),
404  utils::IteratorType::parallel);
405  SmallVector<AffineMap> indexingMaps(
406  1, rewriter.getMultiDimIdentityMap(tensorType.getRank()));
407  auto genericOp = rewriter.create<linalg::GenericOp>(
408  loc, tensorType, /*inputs=*/ValueRange(),
409  /*outputs=*/ValueRange{emptyOp.getResult()}, /*indexingMaps=*/
410  indexingMaps, iteratorTypes);
411  Block *body = rewriter.createBlock(&genericOp->getRegion(0), {},
412  tensorType.getElementType(), loc);
413  rewriter.setInsertionPointToStart(body);
414  SmallVector<Value> bbArgReplacements;
415  for (int64_t i = 0; i < tensorType.getRank(); ++i)
416  bbArgReplacements.push_back(rewriter.create<linalg::IndexOp>(loc, i));
417  rewriter.mergeBlocks(&generateOp.getBody().front(), body, bbArgReplacements);
418 
419  // Update terminator.
420  auto yieldOp = cast<tensor::YieldOp>(body->getTerminator());
421  rewriter.replaceOpWithNewOp<linalg::YieldOp>(yieldOp, yieldOp.getValue());
422 
423  // Replace tensor.generate.
424  rewriter.replaceOp(generateOp, genericOp->getResult(0));
425  return genericOp.getOperation();
426 }
427 
428 /// Lower tensor.pad to linalg.generic + tensor.insert_slice.
429 FailureOr<Operation *>
431  tensor::PadOp padOp) {
432  // Only ops with exactly one block are supported.
433  if (!padOp.getBodyRegion().hasOneBlock())
434  return failure();
435 
436  // Create tensor.empty.
437  Location loc = padOp.getLoc();
438  RankedTensorType resultType = padOp.getResultType();
439  ReifiedRankedShapedTypeDims reifiedShape;
440  if (failed(reifyResultShapes(rewriter, padOp, reifiedShape)))
441  return rewriter.notifyMatchFailure(
442  padOp, "failed to reify tensor.pad op result shape");
443  SmallVector<Value> dynamicSizes;
444  for (int64_t i = 0; i < resultType.getRank(); ++i)
445  if (resultType.isDynamicDim(i))
446  dynamicSizes.push_back(cast<Value>(reifiedShape[0][i]));
447 
448  // If the `padOp` has a nofold attribute and all paddings are known to be 0,
449  // explicitly insert a `linalg.copy`.
450  if (padOp.getNofoldAttr() &&
451  llvm::all_of(padOp.getMixedLowPad(), isZeroInteger) &&
452  llvm::all_of(padOp.getMixedHighPad(), isZeroInteger)) {
453  using bufferization::AllocTensorOp;
454  Value allocated =
455  rewriter.create<AllocTensorOp>(loc, resultType, dynamicSizes);
456  auto copyOp = rewriter.replaceOpWithNewOp<linalg::CopyOp>(
457  padOp, padOp.getSource(), allocated);
458  return copyOp.getOperation();
459  }
460 
461  Value empty = rewriter.create<EmptyOp>(loc, resultType, dynamicSizes);
462  // Create linalg.fill or linalg.generic.
463  Operation *fillOp = movePaddingToFillOrGenericOp(rewriter, loc, padOp, empty);
464  rewriter.setInsertionPointAfter(fillOp);
465 
466  // Create tensor::InsertSliceOp.
467  SmallVector<OpFoldResult> sliceSizes =
468  getMixedSizes(rewriter, loc, padOp.getSource());
469  SmallVector<OpFoldResult> sliceStrides(resultType.getRank(),
470  rewriter.getIndexAttr(1));
471  auto insertSliceOp = rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
472  padOp, padOp.getSource(), fillOp->getResult(0),
473  /*offsets=*/padOp.getMixedLowPad(), sliceSizes, sliceStrides);
474  return insertSliceOp.getOperation();
475 }
476 
479  Operation *op, Attribute memorySpace, Operation *insertionPoint) {
480  using namespace bufferization;
481 
482  // Call specialized overload for certain ops.
483  if (auto padOp = dyn_cast<tensor::PadOp>(op))
484  return bufferizeToAllocation(rewriter, options, padOp, memorySpace);
485  if (auto maskOp = dyn_cast<vector::MaskOp>(op))
486  return bufferizeToAllocation(rewriter, options, maskOp, memorySpace);
487  if (auto allocTensorOp = dyn_cast<bufferization::AllocTensorOp>(op))
488  return bufferizeToAllocation(rewriter, options, allocTensorOp, memorySpace);
489 
490  // Only bufferizable ops are supported.
491  auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
492  if (!bufferizableOp)
493  return nullptr;
494 
495  // Should the bufferization options and states be function arguments?
496  BufferizationOptions bufferizationOptions;
497  AnalysisState analysisState(bufferizationOptions);
498  BufferizationState bufferizationState;
499 
500 #ifndef NDEBUG
501  if (!options.bufferizeDestinationOnly) {
502  // Ops with nested tensor ops are not supported yet. At the moment, this
503  // function just bufferizes the given op itself, but not its body.
504  op->walk([&](Operation *nestedOp) {
505  if (op == nestedOp)
506  return;
507  if (llvm::any_of(nestedOp->getOperands(),
508  [](Value v) { return isa<TensorType>(v.getType()); }))
509  llvm_unreachable("ops with nested tensor ops are not supported yet");
510  if (llvm::any_of(nestedOp->getResults(),
511  [](Value v) { return isa<TensorType>(v.getType()); }))
512  llvm_unreachable("ops with nested tensor ops are not supported yet");
513  });
514  }
515 #endif // NDEBUG
516 
517  // Gather tensor results.
518  SmallVector<OpResult> tensorResults;
519  for (OpResult result : op->getResults()) {
520  if (!isa<TensorType>(result.getType()))
521  continue;
522  // Unranked tensors are not supported
523  if (!isa<RankedTensorType>(result.getType()))
524  return nullptr;
525  // Ops that bufferize to an allocation are not supported.
526  if (bufferizableOp.bufferizesToAllocation(result))
527  return nullptr;
528  tensorResults.push_back(result);
529  }
530 
531  // Gather all operands that should bufferize to a new allocation. I.e.,
532  // bufferize out-of-place.
533  SmallVector<OpOperand *> outOfPlaceOperands, resultUses;
534  auto addOutOfPlaceOperand = [&](OpOperand *operand) {
535  if (!llvm::is_contained(outOfPlaceOperands, operand))
536  outOfPlaceOperands.push_back(operand);
537  };
538  for (OpResult result : tensorResults) {
539  AliasingOpOperandList aliasingOperands =
540  analysisState.getAliasingOpOperands(result);
541  for (const AliasingOpOperand &operand : aliasingOperands) {
542  addOutOfPlaceOperand(operand.opOperand);
543  for (OpOperand &resultUse : result.getUses())
544  resultUses.push_back(&resultUse);
545  }
546  }
547  for (OpOperand &operand : op->getOpOperands()) {
548  if (!analysisState.bufferizesToMemoryWrite(operand))
549  continue;
550  if (!isa<RankedTensorType>(operand.get().getType()))
551  continue;
552  addOutOfPlaceOperand(&operand);
553  }
554  // TODO: Support multiple buffers.
555  if (outOfPlaceOperands.size() != 1)
556  return nullptr;
557 
558  // Allocate buffers.
559  OpBuilder::InsertionGuard g(rewriter);
560  rewriter.setInsertionPoint(insertionPoint ? insertionPoint : op);
561  SmallVector<Value> allocs;
562  for (OpOperand *operand : outOfPlaceOperands) {
564  rewriter, op->getLoc(), operand->get(), options, memorySpace);
565  allocs.push_back(alloc);
566  if (!analysisState.findDefinitions(operand).empty()) {
567  // Initialize buffer with a copy of the operand data. Not needed if the
568  // tensor is uninitialized.
569  createMemcpy(rewriter, op->getLoc(), operand->get(), alloc, options);
570  }
571  rewriter.modifyOpInPlace(op, [&]() {
572  auto toTensorOp = rewriter.create<ToTensorOp>(
573  op->getLoc(), operand->get().getType(), alloc);
574  operand->set(toTensorOp);
575  if (options.bufferizeDestinationOnly) {
576  rewriter.modifyOpInPlace(toTensorOp, [&]() {
577  toTensorOp.setRestrict(true);
578  toTensorOp.setWritable(true);
579  });
580  }
581  });
582  }
583 
584  if (options.bufferizeDestinationOnly)
585  return allocs.front();
586 
587  // Bufferize the op.
588  rewriter.setInsertionPoint(op);
589  if (failed(bufferizableOp.bufferize(rewriter, bufferizationOptions,
590  bufferizationState)))
591  return nullptr;
592 
593  // Set "restrict" attribute, indicating that no other tensor aliases with
594  // this tensor. That is because we just allocated a new buffer for the tensor.
595  for (OpOperand *resultUse : resultUses) {
596  auto toTensorOp = resultUse->get().getDefiningOp<ToTensorOp>();
597  assert(toTensorOp && "expected to_tensor op");
598  rewriter.modifyOpInPlace(toTensorOp, [&]() {
599  toTensorOp.setRestrict(true);
600  toTensorOp.setWritable(true);
601  });
602  }
603  return allocs.front();
604 }
605 
606 namespace {
607 
608 template <typename OpTy>
609 LogicalResult rewriteOpInDestinationPassingStyle(OpTy op,
610  PatternRewriter &rewriter) {
611  return linalg::rewriteInDestinationPassingStyle(rewriter, op);
612 }
613 
614 } // namespace
615 
618  patterns.add(rewriteOpInDestinationPassingStyle<tensor::FromElementsOp>);
619  patterns.add(rewriteOpInDestinationPassingStyle<tensor::GenerateOp>);
620  patterns.add(rewriteOpInDestinationPassingStyle<tensor::PadOp>);
621 }
static Operation * movePaddingToFillOrGenericOp(RewriterBase &rewriter, Location loc, PadOp padOp, Value dest)
static Value createAllocationForTensor(RewriterBase &rewriter, Location loc, Value value, const linalg::BufferizeToAllocationOptions &options, Attribute memorySpace={})
static void createMemcpy(OpBuilder &b, Location loc, Value tensorSource, Value memrefDest, const linalg::BufferizeToAllocationOptions &options)
Create a memcpy from the given source tensor to the given destination memref.
static SmallVector< Value > reifyOrComputeDynamicSizes(OpBuilder &b, Value value)
static Value createInserts(RewriterBase &rewriter, Location loc, int dim, Value destination, ArrayRef< int64_t > shape, ArrayRef< Value > constants, OperandRange::iterator &elementIt, SmallVectorImpl< Value > &indices)
static llvm::ManagedStatic< PassManagerOptions > options
Base class for generic analysis states.
Attributes are known-constant values of operations.
Definition: Attributes.h:25
Block represents an ordered list of Operations.
Definition: Block.h:33
Operation * getTerminator()
Get the terminator operation of this block.
Definition: Block.cpp:246
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:106
AffineMap getMultiDimIdentityMap(unsigned rank)
Definition: Builders.cpp:385
MLIRContext * getContext() const
Definition: Builders.h:55
Dialects are groups of MLIR operations, types and attributes, as well as behavior associated with the...
Definition: Dialect.h:38
virtual Operation * materializeConstant(OpBuilder &builder, Attribute value, Type type, Location loc)
Registered hook to materialize a single constant operation from a given attribute value with the desi...
Definition: Dialect.h:83
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:76
Dialect * getLoadedDialect(StringRef name)
Get a registered IR dialect with the given namespace.
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:346
This class helps build Operations.
Definition: Builders.h:205
Block * createBlock(Region *parent, Region::iterator insertPt={}, TypeRange argTypes={}, ArrayRef< Location > locs={})
Add new block with 'argTypes' arguments and set the insertion point to the end of it.
Definition: Builders.cpp:428
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Definition: Builders.h:429
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:396
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:455
void setInsertionPointAfter(Operation *op)
Sets the insertion point to the node after the specified operation, which will cause subsequent inser...
Definition: Builders.h:410
Block * getInsertionBlock() const
Return the block the current insertion point belongs to.
Definition: Builders.h:440
This class represents an operand of an operation.
Definition: Value.h:257
This is a value defined by a result of an operation.
Definition: Value.h:447
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
std::enable_if_t< llvm::function_traits< std::decay_t< FnT > >::num_args==1, RetT > walk(FnT &&callback)
Walk the operation by calling the callback for each nested operation (including this one),...
Definition: Operation.h:797
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:223
Operation * getParentOp()
Returns the closest surrounding operation that contains this operation or nullptr if this is a top-le...
Definition: Operation.h:234
MutableArrayRef< OpOperand > getOpOperands()
Definition: Operation.h:383
operand_range getOperands()
Returns an iterator on the underlying Value's.
Definition: Operation.h:378
result_range getResults()
Definition: Operation.h:415
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:748
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:358
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
Definition: PatternMatch.h:681
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
void mergeBlocks(Block *source, Block *dest, ValueRange argValues={})
Inline the operations of block 'source' into the end of block 'dest'.
void modifyOpInPlace(Operation *root, CallableT &&callable)
This method is a utility wrapper around an in-place modification of an operation.
Definition: PatternMatch.h:593
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Definition: PatternMatch.h:500
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:387
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Type getType() const
Return the type of this value.
Definition: Value.h:105
Location getLoc() const
Return the location of this value.
Definition: Value.cpp:26
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:93
BufferizationState provides information about the state of the IR during the bufferization process.
BaseMemRefType getMemRefTypeWithStaticIdentityLayout(TensorType tensorType, Attribute memorySpace=nullptr)
Return a MemRef type with a static identity layout (i.e., no layout map).
AliasList< AliasingOpOperand > AliasingOpOperandList
A list of possible aliasing OpOperands.
BaseMemRefType getMemRefTypeWithFullyDynamicLayout(TensorType tensorType, Attribute memorySpace=nullptr)
Return a MemRef type with fully dynamic layout.
Value bufferizeToAllocation(RewriterBase &rewriter, const BufferizeToAllocationOptions &options, tensor::PadOp padOp, Attribute memorySpace={}, Operation *insertionPoint=nullptr)
Materialize a buffer allocation for the given tensor.pad op and lower the op to linalg....
FailureOr< Operation * > rewriteInDestinationPassingStyle(RewriterBase &rewriter, tensor::FromElementsOp fromElementsOp)
Rewrite tensor.from_elements to linalg.generic.
void populateConvertToDestinationStylePatterns(RewritePatternSet &patterns)
Populate patterns that convert non-destination-style ops to destination style ops.
SmallVector< OpFoldResult > getMixedSizes(OpBuilder &builder, Location loc, Value value)
Return the dimensions of the given tensor value.
Definition: TensorOps.cpp:73
Include the generated interface declarations.
bool matchPattern(Value value, const Pattern &pattern)
Entry point for matching a pattern over a Value.
Definition: Matchers.h:490
LogicalResult reifyResultShapes(OpBuilder &b, Operation *op, ReifiedRankedShapedTypeDims &reifiedReturnShapes)
Reify the shape of the result of an operation (typically in terms of the shape of its operands).
const FrozenRewritePatternSet & patterns
bool isZeroInteger(OpFoldResult v)
Return true if v is an IntegerAttr with value 0.
detail::constant_op_matcher m_Constant()
Matches a constant foldable operation.
Definition: Matchers.h:369
Options for BufferizableOpInterface-based bufferization.