MLIR  19.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 
22 #include "mlir/IR/Dialect.h"
23 #include "mlir/IR/Operation.h"
24 
25 using namespace mlir;
26 using namespace mlir::bufferization;
27 using namespace mlir::tensor;
28 
29 namespace mlir {
30 namespace tensor {
31 namespace {
32 
33 struct CastOpInterface
34  : public BufferizableOpInterface::ExternalModel<CastOpInterface,
35  tensor::CastOp> {
36  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
37  const AnalysisState &state) const {
38  return false;
39  }
40 
41  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
42  const AnalysisState &state) const {
43  return false;
44  }
45 
46  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
47  const AnalysisState &state) const {
48  return {{op->getResult(0), BufferRelation::Equivalent}};
49  }
50 
53  SmallVector<Value> &invocationStack) const {
54  auto castOp = cast<tensor::CastOp>(op);
55  auto maybeSrcBufferType = bufferization::getBufferType(
56  castOp.getSource(), options, invocationStack);
57  if (failed(maybeSrcBufferType))
58  return failure();
59  Attribute memorySpace = maybeSrcBufferType->getMemorySpace();
60 
61  // Note: `getMemRefTypeWithFullyDynamicLayout` returns an unranked memref
62  // type in case the input is an unranked tensor type.
63 
64  // Case 1: Casting an unranked tensor
65  if (isa<UnrankedTensorType>(castOp.getSource().getType())) {
66  // When casting to a ranked tensor, we cannot infer any static offset or
67  // strides from the source. Assume fully dynamic.
68  return getMemRefTypeWithFullyDynamicLayout(castOp.getType(), memorySpace);
69  }
70 
71  // Case 2: Casting to an unranked tensor type
72  if (isa<UnrankedTensorType>(castOp.getType())) {
73  return getMemRefTypeWithFullyDynamicLayout(castOp.getType(), memorySpace);
74  }
75 
76  // Case 3: Ranked tensor -> ranked tensor. The offsets and strides do not
77  // change.
78  auto rankedResultType = cast<RankedTensorType>(castOp.getType());
79  return MemRefType::get(
80  rankedResultType.getShape(), rankedResultType.getElementType(),
81  llvm::cast<MemRefType>(*maybeSrcBufferType).getLayout(), memorySpace);
82  }
83 
84  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
85  const BufferizationOptions &options) const {
86  auto castOp = cast<tensor::CastOp>(op);
87 
88  // The result buffer still has the old (pre-cast) type.
89  FailureOr<Value> resultBuffer =
90  getBuffer(rewriter, castOp.getSource(), options);
91  if (failed(resultBuffer))
92  return failure();
93 
94  // Compute the new type.
95  auto resultMemRefType =
96  bufferization::getBufferType(castOp.getResult(), options);
97  if (failed(resultMemRefType))
98  return failure();
99  if (resultBuffer->getType() == *resultMemRefType) {
100  // This cast is a no-op.
101  replaceOpWithBufferizedValues(rewriter, op, *resultBuffer);
102  return success();
103  }
104 
105  // Replace the op with a memref.cast.
106  assert(memref::CastOp::areCastCompatible(resultBuffer->getType(),
107  *resultMemRefType) &&
108  "CallOp::bufferize: cast incompatible");
109  replaceOpWithNewBufferizedOp<memref::CastOp>(
110  rewriter, op, *resultMemRefType, *resultBuffer);
111 
112  return success();
113  }
114 };
115 
116 /// Bufferization of tensor.collapse_shape. Replace with memref.collapse_shape.
117 struct CollapseShapeOpInterface
118  : public BufferizableOpInterface::ExternalModel<CollapseShapeOpInterface,
119  tensor::CollapseShapeOp> {
120  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
121  const AnalysisState &state) const {
122  // tensor.collapse_shape may reallocate, at which point the source buffer is
123  // copied. I.e., there will be a memory read side effect on the bufferized
124  // source. This function conservatively returns "true" because whether a
125  // copy will be created or not is not known at this point.
126  return true;
127  }
128 
129  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
130  const AnalysisState &state) const {
131  return false;
132  }
133 
134  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
135  const AnalysisState &state) const {
136  // TODO: CollapseShapeOp may allocate at runtime.
137  return {{op->getOpResult(0), BufferRelation::Equivalent}};
138  }
139 
142  SmallVector<Value> &invocationStack) const {
143  auto collapseShapeOp = cast<tensor::CollapseShapeOp>(op);
144  auto maybeSrcBufferType = bufferization::getBufferType(
145  collapseShapeOp.getSrc(), options, invocationStack);
146  if (failed(maybeSrcBufferType))
147  return failure();
148  auto srcBufferType = llvm::cast<MemRefType>(*maybeSrcBufferType);
149  bool canBeCollapsed = memref::CollapseShapeOp::isGuaranteedCollapsible(
150  srcBufferType, collapseShapeOp.getReassociationIndices());
151 
152  if (!canBeCollapsed) {
153  // If dims cannot be collapsed, this op bufferizes to a new allocation.
154  RankedTensorType tensorResultType = collapseShapeOp.getResultType();
156  tensorResultType, srcBufferType.getMemorySpace());
157  }
158 
159  return memref::CollapseShapeOp::computeCollapsedType(
160  srcBufferType, collapseShapeOp.getReassociationIndices());
161  }
162 
163  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
164  const BufferizationOptions &options) const {
165  auto collapseShapeOp = cast<tensor::CollapseShapeOp>(op);
166  RankedTensorType tensorResultType = collapseShapeOp.getResultType();
167  FailureOr<Value> maybeBuffer =
168  getBuffer(rewriter, collapseShapeOp.getSrc(), options);
169  if (failed(maybeBuffer))
170  return failure();
171  Value buffer = *maybeBuffer;
172  auto bufferType = cast<MemRefType>(buffer.getType());
173 
174  if (tensorResultType.getRank() == 0) {
175  // 0-d collapses must go through a different op builder.
176  MemRefType resultType;
177 
178  if (bufferType.getLayout().isIdentity()) {
179  // Standard layout: result type has no offset.
180  MemRefLayoutAttrInterface layout;
181  resultType = MemRefType::get({}, tensorResultType.getElementType(),
182  layout, bufferType.getMemorySpace());
183  } else {
184  // Source memref has a layout map: result type has the same offset as
185  // the source type.
186  SmallVector<int64_t> strides;
187  int64_t offset;
188  if (failed(getStridesAndOffset(bufferType, strides, offset)))
189  return failure();
190  resultType = MemRefType::get(
191  {}, tensorResultType.getElementType(),
192  StridedLayoutAttr::get(op->getContext(), offset, {}),
193  bufferType.getMemorySpace());
194  }
195 
196  replaceOpWithNewBufferizedOp<memref::CollapseShapeOp>(
197  rewriter, op, resultType, buffer, collapseShapeOp.getReassociation());
198  return success();
199  }
200 
201  // If the dims are not collapsible (due to an incompatible source layout
202  // map), force an out-of-place bufferization, i.e., a buffer copy. This
203  // newly allocated buffer will have no layout map and thus be collapsible.
204  bool canBeCollapsed = memref::CollapseShapeOp::isGuaranteedCollapsible(
205  bufferType, collapseShapeOp.getReassociationIndices());
206  if (!canBeCollapsed) {
207  // TODO: Create alloc_tensor ops during TensorCopyInsertion.
208  AnalysisState analysisState(options);
210  rewriter, op->getLoc(), collapseShapeOp.getSrc(), options);
211  if (failed(tensorAlloc))
212  return failure();
213  auto memrefType =
214  MemRefType::get(collapseShapeOp.getSrcType().getShape(),
215  collapseShapeOp.getSrcType().getElementType(),
216  AffineMap(), bufferType.getMemorySpace());
217  buffer = rewriter.create<bufferization::ToMemrefOp>(
218  op->getLoc(), memrefType, *tensorAlloc);
219  }
220 
221  // Result type is inferred by the builder.
222  replaceOpWithNewBufferizedOp<memref::CollapseShapeOp>(
223  rewriter, op, buffer, collapseShapeOp.getReassociationIndices());
224  return success();
225  }
226 };
227 
228 /// Bufferization of tensor.dim. Replace with memref.dim.
229 struct DimOpInterface
230  : public BufferizableOpInterface::ExternalModel<DimOpInterface,
231  tensor::DimOp> {
232  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
233  const AnalysisState &state) const {
234  // The op reads the tensor's metadata but not its contents.
235  return false;
236  }
237 
238  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
239  const AnalysisState &state) const {
240  return false;
241  }
242 
243  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
244  const AnalysisState &state) const {
245  return {};
246  }
247 
248  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
249  const BufferizationOptions &options) const {
250  auto dimOp = cast<tensor::DimOp>(op);
251  FailureOr<Value> v = getBuffer(rewriter, dimOp.getSource(), options);
252  if (failed(v))
253  return failure();
254  replaceOpWithNewBufferizedOp<memref::DimOp>(rewriter, op, *v,
255  dimOp.getIndex());
256  return success();
257  }
258 };
259 
260 /// Bufferization of "tensor.empty". Replace with "bufferization.alloc_tensor".
261 struct EmptyOpInterface
262  : public BufferizableOpInterface::ExternalModel<EmptyOpInterface,
263  tensor::EmptyOp> {
264  bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
265 
266  bool resultBufferizesToMemoryWrite(Operation *op, OpResult opResult,
267  const AnalysisState &state) const {
268  // The returned tensor does not have specified contents.
269  return false;
270  }
271 
272  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
273  const BufferizationOptions &options) const {
274  auto emptyOp = cast<tensor::EmptyOp>(op);
275 
276  // Optimization: Fold away the op if it has no uses.
277  if (op->getUses().empty()) {
278  rewriter.eraseOp(op);
279  return success();
280  }
281 
282  // Allocate a tensor. This emits a "bufferization.alloc_tensor" op.
284  rewriter, op->getLoc(), emptyOp.getResult(), options, /*copy=*/false);
285  if (failed(allocTensor))
286  return failure();
287  rewriter.replaceOp(op, *allocTensor);
288  return success();
289  }
290 };
291 
292 /// Bufferization of tensor.expand_shape. Replace with memref.expand_shape.
293 struct ExpandShapeOpInterface
294  : public BufferizableOpInterface::ExternalModel<ExpandShapeOpInterface,
295  tensor::ExpandShapeOp> {
296  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
297  const AnalysisState &state) const {
298  // In contrast to tensor.collapse_shape, this op can always be bufferized
299  // without a copy.
300  return false;
301  }
302 
303  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
304  const AnalysisState &state) const {
305  return false;
306  }
307 
308  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
309  const AnalysisState &state) const {
310  return {{op->getOpResult(0), BufferRelation::Equivalent}};
311  }
312 
315  SmallVector<Value> &invocationStack) const {
316  auto expandShapeOp = cast<tensor::ExpandShapeOp>(op);
317  auto maybeSrcBufferType = bufferization::getBufferType(
318  expandShapeOp.getSrc(), options, invocationStack);
319  if (failed(maybeSrcBufferType))
320  return failure();
321  auto srcBufferType = llvm::cast<MemRefType>(*maybeSrcBufferType);
322  auto maybeResultType = memref::ExpandShapeOp::computeExpandedType(
323  srcBufferType, expandShapeOp.getResultType().getShape(),
324  expandShapeOp.getReassociationIndices());
325  if (failed(maybeResultType))
326  return failure();
327  return *maybeResultType;
328  }
329 
330  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
331  const BufferizationOptions &options) const {
332  auto expandShapeOp = cast<tensor::ExpandShapeOp>(op);
333  auto tensorResultType = expandShapeOp.getResultType();
334  FailureOr<Value> buffer =
335  getBuffer(rewriter, expandShapeOp.getSrc(), options);
336  if (failed(buffer))
337  return failure();
338 
339  // Memref result type is inferred by the builder based on reassociation
340  // indices and result shape.
341  // TODO: Instead of inferring the output shape argument of
342  // memref.expand_shape op, use output_shape argument of tensor.expand_shape
343  // op.
344  replaceOpWithNewBufferizedOp<memref::ExpandShapeOp>(
345  rewriter, op, tensorResultType.getShape(), *buffer,
346  expandShapeOp.getReassociationIndices());
347  return success();
348  }
349 };
350 
351 /// Bufferization of tensor.extract_slice. Replace with memref.subview.
352 struct ExtractSliceOpInterface
353  : public BufferizableOpInterface::ExternalModel<ExtractSliceOpInterface,
354  tensor::ExtractSliceOp> {
355  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
356  const AnalysisState &state) const {
357  return false;
358  }
359 
360  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
361  const AnalysisState &state) const {
362  return false;
363  }
364 
365  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
366  const AnalysisState &state) const {
367  return {{op->getOpResult(0), BufferRelation::Unknown}};
368  }
369 
370  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
371  const BufferizationOptions &options) const {
372  auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
373  SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
374  SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
375  SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
376  Location loc = extractSliceOp.getLoc();
377 
378  // Get source buffer.
379  FailureOr<Value> srcMemref =
380  getBuffer(rewriter, extractSliceOp.getSource(), options);
381  if (failed(srcMemref))
382  return failure();
383 
384  // Take a subview of the source buffer.
385  auto resultMemrefType =
386  bufferization::getBufferType(extractSliceOp.getResult(), options);
387  if (failed(resultMemrefType))
388  return failure();
389  Value subView = rewriter.create<memref::SubViewOp>(
390  loc, llvm::cast<MemRefType>(*resultMemrefType), *srcMemref, mixedOffsets,
391  mixedSizes, mixedStrides);
392 
393  replaceOpWithBufferizedValues(rewriter, op, subView);
394  return success();
395  }
396 
399  SmallVector<Value> &invocationStack) const {
400  auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
401  assert(value == extractSliceOp.getResult() && "invalid value");
402  auto srcMemrefType = bufferization::getBufferType(
403  extractSliceOp.getSource(), options, invocationStack);
404  if (failed(srcMemrefType))
405  return failure();
406  SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
407  SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
408  SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
409  return cast<BaseMemRefType>(memref::SubViewOp::inferRankReducedResultType(
410  extractSliceOp.getType().getShape(), llvm::cast<MemRefType>(*srcMemrefType),
411  mixedOffsets, mixedSizes, 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  // TODO: Implement memory space for this op.
482  if (options.defaultMemorySpaceFn(tensorType) != Attribute())
483  return op->emitError("memory space not implemented yet");
484 
485  // Allocate a buffer for the result.
486  Location loc = op->getLoc();
487  auto shape = tensorType.getShape();
488  // TODO: Create alloc_tensor ops during TensorCopyInsertion.
490  rewriter, loc, fromElementsOp.getResult(), options,
491  /*copy=*/false);
492  if (failed(tensorAlloc))
493  return failure();
494  auto memrefType =
495  MemRefType::get(tensorType.getShape(), tensorType.getElementType());
496  Value buffer = rewriter.create<bufferization::ToMemrefOp>(
497  op->getLoc(), memrefType, *tensorAlloc);
498 
499  // Case: tensor<0xelem_type>.
500  if (fromElementsOp.getElements().empty()) {
501  replaceOpWithBufferizedValues(rewriter, op, buffer);
502  return success();
503  }
504 
505  // Case: tensor<elem_type>.
506  if (shape.empty()) {
507  rewriter.create<memref::StoreOp>(
508  loc, fromElementsOp.getElements().front(), buffer);
509  replaceOpWithBufferizedValues(rewriter, op, buffer);
510  return success();
511  }
512 
513  // Create constants for the range of possible indices [0, max{shape_i}).
514  auto maxDim = *llvm::max_element(shape);
515  SmallVector<Value, 2> constants;
516  constants.reserve(maxDim);
517  for (int i = 0; i < maxDim; ++i)
518  constants.push_back(rewriter.create<arith::ConstantIndexOp>(loc, i));
519 
520  // Traverse all `elements` and create `memref.store` ops.
521  auto elementIt = fromElementsOp.getElements().begin();
522  SmallVector<Value, 2> indices(tensorType.getRank(), constants[0]);
523  createStores(rewriter, loc, /*dim=*/0, buffer, shape, constants, elementIt,
524  indices);
525 
526  replaceOpWithBufferizedValues(rewriter, op, buffer);
527 
528  return success();
529  }
530 };
531 
532 /// Lower the body of a tensor.generate like op (one index-typed bbArg per dim).
533 /// Such ops are lowered to linalg.map with the given tensor as a destination.
534 ///
535 /// Example:
536 /// ```
537 /// %r = tensor.generate %x, %y {
538 /// ^bb0(%arg0: index, %arg1: index):
539 /// %0 = "some_op"(%arg0, %arg1) : (index, index) -> (index)
540 /// tensor.yield %0 : index
541 /// } : tensor<?x?xindex>
542 /// ```
543 ///
544 /// Is lowered to:
545 /// ```
546 /// linalg.map ins() outs(%dest) {
547 /// %d0 = linalg.index 0 : index
548 /// %d1 = linalg.index 1 : index
549 /// %0 = "some_op"(%d0, %d1) : (index, index) -> (index)
550 /// linalg.yield %0 : index
551 /// }
552 /// ```
553 static Value lowerGenerateLikeOpBody(RewriterBase &rewriter, Location loc,
554  Value tensorDestination,
555  ValueRange dynamicSizes,
556  Region &generateBody) {
557  assert(generateBody.hasOneBlock() && "expected body with single block");
558  auto tensorType = cast<RankedTensorType>(tensorDestination.getType());
559  assert(generateBody.getNumArguments() == tensorType.getRank() &&
560  "rank mismatch");
561 
562  // Create linalg::MapOp.
563  OpBuilder::InsertionGuard g(rewriter);
564  auto linalgOp =
565  rewriter.create<linalg::MapOp>(loc, tensorType, /*inputs=*/ValueRange(),
566  /*init=*/tensorDestination);
567  Block &linalgBody = linalgOp.getMapper().emplaceBlock();
568 
569  // Create linalg::IndexOps.
570  rewriter.setInsertionPointToStart(&linalgBody);
571  SmallVector<Value> indices;
572  for (int64_t dim = 0; dim < tensorType.getRank(); ++dim)
573  indices.push_back(rewriter.create<linalg::IndexOp>(loc, dim));
574 
575  // Move over body.
576  rewriter.mergeBlocks(&generateBody.front(), &linalgBody, indices);
577  auto yieldOp = cast<tensor::YieldOp>(linalgBody.getTerminator());
578  rewriter.replaceOpWithNewOp<linalg::YieldOp>(yieldOp, yieldOp.getValue());
579 
580  return linalgOp.getResult()[0];
581 }
582 
583 /// Bufferization of tensor.generate.
584 struct GenerateOpInterface
585  : public BufferizableOpInterface::ExternalModel<GenerateOpInterface,
586  tensor::GenerateOp> {
587 
588  bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
589 
590  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
591  const BufferizationOptions &options) const {
592  auto generateOp = cast<tensor::GenerateOp>(op);
593 
594  auto type = generateOp.getResult().getType();
595 
596  // TODO: Implement memory space for this op.
597  if (options.defaultMemorySpaceFn(type) != Attribute())
598  return op->emitError("memory space not implemented yet");
599 
600  // Allocate memory.
601  Location loc = op->getLoc();
603  rewriter, loc, generateOp.getResult(), options,
604  /*copy=*/false);
605  if (failed(tensorAlloc))
606  return failure();
607 
608  Value result = lowerGenerateLikeOpBody(rewriter, loc, *tensorAlloc,
609  generateOp.getDynamicExtents(),
610  generateOp.getBody());
611  rewriter.replaceOp(generateOp, result);
612 
613  return success();
614  }
615 };
616 
617 /// Bufferization of tensor.insert. Replace with memref.store.
618 ///
619 /// Note: DstBufferizableOpInterfaceExternalModel provides many default method
620 /// implementations for DestinationStyle ops.
621 struct InsertOpInterface
622  : public DstBufferizableOpInterfaceExternalModel<InsertOpInterface,
623  tensor::InsertOp> {
624  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
625  const BufferizationOptions &options) const {
626  auto insertOp = cast<tensor::InsertOp>(op);
627  FailureOr<Value> destMemref =
628  getBuffer(rewriter, insertOp.getDest(), options);
629  if (failed(destMemref))
630  return failure();
631  rewriter.create<memref::StoreOp>(insertOp.getLoc(), insertOp.getScalar(),
632  *destMemref, insertOp.getIndices());
633  replaceOpWithBufferizedValues(rewriter, op, *destMemref);
634  return success();
635  }
636 };
637 
638 /// Bufferization of tensor.insert_slice. Replace with a memory copy. Under
639 /// certain circumstances, this op can also be a no-op.
640 ///
641 /// Note: DstBufferizableOpInterfaceExternalModel provides many default method
642 /// implementations for DestinationStyle ops.
643 struct InsertSliceOpInterface
644  : public DstBufferizableOpInterfaceExternalModel<InsertSliceOpInterface,
645  tensor::InsertSliceOp> {
646  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
647  const AnalysisState &state) const {
648  auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
649  RankedTensorType destType = insertSliceOp.getDestType();
650 
651  // The source is always read.
652  if (opOperand == insertSliceOp.getSourceMutable())
653  return true;
654 
655  // For the destination, it depends...
656  assert(opOperand == insertSliceOp.getDestMutable() && "expected dest");
657 
658  // Dest is not read if it is entirely overwritten. E.g.:
659  // tensor.insert_slice %a into %t[0][10][1] : ... into tensor<10xf32>
660  bool allOffsetsZero =
661  llvm::all_of(insertSliceOp.getMixedOffsets(), [](OpFoldResult ofr) {
662  return isConstantIntValue(ofr, 0);
663  });
664  bool sizesMatchDestSizes = llvm::all_of(
665  llvm::enumerate(insertSliceOp.getMixedSizes()), [&](const auto &it) {
666  return getConstantIntValue(it.value()) ==
667  destType.getDimSize(it.index());
668  });
669  bool allStridesOne =
670  llvm::all_of(insertSliceOp.getMixedStrides(), [](OpFoldResult ofr) {
671  return isConstantIntValue(ofr, 1);
672  });
673  return !(allOffsetsZero && sizesMatchDestSizes && allStridesOne);
674  }
675 
676  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
677  const BufferizationOptions &options) const {
678  // insert_slice ops arise from tiling and bufferizing them out-of-place is
679  // generally a deal breaker. When used with loops, this ends up cloning the
680  // whole tensor on every single iteration and is a symptom of a
681  // catastrophically bad scheduling decision.
682  // TODO: be very loud about it or even consider failing the pass.
683  auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
684  SmallVector<OpFoldResult> mixedOffsets = insertSliceOp.getMixedOffsets();
685  SmallVector<OpFoldResult> mixedSizes = insertSliceOp.getMixedSizes();
686  SmallVector<OpFoldResult> mixedStrides = insertSliceOp.getMixedStrides();
687  Location loc = insertSliceOp.getLoc();
688 
689  // Get destination buffer.
690  FailureOr<Value> dstMemref =
691  getBuffer(rewriter, insertSliceOp.getDest(), options);
692  if (failed(dstMemref))
693  return failure();
694 
695  // Take a subview of the destination buffer.
696  auto dstMemrefType = cast<MemRefType>(dstMemref->getType());
697  auto subviewMemRefType =
698  cast<MemRefType>(memref::SubViewOp::inferRankReducedResultType(
699  insertSliceOp.getSourceType().getShape(), dstMemrefType,
700  mixedOffsets, mixedSizes, mixedStrides));
701  Value subView = rewriter.create<memref::SubViewOp>(
702  loc, subviewMemRefType, *dstMemref, mixedOffsets, mixedSizes,
703  mixedStrides);
704 
705  // Copy tensor. If this tensor.insert_slice has a matching
706  // tensor.extract_slice, the copy operation will eventually fold away.
707  FailureOr<Value> srcMemref =
708  getBuffer(rewriter, insertSliceOp.getSource(), options);
709  if (failed(srcMemref))
710  return failure();
711  if (failed(options.createMemCpy(rewriter, loc, *srcMemref, subView)))
712  return failure();
713 
714  replaceOpWithBufferizedValues(rewriter, op, *dstMemref);
715  return success();
716  }
717 };
718 
719 /// Bufferization of tensor.pad. Replace with bufferization.alloc_tensor +
720 /// linalg.map + insert_slice.
721 /// For best performance, vectorize before bufferization (better performance in
722 /// case of padding with a constant).
723 struct PadOpInterface
724  : public BufferizableOpInterface::ExternalModel<PadOpInterface,
725  tensor::PadOp> {
726  bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
727 
728  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
729  const AnalysisState &state) const {
730  return true;
731  }
732 
733  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
734  const AnalysisState &state) const {
735  return false;
736  }
737 
738  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
739  const AnalysisState &state) const {
740  return {};
741  }
742 
745  SmallVector<Value> &invocationStack) const {
746  // Infer memory space from the source tensor.
747  auto padOp = cast<tensor::PadOp>(op);
748  auto maybeSrcBufferType = bufferization::getBufferType(
749  padOp.getSource(), options, invocationStack);
750  if (failed(maybeSrcBufferType))
751  return failure();
752  MemRefLayoutAttrInterface layout;
753  return MemRefType::get(padOp.getResultType().getShape(),
754  padOp.getResultType().getElementType(), layout,
755  maybeSrcBufferType->getMemorySpace());
756  }
757 
758  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
759  const BufferizationOptions &options) const {
760  auto padOp = cast<tensor::PadOp>(op);
761  Location loc = padOp.getLoc();
762  RankedTensorType resultType = padOp.getResultType();
763  RankedTensorType srcType = padOp.getSourceType();
764 
765  auto toValue = [&](OpFoldResult ofr) {
766  if (ofr.is<Value>())
767  return ofr.get<Value>();
768  return rewriter
769  .create<arith::ConstantIndexOp>(loc, *getConstantIntValue(ofr))
770  .getResult();
771  };
772 
773  // Compute dynamic result dimensions.
774  SmallVector<OpFoldResult> mixedLowPad = padOp.getMixedLowPad();
775  SmallVector<OpFoldResult> mixedHighPad = padOp.getMixedHighPad();
776  SmallVector<Value> dynamicSizes;
777  for (int64_t i = 0; i < resultType.getRank(); ++i) {
778  if (!resultType.isDynamicDim(i))
779  continue;
780  Value srcDim = rewriter.create<tensor::DimOp>(loc, padOp.getSource(), i);
781  Value lowPad = toValue(mixedLowPad[i]);
782  Value highPad = toValue(mixedHighPad[i]);
783  AffineExpr s0, s1, s2;
784  bindSymbols(op->getContext(), s0, s1, s2);
785  AffineExpr sumExpr = s0 + s1 + s2;
786  Value sum = rewriter.create<affine::AffineApplyOp>(
787  loc, sumExpr, ValueRange{srcDim, lowPad, highPad});
788  dynamicSizes.push_back(sum);
789  }
790 
791  // Allocate a buffer for the padded result.
792  FailureOr<Value> tensorAlloc =
793  allocateTensorForShapedValue(rewriter, loc, padOp.getResult(), options,
794  /*copy=*/false);
795  if (failed(tensorAlloc))
796  return failure();
797 
798  // tensor::PadOp is like tensor::GenerateOp: The only difference is that
799  // only a part of the generated tensor is needed. For simplicity, we reuse
800  // the same functionality here.
801  Value filledBuffer = lowerGenerateLikeOpBody(
802  rewriter, loc, *tensorAlloc, dynamicSizes, padOp.getBodyRegion());
803 
804  // Create tensor::InsertSliceOp.
805  SmallVector<OpFoldResult> sliceSizes =
806  getMixedSizes(rewriter, loc, padOp.getSource());
807  SmallVector<OpFoldResult> sliceStrides(srcType.getRank(),
808  rewriter.getIndexAttr(1));
809  rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
810  padOp, padOp.getSource(), filledBuffer,
811  /*offsets=*/padOp.getMixedLowPad(), sliceSizes, sliceStrides);
812 
813  return success();
814  }
815 };
816 
817 /// Bufferization of tensor.rank. Replace with memref.rank.
818 struct RankOpInterface
819  : public BufferizableOpInterface::ExternalModel<RankOpInterface,
820  tensor::RankOp> {
821  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
822  const AnalysisState &state) const {
823  // The op reads the tensor's metadata but not its contents.
824  return false;
825  }
826 
827  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
828  const AnalysisState &state) const {
829  return false;
830  }
831 
832  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
833  const AnalysisState &state) const {
834  return {};
835  }
836 
837  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
838  const BufferizationOptions &options) const {
839  auto rankOp = cast<tensor::RankOp>(op);
840  FailureOr<Value> v = getBuffer(rewriter, rankOp.getTensor(), options);
841  if (failed(v))
842  return failure();
843  replaceOpWithNewBufferizedOp<memref::RankOp>(rewriter, op, rankOp.getType(),
844  *v);
845  return success();
846  }
847 };
848 
849 /// Bufferization of tensor.reshape. Replace with memref.reshape.
850 struct ReshapeOpInterface
851  : public BufferizableOpInterface::ExternalModel<ReshapeOpInterface,
852  tensor::ReshapeOp> {
853  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
854  const AnalysisState &state) const {
855  // Depending on the layout map, the source buffer may have to be copied.
856  auto reshapeOp = cast<tensor::ReshapeOp>(op);
857  return opOperand == reshapeOp.getShapeMutable();
858  }
859 
860  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
861  const AnalysisState &state) const {
862  return false;
863  }
864 
865  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
866  const AnalysisState &state) const {
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()) {
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 
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 true;
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  auto subviewMemRefType =
965  cast<MemRefType>(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 
1002 /// Bufferization of tensor.splat. Bufferizes to a new allocation that is filled
1003 /// with a linalg.map. Similar to tensor.generate.
1004 struct SplatOpInterface
1005  : public BufferizableOpInterface::ExternalModel<SplatOpInterface,
1006  tensor::SplatOp> {
1007 
1008  bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
1009 
1010  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
1011  const BufferizationOptions &options) const {
1012  OpBuilder::InsertionGuard g(rewriter);
1013  auto splatOp = cast<tensor::SplatOp>(op);
1014 
1015  // Allocate memory.
1016  Location loc = op->getLoc();
1018  rewriter, loc, splatOp.getResult(), options,
1019  /*copy=*/false);
1020  if (failed(tensorAlloc))
1021  return failure();
1022 
1023  // Create linalg::MapOp.
1024  auto tensorType = cast<RankedTensorType>(tensorAlloc->getType());
1025 
1026  // TODO: Implement memory space for this op.
1027  if (options.defaultMemorySpaceFn(tensorType) != Attribute())
1028  return op->emitError("memory space not implemented yet");
1029 
1030  auto linalgOp =
1031  rewriter.create<linalg::MapOp>(loc, tensorType, /*inputs=*/ValueRange(),
1032  /*init=*/*tensorAlloc);
1033  Block &linalgBody = linalgOp.getMapper().emplaceBlock();
1034 
1035  // Create linalg::IndexOps.
1036  rewriter.setInsertionPointToStart(&linalgBody);
1037  rewriter.create<linalg::YieldOp>(loc, splatOp.getInput());
1038  rewriter.replaceOp(splatOp, linalgOp.getResult()[0]);
1039 
1040  return success();
1041  }
1042 };
1043 
1044 } // namespace
1045 } // namespace tensor
1046 } // namespace mlir
1047 
1049  DialectRegistry &registry) {
1050  registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
1051  CastOp::attachInterface<CastOpInterface>(*ctx);
1052  CollapseShapeOp::attachInterface<CollapseShapeOpInterface>(*ctx);
1053  DimOp::attachInterface<DimOpInterface>(*ctx);
1054  EmptyOp::attachInterface<EmptyOpInterface>(*ctx);
1055  ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
1056  ExtractSliceOp::attachInterface<ExtractSliceOpInterface>(*ctx);
1057  ExtractOp::attachInterface<ExtractOpInterface>(*ctx);
1058  FromElementsOp::attachInterface<FromElementsOpInterface>(*ctx);
1059  GenerateOp::attachInterface<GenerateOpInterface>(*ctx);
1060  InsertOp::attachInterface<InsertOpInterface>(*ctx);
1061  InsertSliceOp::attachInterface<InsertSliceOpInterface>(*ctx);
1062  PadOp::attachInterface<PadOpInterface>(*ctx);
1063  ParallelInsertSliceOp::attachInterface<ParallelInsertSliceOpInterface>(
1064  *ctx);
1065  RankOp::attachInterface<RankOpInterface>(*ctx);
1066  ReshapeOp::attachInterface<ReshapeOpInterface>(*ctx);
1067  SplatOp::attachInterface<SplatOpInterface>(*ctx);
1068 
1069  // Load additional dialects of which ops may get created.
1070  ctx->loadDialect<arith::ArithDialect, linalg::LinalgDialect>();
1071  });
1072 
1073  // Bufferization requires SubsetInsertionOpInterface models. Make sure that
1074  // they are registered.
1076 }
static llvm::ManagedStatic< PassManagerOptions > options
Base type for affine expression.
Definition: AffineExpr.h:69
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:47
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:31
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:124
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
void addExtension(std::unique_ptr< DialectExtensionBase > extension)
Add the given extension to the registry.
This class provides support for representing a failure result, or a valid value of type T.
Definition: LogicalResult.h:78
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:63
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
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:350
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Definition: Builders.h:433
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:400
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:464
This class represents a single result from folding an operation.
Definition: OpDefinition.h:268
This class represents an operand of an operation.
Definition: Value.h:267
This is a value defined by a result of an operation.
Definition: Value.h:457
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
OpResult getOpResult(unsigned idx)
Definition: Operation.h:416
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
Definition: Operation.h:402
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
use_range getUses()
Returns a range of all uses, which is useful for iterating over all uses.
Definition: Operation.h:842
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:400
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:536
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:381
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:129
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.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:285
SmallVector< OpFoldResult > getMixedSizes(OpBuilder &builder, Location loc, Value value)
Return the dimensions of the given memref value.
Definition: MemRefOps.cpp:77
void registerSubsetOpInterfaceExternalModels(DialectRegistry &registry)
void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry)
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
std::optional< int64_t > getConstantIntValue(OpFoldResult ofr)
If ofr is a constant integer or an IntegerAttr, return the integer.
LogicalResult getStridesAndOffset(MemRefType t, SmallVectorImpl< int64_t > &strides, int64_t &offset)
Returns the strides of the MemRef if the layout map is in strided form.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
void bindSymbols(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to SymbolExpr at positions: [0 .
Definition: AffineExpr.h:363
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
Definition: LogicalResult.h:72
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
Options for BufferizableOpInterface-based bufferization.
Bufferizable ops that implement the DestinationStyleOpInterface can use this external model base clas...