MLIR  21.0.0git
BufferizableOpInterfaceImpl.cpp
Go to the documentation of this file.
1 //===- BufferizableOpInterfaceImpl.cpp - 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 
10 
23 #include "mlir/IR/Dialect.h"
24 #include "mlir/IR/Operation.h"
25 
26 using namespace mlir;
27 using namespace mlir::bufferization;
28 using namespace mlir::tensor;
29 
30 namespace mlir {
31 namespace tensor {
32 namespace {
33 
34 struct CastOpInterface
35  : public BufferizableOpInterface::ExternalModel<CastOpInterface,
36  tensor::CastOp> {
37  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
38  const AnalysisState &state) const {
39  return false;
40  }
41 
42  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
43  const AnalysisState &state) const {
44  return false;
45  }
46 
47  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
48  const AnalysisState &state) const {
49  return {{op->getResult(0), BufferRelation::Equivalent}};
50  }
51 
52  FailureOr<BaseMemRefType>
54  SmallVector<Value> &invocationStack) const {
55  auto castOp = cast<tensor::CastOp>(op);
56  auto maybeSrcBufferType = bufferization::getBufferType(
57  castOp.getSource(), options, invocationStack);
58  if (failed(maybeSrcBufferType))
59  return failure();
60  Attribute memorySpace = maybeSrcBufferType->getMemorySpace();
61 
62  // Note: `getMemRefTypeWithFullyDynamicLayout` returns an unranked memref
63  // type in case the input is an unranked tensor type.
64 
65  // Case 1: Casting an unranked tensor
66  if (isa<UnrankedTensorType>(castOp.getSource().getType())) {
67  // When casting to a ranked tensor, we cannot infer any static offset or
68  // strides from the source. Assume fully dynamic.
69  return getMemRefTypeWithFullyDynamicLayout(castOp.getType(), memorySpace);
70  }
71 
72  // Case 2: Casting to an unranked tensor type
73  if (isa<UnrankedTensorType>(castOp.getType())) {
74  return getMemRefTypeWithFullyDynamicLayout(castOp.getType(), memorySpace);
75  }
76 
77  // Case 3: Ranked tensor -> ranked tensor. The offsets and strides do not
78  // change.
79  auto rankedResultType = cast<RankedTensorType>(castOp.getType());
80  return MemRefType::get(
81  rankedResultType.getShape(), rankedResultType.getElementType(),
82  llvm::cast<MemRefType>(*maybeSrcBufferType).getLayout(), memorySpace);
83  }
84 
85  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
86  const BufferizationOptions &options) const {
87  auto castOp = cast<tensor::CastOp>(op);
88 
89  // The result buffer still has the old (pre-cast) type.
90  FailureOr<Value> resultBuffer =
91  getBuffer(rewriter, castOp.getSource(), options);
92  if (failed(resultBuffer))
93  return failure();
94 
95  // Compute the new type.
96  auto resultMemRefType =
97  bufferization::getBufferType(castOp.getResult(), options);
98  if (failed(resultMemRefType))
99  return failure();
100  if (resultBuffer->getType() == *resultMemRefType) {
101  // This cast is a no-op.
102  replaceOpWithBufferizedValues(rewriter, op, *resultBuffer);
103  return success();
104  }
105 
106  // Replace the op with a memref.cast.
107  assert(memref::CastOp::areCastCompatible(resultBuffer->getType(),
108  *resultMemRefType) &&
109  "CallOp::bufferize: cast incompatible");
110  replaceOpWithNewBufferizedOp<memref::CastOp>(
111  rewriter, op, *resultMemRefType, *resultBuffer);
112 
113  return success();
114  }
115 };
116 
117 /// Bufferization of tensor.collapse_shape. Replace with memref.collapse_shape.
118 struct CollapseShapeOpInterface
119  : public BufferizableOpInterface::ExternalModel<CollapseShapeOpInterface,
120  tensor::CollapseShapeOp> {
121  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
122  const AnalysisState &state) const {
123  // tensor.collapse_shape may reallocate, at which point the source buffer is
124  // copied. I.e., there will be a memory read side effect on the bufferized
125  // source. This function conservatively returns "true" because whether a
126  // copy will be created or not is not known at this point.
127  return true;
128  }
129 
130  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
131  const AnalysisState &state) const {
132  return false;
133  }
134 
135  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
136  const AnalysisState &state) const {
137  // TODO: CollapseShapeOp may allocate at runtime.
138  return {{op->getOpResult(0), BufferRelation::Equivalent}};
139  }
140 
141  FailureOr<BaseMemRefType>
143  SmallVector<Value> &invocationStack) const {
144  auto collapseShapeOp = cast<tensor::CollapseShapeOp>(op);
145  auto maybeSrcBufferType = bufferization::getBufferType(
146  collapseShapeOp.getSrc(), options, invocationStack);
147  if (failed(maybeSrcBufferType))
148  return failure();
149  auto srcBufferType = llvm::cast<MemRefType>(*maybeSrcBufferType);
150  bool canBeCollapsed = memref::CollapseShapeOp::isGuaranteedCollapsible(
151  srcBufferType, collapseShapeOp.getReassociationIndices());
152 
153  if (!canBeCollapsed) {
154  // If dims cannot be collapsed, this op bufferizes to a new allocation.
155  RankedTensorType tensorResultType = collapseShapeOp.getResultType();
157  tensorResultType, srcBufferType.getMemorySpace());
158  }
159 
160  return memref::CollapseShapeOp::computeCollapsedType(
161  srcBufferType, collapseShapeOp.getReassociationIndices());
162  }
163 
164  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
165  const BufferizationOptions &options) const {
166  auto collapseShapeOp = cast<tensor::CollapseShapeOp>(op);
167  RankedTensorType tensorResultType = collapseShapeOp.getResultType();
168  FailureOr<Value> maybeBuffer =
169  getBuffer(rewriter, collapseShapeOp.getSrc(), options);
170  if (failed(maybeBuffer))
171  return failure();
172  Value buffer = *maybeBuffer;
173  auto bufferType = cast<MemRefType>(buffer.getType());
174 
175  if (tensorResultType.getRank() == 0) {
176  // 0-d collapses must go through a different op builder.
177  MemRefType resultType;
178 
179  if (bufferType.getLayout().isIdentity()) {
180  // Standard layout: result type has no offset.
181  MemRefLayoutAttrInterface layout;
182  resultType = MemRefType::get({}, tensorResultType.getElementType(),
183  layout, bufferType.getMemorySpace());
184  } else {
185  // Source memref has a layout map: result type has the same offset as
186  // the source type.
187  SmallVector<int64_t> strides;
188  int64_t offset;
189  if (failed(bufferType.getStridesAndOffset(strides, offset)))
190  return failure();
191  resultType = MemRefType::get(
192  {}, tensorResultType.getElementType(),
193  StridedLayoutAttr::get(op->getContext(), offset, {}),
194  bufferType.getMemorySpace());
195  }
196 
197  replaceOpWithNewBufferizedOp<memref::CollapseShapeOp>(
198  rewriter, op, resultType, buffer, collapseShapeOp.getReassociation());
199  return success();
200  }
201 
202  // If the dims are not collapsible (due to an incompatible source layout
203  // map), force an out-of-place bufferization, i.e., a buffer copy. This
204  // newly allocated buffer will have no layout map and thus be collapsible.
205  bool canBeCollapsed = memref::CollapseShapeOp::isGuaranteedCollapsible(
206  bufferType, collapseShapeOp.getReassociationIndices());
207  if (!canBeCollapsed) {
208  // TODO: Create alloc_tensor ops during TensorCopyInsertion.
209  AnalysisState analysisState(options);
210  FailureOr<Value> tensorAlloc = allocateTensorForShapedValue(
211  rewriter, op->getLoc(), collapseShapeOp.getSrc(), options);
212  if (failed(tensorAlloc))
213  return failure();
214  auto memrefType =
215  MemRefType::get(collapseShapeOp.getSrcType().getShape(),
216  collapseShapeOp.getSrcType().getElementType(),
217  AffineMap(), bufferType.getMemorySpace());
218  buffer = rewriter.create<bufferization::ToMemrefOp>(
219  op->getLoc(), memrefType, *tensorAlloc);
220  }
221 
222  // Result type is inferred by the builder.
223  replaceOpWithNewBufferizedOp<memref::CollapseShapeOp>(
224  rewriter, op, buffer, collapseShapeOp.getReassociationIndices());
225  return success();
226  }
227 };
228 
229 /// Bufferization of tensor.dim. Replace with memref.dim.
230 struct DimOpInterface
231  : public BufferizableOpInterface::ExternalModel<DimOpInterface,
232  tensor::DimOp> {
233  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
234  const AnalysisState &state) const {
235  // The op reads the tensor's metadata but not its contents.
236  return false;
237  }
238 
239  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
240  const AnalysisState &state) const {
241  return false;
242  }
243 
244  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
245  const AnalysisState &state) const {
246  return {};
247  }
248 
249  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
250  const BufferizationOptions &options) const {
251  auto dimOp = cast<tensor::DimOp>(op);
252  FailureOr<Value> v = getBuffer(rewriter, dimOp.getSource(), options);
253  if (failed(v))
254  return failure();
255  replaceOpWithNewBufferizedOp<memref::DimOp>(rewriter, op, *v,
256  dimOp.getIndex());
257  return success();
258  }
259 };
260 
261 /// Bufferization of "tensor.empty". Replace with "bufferization.alloc_tensor".
262 struct EmptyOpInterface
263  : public BufferizableOpInterface::ExternalModel<EmptyOpInterface,
264  tensor::EmptyOp> {
265  bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
266 
267  bool resultBufferizesToMemoryWrite(Operation *op, OpResult opResult,
268  const AnalysisState &state) const {
269  // The returned tensor does not have specified contents.
270  return false;
271  }
272 
273  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
274  const BufferizationOptions &options) const {
275  auto emptyOp = cast<tensor::EmptyOp>(op);
276 
277  // Optimization: Fold away the op if it has no uses.
278  if (op->getUses().empty()) {
279  rewriter.eraseOp(op);
280  return success();
281  }
282 
283  // Allocate a tensor. This emits a "bufferization.alloc_tensor" op.
284  FailureOr<Value> allocTensor = allocateTensorForShapedValue(
285  rewriter, op->getLoc(), emptyOp.getResult(), options, /*copy=*/false);
286  if (failed(allocTensor))
287  return failure();
288  rewriter.replaceOp(op, *allocTensor);
289  return success();
290  }
291 };
292 
293 /// Bufferization of tensor.expand_shape. Replace with memref.expand_shape.
294 struct ExpandShapeOpInterface
295  : public BufferizableOpInterface::ExternalModel<ExpandShapeOpInterface,
296  tensor::ExpandShapeOp> {
297  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
298  const AnalysisState &state) const {
299  // In contrast to tensor.collapse_shape, this op can always be bufferized
300  // without a copy.
301  return false;
302  }
303 
304  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
305  const AnalysisState &state) const {
306  return false;
307  }
308 
309  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
310  const AnalysisState &state) const {
311  return {{op->getOpResult(0), BufferRelation::Equivalent}};
312  }
313 
314  FailureOr<BaseMemRefType>
316  SmallVector<Value> &invocationStack) const {
317  auto expandShapeOp = cast<tensor::ExpandShapeOp>(op);
318  auto maybeSrcBufferType = bufferization::getBufferType(
319  expandShapeOp.getSrc(), options, invocationStack);
320  if (failed(maybeSrcBufferType))
321  return failure();
322  auto srcBufferType = llvm::cast<MemRefType>(*maybeSrcBufferType);
323  auto maybeResultType = memref::ExpandShapeOp::computeExpandedType(
324  srcBufferType, expandShapeOp.getResultType().getShape(),
325  expandShapeOp.getReassociationIndices());
326  if (failed(maybeResultType))
327  return failure();
328  return *maybeResultType;
329  }
330 
331  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
332  const BufferizationOptions &options) const {
333  auto expandShapeOp = cast<tensor::ExpandShapeOp>(op);
334  auto tensorResultType = expandShapeOp.getResultType();
335  FailureOr<Value> buffer =
336  getBuffer(rewriter, expandShapeOp.getSrc(), options);
337  if (failed(buffer))
338  return failure();
339 
340  auto memrefExpandShape = rewriter.create<memref::ExpandShapeOp>(
341  op->getLoc(), tensorResultType.getShape(), *buffer,
342  expandShapeOp.getReassociationIndices(),
343  expandShapeOp.getMixedOutputShape());
344  replaceOpWithBufferizedValues(rewriter, op,
345  memrefExpandShape->getResults());
346  return success();
347  }
348 };
349 
350 /// Bufferization of tensor.extract_slice. Replace with memref.subview.
351 struct ExtractSliceOpInterface
352  : public BufferizableOpInterface::ExternalModel<ExtractSliceOpInterface,
353  tensor::ExtractSliceOp> {
354  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
355  const AnalysisState &state) const {
356  return false;
357  }
358 
359  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
360  const AnalysisState &state) const {
361  return false;
362  }
363 
364  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
365  const AnalysisState &state) const {
366  return {{op->getOpResult(0), BufferRelation::Unknown}};
367  }
368 
369  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
370  const BufferizationOptions &options) const {
371  auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
372  SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
373  SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
374  SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
375  Location loc = extractSliceOp.getLoc();
376 
377  // Get source buffer.
378  FailureOr<Value> srcMemref =
379  getBuffer(rewriter, extractSliceOp.getSource(), options);
380  if (failed(srcMemref))
381  return failure();
382 
383  // Take a subview of the source buffer.
384  auto resultMemrefType =
385  bufferization::getBufferType(extractSliceOp.getResult(), options);
386  if (failed(resultMemrefType))
387  return failure();
388  Value subView = rewriter.create<memref::SubViewOp>(
389  loc, llvm::cast<MemRefType>(*resultMemrefType), *srcMemref,
390  mixedOffsets, mixedSizes, mixedStrides);
391 
392  replaceOpWithBufferizedValues(rewriter, op, subView);
393  return success();
394  }
395 
396  FailureOr<BaseMemRefType>
398  SmallVector<Value> &invocationStack) const {
399  auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
400  assert(value == extractSliceOp.getResult() && "invalid value");
401  auto srcMemrefType = bufferization::getBufferType(
402  extractSliceOp.getSource(), options, invocationStack);
403  if (failed(srcMemrefType))
404  return failure();
405  SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
406  SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
407  SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
408  return memref::SubViewOp::inferRankReducedResultType(
409  extractSliceOp.getType().getShape(),
410  llvm::cast<MemRefType>(*srcMemrefType), mixedOffsets, mixedSizes,
411  mixedStrides);
412  }
413 };
414 
415 /// Bufferization of tensor.extract. Replace with memref.load.
416 struct ExtractOpInterface
417  : public BufferizableOpInterface::ExternalModel<ExtractOpInterface,
418  tensor::ExtractOp> {
419  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
420  const AnalysisState &state) const {
421  return true;
422  }
423 
424  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
425  const AnalysisState &state) const {
426  return false;
427  }
428 
429  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
430  const AnalysisState &state) const {
431  return {};
432  }
433 
434  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
435  const BufferizationOptions &options) const {
436  auto extractOp = cast<tensor::ExtractOp>(op);
437  FailureOr<Value> srcMemref =
438  getBuffer(rewriter, extractOp.getTensor(), options);
439  if (failed(srcMemref))
440  return failure();
441  replaceOpWithNewBufferizedOp<memref::LoadOp>(rewriter, op, *srcMemref,
442  extractOp.getIndices());
443  return success();
444  }
445 };
446 
447 // Implements backtracking to traverse indices of the output buffer while
448 // iterating over op.elements().
449 static void createStores(RewriterBase &rewriter, Location loc, int dim,
450  Value buffer, ArrayRef<int64_t> shape,
451  ArrayRef<Value> constants,
452  OperandRange::iterator &elementIt,
453  SmallVectorImpl<Value> &indices) {
454  if (dim == static_cast<int>(shape.size()) - 1) {
455  for (int i = 0; i < shape.back(); ++i) {
456  indices.back() = constants[i];
457  rewriter.create<memref::StoreOp>(loc, *elementIt, buffer, indices);
458  ++elementIt;
459  }
460  return;
461  }
462  for (int i = 0; i < shape[dim]; ++i) {
463  indices[dim] = constants[i];
464  createStores(rewriter, loc, dim + 1, buffer, shape, constants, elementIt,
465  indices);
466  }
467 }
468 
469 /// Bufferization of tensor.from_elements.
470 struct FromElementsOpInterface
471  : public BufferizableOpInterface::ExternalModel<FromElementsOpInterface,
472  tensor::FromElementsOp> {
473 
474  bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
475 
476  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
477  const BufferizationOptions &options) const {
478  auto fromElementsOp = cast<tensor::FromElementsOp>(op);
479  auto tensorType = cast<RankedTensorType>(fromElementsOp.getType());
480 
481  // Allocate a buffer for the result.
482  Location loc = op->getLoc();
483  auto shape = tensorType.getShape();
484  // TODO: Create alloc_tensor ops during TensorCopyInsertion.
485  FailureOr<Value> tensorAlloc = allocateTensorForShapedValue(
486  rewriter, loc, fromElementsOp.getResult(), options,
487  /*copy=*/false);
488  if (failed(tensorAlloc))
489  return failure();
490  FailureOr<BaseMemRefType> memrefType =
491  bufferization::getBufferType(*tensorAlloc, options);
492  if (failed(memrefType))
493  return failure();
494  Value buffer = rewriter.create<bufferization::ToMemrefOp>(
495  op->getLoc(), *memrefType, *tensorAlloc);
496 
497  // Case: tensor<0xelem_type>.
498  if (fromElementsOp.getElements().empty()) {
499  replaceOpWithBufferizedValues(rewriter, op, buffer);
500  return success();
501  }
502 
503  // Case: tensor<elem_type>.
504  if (shape.empty()) {
505  rewriter.create<memref::StoreOp>(
506  loc, fromElementsOp.getElements().front(), buffer);
507  replaceOpWithBufferizedValues(rewriter, op, buffer);
508  return success();
509  }
510 
511  // Create constants for the range of possible indices [0, max{shape_i}).
512  auto maxDim = *llvm::max_element(shape);
513  SmallVector<Value, 2> constants;
514  constants.reserve(maxDim);
515  for (int i = 0; i < maxDim; ++i)
516  constants.push_back(rewriter.create<arith::ConstantIndexOp>(loc, i));
517 
518  // Traverse all `elements` and create `memref.store` ops.
519  auto elementIt = fromElementsOp.getElements().begin();
520  SmallVector<Value, 2> indices(tensorType.getRank(), constants[0]);
521  createStores(rewriter, loc, /*dim=*/0, buffer, shape, constants, elementIt,
522  indices);
523 
524  replaceOpWithBufferizedValues(rewriter, op, buffer);
525 
526  return success();
527  }
528 };
529 
530 /// Lower the body of a tensor.generate like op (one index-typed bbArg per dim).
531 /// Such ops are lowered to linalg.map with the given tensor as a destination.
532 ///
533 /// Example:
534 /// ```
535 /// %r = tensor.generate %x, %y {
536 /// ^bb0(%arg0: index, %arg1: index):
537 /// %0 = "some_op"(%arg0, %arg1) : (index, index) -> (index)
538 /// tensor.yield %0 : index
539 /// } : tensor<?x?xindex>
540 /// ```
541 ///
542 /// Is lowered to:
543 /// ```
544 /// linalg.map ins() outs(%dest) {
545 /// %d0 = linalg.index 0 : index
546 /// %d1 = linalg.index 1 : index
547 /// %0 = "some_op"(%d0, %d1) : (index, index) -> (index)
548 /// linalg.yield %0 : index
549 /// }
550 /// ```
551 static Value lowerGenerateLikeOpBody(RewriterBase &rewriter, Location loc,
552  Value tensorDestination,
553  ValueRange dynamicSizes,
554  Region &generateBody) {
555  assert(generateBody.hasOneBlock() && "expected body with single block");
556  auto tensorType = cast<RankedTensorType>(tensorDestination.getType());
557  assert(generateBody.getNumArguments() == tensorType.getRank() &&
558  "rank mismatch");
559 
560  // Create linalg::MapOp.
561  OpBuilder::InsertionGuard g(rewriter);
562  auto linalgOp =
563  rewriter.create<linalg::MapOp>(loc, tensorType, /*inputs=*/ValueRange(),
564  /*init=*/tensorDestination);
565  Block &linalgBody = linalgOp.getMapper().emplaceBlock();
566 
567  // Create linalg::IndexOps.
568  rewriter.setInsertionPointToStart(&linalgBody);
569  SmallVector<Value> indices;
570  for (int64_t dim = 0; dim < tensorType.getRank(); ++dim)
571  indices.push_back(rewriter.create<linalg::IndexOp>(loc, dim));
572 
573  // Move over body.
574  rewriter.mergeBlocks(&generateBody.front(), &linalgBody, indices);
575  auto yieldOp = cast<tensor::YieldOp>(linalgBody.getTerminator());
576  rewriter.replaceOpWithNewOp<linalg::YieldOp>(yieldOp, yieldOp.getValue());
577 
578  return linalgOp.getResult()[0];
579 }
580 
581 /// Bufferization of tensor.generate.
582 struct GenerateOpInterface
583  : public BufferizableOpInterface::ExternalModel<GenerateOpInterface,
584  tensor::GenerateOp> {
585 
586  bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
587 
588  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
589  const BufferizationOptions &options) const {
590  auto generateOp = cast<tensor::GenerateOp>(op);
591 
592  auto type = generateOp.getResult().getType();
593 
594  // TODO: Implement memory space for this op.
595  if (options.defaultMemorySpaceFn(type) != Attribute())
596  return op->emitError("memory space not implemented yet");
597 
598  // Allocate memory.
599  Location loc = op->getLoc();
600  FailureOr<Value> tensorAlloc = allocateTensorForShapedValue(
601  rewriter, loc, generateOp.getResult(), options,
602  /*copy=*/false);
603  if (failed(tensorAlloc))
604  return failure();
605 
606  Value result = lowerGenerateLikeOpBody(rewriter, loc, *tensorAlloc,
607  generateOp.getDynamicExtents(),
608  generateOp.getBody());
609  rewriter.replaceOp(generateOp, result);
610 
611  return success();
612  }
613 };
614 
615 /// Bufferization of tensor.insert. Replace with memref.store.
616 ///
617 /// Note: DstBufferizableOpInterfaceExternalModel provides many default method
618 /// implementations for DestinationStyle ops.
619 struct InsertOpInterface
620  : public DstBufferizableOpInterfaceExternalModel<InsertOpInterface,
621  tensor::InsertOp> {
622  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
623  const BufferizationOptions &options) const {
624  auto insertOp = cast<tensor::InsertOp>(op);
625  FailureOr<Value> destMemref =
626  getBuffer(rewriter, insertOp.getDest(), options);
627  if (failed(destMemref))
628  return failure();
629  rewriter.create<memref::StoreOp>(insertOp.getLoc(), insertOp.getScalar(),
630  *destMemref, insertOp.getIndices());
631  replaceOpWithBufferizedValues(rewriter, op, *destMemref);
632  return success();
633  }
634 };
635 
636 template <typename InsertOpTy>
637 static bool insertSliceOpRequiresRead(InsertOpTy insertSliceOp,
638  OpOperand &opOperand) {
639  // The source is always read.
640  if (opOperand == insertSliceOp.getSourceMutable())
641  return true;
642 
643  // For the destination, it depends...
644  assert(opOperand == insertSliceOp.getDestMutable() && "expected dest");
645 
646  // Dest is not read if it is entirely overwritten. E.g.:
647  // tensor.insert_slice %a into %t[0][10][1] : ... into tensor<10xf32>
648  bool allOffsetsZero =
649  llvm::all_of(insertSliceOp.getMixedOffsets(), isZeroIndex);
650  RankedTensorType destType = insertSliceOp.getDestType();
651  bool sizesMatchDestSizes =
652  areConstantIntValues(insertSliceOp.getMixedSizes(), destType.getShape());
653  bool allStridesOne =
654  areAllConstantIntValue(insertSliceOp.getMixedStrides(), 1);
655  return !(allOffsetsZero && sizesMatchDestSizes && allStridesOne);
656 }
657 
658 /// Bufferization of tensor.insert_slice. Replace with a memory copy. Under
659 /// certain circumstances, this op can also be a no-op.
660 ///
661 /// Note: DstBufferizableOpInterfaceExternalModel provides many default method
662 /// implementations for DestinationStyle ops.
663 struct InsertSliceOpInterface
664  : public DstBufferizableOpInterfaceExternalModel<InsertSliceOpInterface,
665  tensor::InsertSliceOp> {
666  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
667  const AnalysisState &state) const {
668  return insertSliceOpRequiresRead(cast<tensor::InsertSliceOp>(op),
669  opOperand);
670  }
671 
672  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
673  const BufferizationOptions &options) const {
674  // insert_slice ops arise from tiling and bufferizing them out-of-place is
675  // generally a deal breaker. When used with loops, this ends up cloning the
676  // whole tensor on every single iteration and is a symptom of a
677  // catastrophically bad scheduling decision.
678  // TODO: be very loud about it or even consider failing the pass.
679  auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
680  SmallVector<OpFoldResult> mixedOffsets = insertSliceOp.getMixedOffsets();
681  SmallVector<OpFoldResult> mixedSizes = insertSliceOp.getMixedSizes();
682  SmallVector<OpFoldResult> mixedStrides = insertSliceOp.getMixedStrides();
683  Location loc = insertSliceOp.getLoc();
684 
685  // Get destination buffer.
686  FailureOr<Value> dstMemref =
687  getBuffer(rewriter, insertSliceOp.getDest(), options);
688  if (failed(dstMemref))
689  return failure();
690 
691  // Take a subview of the destination buffer.
692  auto dstMemrefType = cast<MemRefType>(dstMemref->getType());
693  MemRefType subviewMemRefType =
694  memref::SubViewOp::inferRankReducedResultType(
695  insertSliceOp.getSourceType().getShape(), dstMemrefType,
696  mixedOffsets, mixedSizes, mixedStrides);
697  Value subView = rewriter.create<memref::SubViewOp>(
698  loc, subviewMemRefType, *dstMemref, mixedOffsets, mixedSizes,
699  mixedStrides);
700 
701  // Copy tensor. If this tensor.insert_slice has a matching
702  // tensor.extract_slice, the copy operation will eventually fold away.
703  FailureOr<Value> srcMemref =
704  getBuffer(rewriter, insertSliceOp.getSource(), options);
705  if (failed(srcMemref))
706  return failure();
707  if (failed(options.createMemCpy(rewriter, loc, *srcMemref, subView)))
708  return failure();
709 
710  replaceOpWithBufferizedValues(rewriter, op, *dstMemref);
711  return success();
712  }
713 };
714 
715 /// Bufferization of tensor.pad. Replace with bufferization.alloc_tensor +
716 /// linalg.map + insert_slice.
717 /// For best performance, vectorize before bufferization (better performance in
718 /// case of padding with a constant).
719 struct PadOpInterface
720  : public BufferizableOpInterface::ExternalModel<PadOpInterface,
721  tensor::PadOp> {
722  bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
723 
724  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
725  const AnalysisState &state) const {
726  return true;
727  }
728 
729  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
730  const AnalysisState &state) const {
731  return false;
732  }
733 
734  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
735  const AnalysisState &state) const {
736  return {};
737  }
738 
739  FailureOr<BaseMemRefType>
741  SmallVector<Value> &invocationStack) const {
742  // Infer memory space from the source tensor.
743  auto padOp = cast<tensor::PadOp>(op);
744  auto maybeSrcBufferType = bufferization::getBufferType(
745  padOp.getSource(), options, invocationStack);
746  if (failed(maybeSrcBufferType))
747  return failure();
748  MemRefLayoutAttrInterface layout;
749  return MemRefType::get(padOp.getResultType().getShape(),
750  padOp.getResultType().getElementType(), layout,
751  maybeSrcBufferType->getMemorySpace());
752  }
753 
754  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
755  const BufferizationOptions &options) const {
756  auto padOp = cast<tensor::PadOp>(op);
757  Location loc = padOp.getLoc();
758  RankedTensorType resultType = padOp.getResultType();
759  RankedTensorType srcType = padOp.getSourceType();
760 
761  auto toValue = [&](OpFoldResult ofr) {
762  if (auto value = dyn_cast<Value>(ofr))
763  return value;
764  return rewriter
765  .create<arith::ConstantIndexOp>(loc, *getConstantIntValue(ofr))
766  .getResult();
767  };
768 
769  // Compute dynamic result dimensions.
770  SmallVector<OpFoldResult> mixedLowPad = padOp.getMixedLowPad();
771  SmallVector<OpFoldResult> mixedHighPad = padOp.getMixedHighPad();
772  SmallVector<Value> dynamicSizes;
773  for (int64_t i = 0; i < resultType.getRank(); ++i) {
774  if (!resultType.isDynamicDim(i))
775  continue;
776  Value srcDim = rewriter.create<tensor::DimOp>(loc, padOp.getSource(), i);
777  Value lowPad = toValue(mixedLowPad[i]);
778  Value highPad = toValue(mixedHighPad[i]);
779  AffineExpr s0, s1, s2;
780  bindSymbols(op->getContext(), s0, s1, s2);
781  AffineExpr sumExpr = s0 + s1 + s2;
782  Value sum = rewriter.create<affine::AffineApplyOp>(
783  loc, sumExpr, ValueRange{srcDim, lowPad, highPad});
784  dynamicSizes.push_back(sum);
785  }
786 
787  // Allocate a buffer for the padded result.
788  FailureOr<Value> tensorAlloc =
789  allocateTensorForShapedValue(rewriter, loc, padOp.getResult(), options,
790  /*copy=*/false);
791  if (failed(tensorAlloc))
792  return failure();
793 
794  // tensor::PadOp is like tensor::GenerateOp: The only difference is that
795  // only a part of the generated tensor is needed. For simplicity, we reuse
796  // the same functionality here.
797  Value filledBuffer = lowerGenerateLikeOpBody(
798  rewriter, loc, *tensorAlloc, dynamicSizes, padOp.getBodyRegion());
799 
800  // Create tensor::InsertSliceOp.
801  SmallVector<OpFoldResult> sliceSizes =
802  getMixedSizes(rewriter, loc, padOp.getSource());
803  SmallVector<OpFoldResult> sliceStrides(srcType.getRank(),
804  rewriter.getIndexAttr(1));
805  rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
806  padOp, padOp.getSource(), filledBuffer,
807  /*offsets=*/padOp.getMixedLowPad(), sliceSizes, sliceStrides);
808 
809  return success();
810  }
811 };
812 
813 /// Bufferization of tensor.rank. Replace with memref.rank.
814 struct RankOpInterface
815  : public BufferizableOpInterface::ExternalModel<RankOpInterface,
816  tensor::RankOp> {
817  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
818  const AnalysisState &state) const {
819  // The op reads the tensor's metadata but not its contents.
820  return false;
821  }
822 
823  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
824  const AnalysisState &state) const {
825  return false;
826  }
827 
828  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
829  const AnalysisState &state) const {
830  return {};
831  }
832 
833  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
834  const BufferizationOptions &options) const {
835  auto rankOp = cast<tensor::RankOp>(op);
836  FailureOr<Value> v = getBuffer(rewriter, rankOp.getTensor(), options);
837  if (failed(v))
838  return failure();
839  replaceOpWithNewBufferizedOp<memref::RankOp>(rewriter, op, rankOp.getType(),
840  *v);
841  return success();
842  }
843 };
844 
845 /// Bufferization of tensor.reshape. Replace with memref.reshape.
846 struct ReshapeOpInterface
847  : public BufferizableOpInterface::ExternalModel<ReshapeOpInterface,
848  tensor::ReshapeOp> {
849  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
850  const AnalysisState &state) const {
851  // Depending on the layout map, the source buffer may have to be copied.
852  auto reshapeOp = cast<tensor::ReshapeOp>(op);
853  return opOperand == reshapeOp.getShapeMutable();
854  }
855 
856  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
857  const AnalysisState &state) const {
858  return false;
859  }
860 
861  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
862  const AnalysisState &state) const {
863  // Only the 'source' operand aliases the result.
864  auto reshapeOp = cast<tensor::ReshapeOp>(op);
865  if (reshapeOp.getSourceMutable() != opOperand)
866  return {};
867  return {{op->getOpResult(0), BufferRelation::Equivalent}};
868  }
869 
870  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
871  const BufferizationOptions &options) const {
872  auto reshapeOp = cast<tensor::ReshapeOp>(op);
873  FailureOr<Value> srcBuffer =
874  getBuffer(rewriter, reshapeOp.getSource(), options);
875  FailureOr<Value> shapeBuffer =
876  getBuffer(rewriter, reshapeOp.getShape(), options);
877  if (failed(srcBuffer) || failed(shapeBuffer))
878  return failure();
879  auto maybeResultMemRefType =
880  bufferization::getBufferType(reshapeOp.getResult(), options);
881  if (failed(maybeResultMemRefType))
882  return failure();
883 
884  // memref.reshape requires the source buffer to have an identity layout.
885  // If the source memref does not have an identity layout, copy the source
886  // into a new buffer with an identity layout.
887  auto srcType = llvm::dyn_cast<MemRefType>(srcBuffer->getType());
888  if (srcType && !srcType.getLayout().isIdentity()) {
889  FailureOr<Value> tensorAlloc = allocateTensorForShapedValue(
890  rewriter, op->getLoc(), reshapeOp.getSource(), options);
891  if (failed(tensorAlloc))
892  return failure();
893  auto memrefType = MemRefType::get(
894  srcType.getShape(), srcType.getElementType(), AffineMap(),
895  cast<BaseMemRefType>(srcBuffer->getType()).getMemorySpace());
896  srcBuffer = rewriter
897  .create<bufferization::ToMemrefOp>(
898  op->getLoc(), memrefType, *tensorAlloc)
899  .getResult();
900  }
901 
902  replaceOpWithNewBufferizedOp<memref::ReshapeOp>(
903  rewriter, op, maybeResultMemRefType.value(), *srcBuffer, *shapeBuffer);
904  return success();
905  }
906 
907  FailureOr<BaseMemRefType>
909  SmallVector<Value> &invocationStack) const {
910  auto reshapeOp = cast<tensor::ReshapeOp>(op);
911  assert(value == reshapeOp.getResult() && "unexpected value provided");
912  auto maybeSourceBufferType = bufferization::getBufferType(
913  reshapeOp.getSource(), options, invocationStack);
914  if (failed(maybeSourceBufferType))
915  return failure();
917  reshapeOp.getResult().getType(),
918  cast<BaseMemRefType>(maybeSourceBufferType.value()).getMemorySpace());
919  }
920 };
921 
922 /// Analysis of ParallelInsertSliceOp.
923 struct ParallelInsertSliceOpInterface
924  : public BufferizableOpInterface::ExternalModel<
925  ParallelInsertSliceOpInterface, ParallelInsertSliceOp> {
926  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
927  const AnalysisState &state) const {
928  return {};
929  }
930 
931  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
932  const AnalysisState &state) const {
933  return opOperand == cast<ParallelInsertSliceOp>(op).getSourceMutable();
934  }
935 
936  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
937  const AnalysisState &state) const {
938  auto parallelInsertSliceOp = cast<ParallelInsertSliceOp>(op);
939  return opOperand == parallelInsertSliceOp.getDestMutable();
940  }
941 
942  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
943  const BufferizationOptions &options) const {
944  OpBuilder::InsertionGuard g(rewriter);
945  auto parallelInsertSliceOp = cast<ParallelInsertSliceOp>(op);
946  ParallelCombiningOpInterface parallelCombiningParent =
947  parallelInsertSliceOp.getParallelCombiningParent();
948 
949  // Bufferize the op outside of the parallel combining terminator.
950  rewriter.setInsertionPoint(parallelCombiningParent);
951 
952  // Get source and destination buffers.
953  FailureOr<Value> destBuffer =
954  getBuffer(rewriter, parallelInsertSliceOp.getDest(), options);
955  if (failed(destBuffer))
956  return failure();
957  FailureOr<Value> srcBuffer =
958  getBuffer(rewriter, parallelInsertSliceOp.getSource(), options);
959  if (failed(srcBuffer))
960  return failure();
961 
962  // Take a subview of the destination buffer.
963  auto destBufferType = cast<MemRefType>(destBuffer->getType());
964  MemRefType subviewMemRefType =
965  memref::SubViewOp::inferRankReducedResultType(
966  parallelInsertSliceOp.getSourceType().getShape(), destBufferType,
967  parallelInsertSliceOp.getMixedOffsets(),
968  parallelInsertSliceOp.getMixedSizes(),
969  parallelInsertSliceOp.getMixedStrides());
970  Value subview = rewriter.create<memref::SubViewOp>(
971  parallelInsertSliceOp.getLoc(), subviewMemRefType, *destBuffer,
972  parallelInsertSliceOp.getMixedOffsets(),
973  parallelInsertSliceOp.getMixedSizes(),
974  parallelInsertSliceOp.getMixedStrides());
975 
976  // This memcpy will fold away if everything bufferizes in-place.
977  if (failed(options.createMemCpy(rewriter, parallelInsertSliceOp.getLoc(),
978  *srcBuffer, subview)))
979  return failure();
980 
981  // In case the source was allocated in the same block, make sure that the
982  // deallocation op (if any) appears after the memcpy. By default, deallocs
983  // are placed before the terminator, but this does not work for ForallOp
984  // because the terminator does more than just yielding a value.
985  //
986  // Note: This is not a problem for the destination buffer because these are
987  // assumed to always bufferize in-place.
988  for (Operation *user : srcBuffer->getUsers()) {
989  if (hasEffect<MemoryEffects::Free>(user)) {
990  if (user->getBlock() == parallelCombiningParent->getBlock())
991  rewriter.moveOpBefore(user, user->getBlock()->getTerminator());
992  break;
993  }
994  }
995 
996  // Delete the op.
997  rewriter.eraseOp(op);
998  return success();
999  }
1000 
1001  /// tensor.parallel_insert_slice op has implicit inplace behavior. We
1002  /// shouldn't create copy to resolve conflict.
1003  LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
1004  const AnalysisState &state) const {
1005  return success();
1006  }
1007 };
1008 
1009 /// Bufferization of tensor.splat. Bufferizes to a new allocation that is filled
1010 /// with a linalg.map. Similar to tensor.generate.
1011 struct SplatOpInterface
1012  : public BufferizableOpInterface::ExternalModel<SplatOpInterface,
1013  tensor::SplatOp> {
1014 
1015  bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
1016 
1017  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
1018  const BufferizationOptions &options) const {
1019  OpBuilder::InsertionGuard g(rewriter);
1020  auto splatOp = cast<tensor::SplatOp>(op);
1021 
1022  // Allocate memory.
1023  Location loc = op->getLoc();
1024  FailureOr<Value> tensorAlloc = allocateTensorForShapedValue(
1025  rewriter, loc, splatOp.getResult(), options,
1026  /*copy=*/false);
1027  if (failed(tensorAlloc))
1028  return failure();
1029 
1030  // Create linalg::MapOp.
1031  auto tensorType = cast<RankedTensorType>(tensorAlloc->getType());
1032 
1033  // TODO: Implement memory space for this op.
1034  if (options.defaultMemorySpaceFn(tensorType) != Attribute())
1035  return op->emitError("memory space not implemented yet");
1036 
1037  auto linalgOp =
1038  rewriter.create<linalg::MapOp>(loc, tensorType, /*inputs=*/ValueRange(),
1039  /*init=*/*tensorAlloc);
1040  Block &linalgBody = linalgOp.getMapper().emplaceBlock();
1041 
1042  // Create linalg::IndexOps.
1043  rewriter.setInsertionPointToStart(&linalgBody);
1044  rewriter.create<linalg::YieldOp>(loc, splatOp.getInput());
1045  rewriter.replaceOp(splatOp, linalgOp.getResult()[0]);
1046 
1047  return success();
1048  }
1049 };
1050 
1051 } // namespace
1052 } // namespace tensor
1053 } // namespace mlir
1054 
1056  DialectRegistry &registry) {
1057  registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
1058  CastOp::attachInterface<CastOpInterface>(*ctx);
1059  CollapseShapeOp::attachInterface<CollapseShapeOpInterface>(*ctx);
1060  DimOp::attachInterface<DimOpInterface>(*ctx);
1061  EmptyOp::attachInterface<EmptyOpInterface>(*ctx);
1062  ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
1063  ExtractSliceOp::attachInterface<ExtractSliceOpInterface>(*ctx);
1064  ExtractOp::attachInterface<ExtractOpInterface>(*ctx);
1065  FromElementsOp::attachInterface<FromElementsOpInterface>(*ctx);
1066  GenerateOp::attachInterface<GenerateOpInterface>(*ctx);
1067  InsertOp::attachInterface<InsertOpInterface>(*ctx);
1068  InsertSliceOp::attachInterface<InsertSliceOpInterface>(*ctx);
1069  PadOp::attachInterface<PadOpInterface>(*ctx);
1070  ParallelInsertSliceOp::attachInterface<ParallelInsertSliceOpInterface>(
1071  *ctx);
1072  RankOp::attachInterface<RankOpInterface>(*ctx);
1073  ReshapeOp::attachInterface<ReshapeOpInterface>(*ctx);
1074  SplatOp::attachInterface<SplatOpInterface>(*ctx);
1075 
1076  // Load additional dialects of which ops may get created.
1077  ctx->loadDialect<arith::ArithDialect, linalg::LinalgDialect>();
1078  });
1079 
1080  // Bufferization requires SubsetInsertionOpInterface models. Make sure that
1081  // they are registered.
1083 }
static llvm::ManagedStatic< PassManagerOptions > options
Base type for affine expression.
Definition: AffineExpr.h:68
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:46
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
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:104
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
bool addExtension(TypeID extensionID, std::unique_ptr< DialectExtensionBase > extension)
Add the given extension to the registry.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:66
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
void loadDialect()
Load a dialect in the context.
Definition: MLIRContext.h:107
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:346
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:453
This class represents a single result from folding an operation.
Definition: OpDefinition.h:271
This class represents an operand of an operation.
Definition: Value.h:243
This is a value defined by a result of an operation.
Definition: Value.h:433
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
OpResult getOpResult(unsigned idx)
Definition: Operation.h:421
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
Definition: Operation.h:407
MLIRContext * getContext()
Return the context this operation is associated with.
Definition: Operation.h:216
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:223
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Definition: Operation.cpp:268
user_range getUsers()
Returns a range of all users.
Definition: Operation.h:874
use_range getUses()
Returns a range of all uses, which is useful for iterating over all uses.
Definition: Operation.h:847
This class contains a list of basic blocks and a link to the parent operation it is attached to.
Definition: Region.h:26
unsigned getNumArguments()
Definition: Region.h:123
Block & front()
Definition: Region.h:65
bool hasOneBlock()
Return true if this region has exactly one block.
Definition: Region.h:68
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:358
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
void mergeBlocks(Block *source, Block *dest, ValueRange argValues=std::nullopt)
Inline the operations of block 'source' into the end of block 'dest'.
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
void moveOpBefore(Operation *op, Operation *existingOp)
Unlink this operation from its current block and insert it right before existingOp which may be in th...
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
void replaceOpWithBufferizedValues(RewriterBase &rewriter, Operation *op, ValueRange values)
Replace an op with replacement values.
BaseMemRefType getMemRefTypeWithStaticIdentityLayout(TensorType tensorType, Attribute memorySpace=nullptr)
Return a MemRef type with a static identity layout (i.e., no layout map).
FailureOr< Value > allocateTensorForShapedValue(OpBuilder &b, Location loc, Value shapedValue, const BufferizationOptions &options, bool copy=true)
Create an AllocTensorOp for the given shaped value (memref or tensor).
FailureOr< BaseMemRefType > getBufferType(Value value, const BufferizationOptions &options)
Return the buffer type for a given Value (tensor) after bufferization without bufferizing any IR.
FailureOr< Value > getBuffer(RewriterBase &rewriter, Value value, const BufferizationOptions &options)
Lookup the buffer for the given value.
BaseMemRefType getMemRefTypeWithFullyDynamicLayout(TensorType tensorType, Attribute memorySpace=nullptr)
Return a MemRef type with fully dynamic layout.
SmallVector< OpFoldResult > getMixedSizes(OpBuilder &builder, Location loc, Value value)
Return the dimensions of the given memref value.
Definition: MemRefOps.cpp:78
void registerSubsetOpInterfaceExternalModels(DialectRegistry &registry)
void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry)
Include the generated interface declarations.
bool isZeroIndex(OpFoldResult v)
Return true if v is an IntegerAttr with value 0 of a ConstantIndexOp with attribute with value 0.
bool areConstantIntValues(ArrayRef< OpFoldResult > ofrs, ArrayRef< int64_t > values)
Return true if all of ofrs are constant integers equal to the corresponding value in values.
std::optional< int64_t > getConstantIntValue(OpFoldResult ofr)
If ofr is a constant integer or an IntegerAttr, return the integer.
bool areAllConstantIntValue(ArrayRef< OpFoldResult > ofrs, int64_t value)
Return true if all of ofrs are constant integers equal to value.
void bindSymbols(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to SymbolExpr at positions: [0 .
Definition: AffineExpr.h:325
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
Options for BufferizableOpInterface-based bufferization.
Bufferizable ops that implement the DestinationStyleOpInterface can use this external model base clas...