MLIR  18.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  replaceOpWithNewBufferizedOp<memref::ExpandShapeOp>(
342  rewriter, op, tensorResultType.getShape(), *buffer,
343  expandShapeOp.getReassociationIndices());
344  return success();
345  }
346 };
347 
348 /// Bufferization of tensor.extract_slice. Replace with memref.subview.
349 struct ExtractSliceOpInterface
350  : public BufferizableOpInterface::ExternalModel<ExtractSliceOpInterface,
351  tensor::ExtractSliceOp> {
352  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
353  const AnalysisState &state) const {
354  return false;
355  }
356 
357  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
358  const AnalysisState &state) const {
359  return false;
360  }
361 
362  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
363  const AnalysisState &state) const {
364  return {{op->getOpResult(0), BufferRelation::Unknown}};
365  }
366 
367  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
368  const BufferizationOptions &options) const {
369  auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
370  SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
371  SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
372  SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
373  Location loc = extractSliceOp.getLoc();
374 
375  // Get source buffer.
376  FailureOr<Value> srcMemref =
377  getBuffer(rewriter, extractSliceOp.getSource(), options);
378  if (failed(srcMemref))
379  return failure();
380 
381  // Take a subview of the source buffer.
382  auto resultMemrefType =
383  bufferization::getBufferType(extractSliceOp.getResult(), options);
384  if (failed(resultMemrefType))
385  return failure();
386  Value subView = rewriter.create<memref::SubViewOp>(
387  loc, llvm::cast<MemRefType>(*resultMemrefType), *srcMemref, mixedOffsets,
388  mixedSizes, mixedStrides);
389 
390  replaceOpWithBufferizedValues(rewriter, op, subView);
391  return success();
392  }
393 
396  SmallVector<Value> &invocationStack) const {
397  auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
398  assert(value == extractSliceOp.getResult() && "invalid value");
399  auto srcMemrefType = bufferization::getBufferType(
400  extractSliceOp.getSource(), options, invocationStack);
401  if (failed(srcMemrefType))
402  return failure();
403  SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
404  SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
405  SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
406  return cast<BaseMemRefType>(memref::SubViewOp::inferRankReducedResultType(
407  extractSliceOp.getType().getShape(), llvm::cast<MemRefType>(*srcMemrefType),
408  mixedOffsets, mixedSizes, mixedStrides));
409  }
410 };
411 
412 /// Bufferization of tensor.extract. Replace with memref.load.
413 struct ExtractOpInterface
414  : public BufferizableOpInterface::ExternalModel<ExtractOpInterface,
415  tensor::ExtractOp> {
416  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
417  const AnalysisState &state) const {
418  return true;
419  }
420 
421  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
422  const AnalysisState &state) const {
423  return false;
424  }
425 
426  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
427  const AnalysisState &state) const {
428  return {};
429  }
430 
431  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
432  const BufferizationOptions &options) const {
433  auto extractOp = cast<tensor::ExtractOp>(op);
434  FailureOr<Value> srcMemref =
435  getBuffer(rewriter, extractOp.getTensor(), options);
436  if (failed(srcMemref))
437  return failure();
438  replaceOpWithNewBufferizedOp<memref::LoadOp>(rewriter, op, *srcMemref,
439  extractOp.getIndices());
440  return success();
441  }
442 };
443 
444 // Implements backtracking to traverse indices of the output buffer while
445 // iterating over op.elements().
446 static void createStores(RewriterBase &rewriter, Location loc, int dim,
447  Value buffer, ArrayRef<int64_t> shape,
448  ArrayRef<Value> constants,
449  OperandRange::iterator &elementIt,
450  SmallVectorImpl<Value> &indices) {
451  if (dim == static_cast<int>(shape.size()) - 1) {
452  for (int i = 0; i < shape.back(); ++i) {
453  indices.back() = constants[i];
454  rewriter.create<memref::StoreOp>(loc, *elementIt, buffer, indices);
455  ++elementIt;
456  }
457  return;
458  }
459  for (int i = 0; i < shape[dim]; ++i) {
460  indices[dim] = constants[i];
461  createStores(rewriter, loc, dim + 1, buffer, shape, constants, elementIt,
462  indices);
463  }
464 }
465 
466 /// Bufferization of tensor.from_elements.
467 struct FromElementsOpInterface
468  : public BufferizableOpInterface::ExternalModel<FromElementsOpInterface,
469  tensor::FromElementsOp> {
470 
471  bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
472 
473  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
474  const BufferizationOptions &options) const {
475  auto fromElementsOp = cast<tensor::FromElementsOp>(op);
476 
477  // TODO: Implement memory space for this op.
478  if (options.defaultMemorySpace != Attribute())
479  return op->emitError("memory space not implemented yet");
480 
481  // Allocate a buffer for the result.
482  Location loc = op->getLoc();
483  auto tensorType = cast<RankedTensorType>(fromElementsOp.getType());
484  auto shape = tensorType.getShape();
485  // TODO: Create alloc_tensor ops during TensorCopyInsertion.
487  rewriter, loc, fromElementsOp.getResult(), options,
488  /*copy=*/false);
489  if (failed(tensorAlloc))
490  return failure();
491  auto memrefType =
492  MemRefType::get(tensorType.getShape(), tensorType.getElementType());
493  Value buffer = rewriter.create<bufferization::ToMemrefOp>(
494  op->getLoc(), memrefType, *tensorAlloc);
495 
496  // Case: tensor<0xelem_type>.
497  if (fromElementsOp.getElements().empty()) {
498  replaceOpWithBufferizedValues(rewriter, op, buffer);
499  return success();
500  }
501 
502  // Case: tensor<elem_type>.
503  if (shape.empty()) {
504  rewriter.create<memref::StoreOp>(
505  loc, fromElementsOp.getElements().front(), buffer);
506  replaceOpWithBufferizedValues(rewriter, op, buffer);
507  return success();
508  }
509 
510  // Create constants for the range of possible indices [0, max{shape_i}).
511  auto maxDim = *std::max_element(shape.begin(), shape.end());
512  SmallVector<Value, 2> constants;
513  constants.reserve(maxDim);
514  for (int i = 0; i < maxDim; ++i)
515  constants.push_back(rewriter.create<arith::ConstantIndexOp>(loc, i));
516 
517  // Traverse all `elements` and create `memref.store` ops.
518  auto elementIt = fromElementsOp.getElements().begin();
519  SmallVector<Value, 2> indices(tensorType.getRank(), constants[0]);
520  createStores(rewriter, loc, /*dim=*/0, buffer, shape, constants, elementIt,
521  indices);
522 
523  replaceOpWithBufferizedValues(rewriter, op, buffer);
524 
525  return success();
526  }
527 };
528 
529 /// Lower the body of a tensor.generate like op (one index-typed bbArg per dim).
530 /// Such ops are lowered to linalg.map with the given tensor as a destination.
531 ///
532 /// Example:
533 /// ```
534 /// %r = tensor.generate %x, %y {
535 /// ^bb0(%arg0: index, %arg1: index):
536 /// %0 = "some_op"(%arg0, %arg1) : (index, index) -> (index)
537 /// tensor.yield %0 : index
538 /// } : tensor<?x?xindex>
539 /// ```
540 ///
541 /// Is lowered to:
542 /// ```
543 /// linalg.map ins() outs(%dest) {
544 /// %d0 = linalg.index 0 : index
545 /// %d1 = linalg.index 1 : index
546 /// %0 = "some_op"(%d0, %d1) : (index, index) -> (index)
547 /// linalg.yield %0 : index
548 /// }
549 /// ```
550 static Value lowerGenerateLikeOpBody(RewriterBase &rewriter, Location loc,
551  Value tensorDestination,
552  ValueRange dynamicSizes,
553  Region &generateBody) {
554  assert(generateBody.hasOneBlock() && "expected body with single block");
555  auto tensorType = cast<RankedTensorType>(tensorDestination.getType());
556  assert(generateBody.getNumArguments() == tensorType.getRank() &&
557  "rank mismatch");
558 
559  // Create linalg::MapOp.
560  OpBuilder::InsertionGuard g(rewriter);
561  auto linalgOp =
562  rewriter.create<linalg::MapOp>(loc, tensorType, /*inputs=*/ValueRange(),
563  /*init=*/tensorDestination);
564  Block &linalgBody = linalgOp.getMapper().emplaceBlock();
565 
566  // Create linalg::IndexOps.
567  rewriter.setInsertionPointToStart(&linalgBody);
568  SmallVector<Value> indices;
569  for (int64_t dim = 0; dim < tensorType.getRank(); ++dim)
570  indices.push_back(rewriter.create<linalg::IndexOp>(loc, dim));
571 
572  // Move over body.
573  rewriter.mergeBlocks(&generateBody.front(), &linalgBody, indices);
574  auto yieldOp = cast<tensor::YieldOp>(linalgBody.getTerminator());
575  rewriter.replaceOpWithNewOp<linalg::YieldOp>(yieldOp, yieldOp.getValue());
576 
577  return linalgOp.getResult()[0];
578 }
579 
580 /// Bufferization of tensor.generate.
581 struct GenerateOpInterface
582  : public BufferizableOpInterface::ExternalModel<GenerateOpInterface,
583  tensor::GenerateOp> {
584 
585  bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
586 
587  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
588  const BufferizationOptions &options) const {
589  auto generateOp = cast<tensor::GenerateOp>(op);
590 
591  // TODO: Implement memory space for this op.
592  if (options.defaultMemorySpace != Attribute())
593  return op->emitError("memory space not implemented yet");
594 
595  // Allocate memory.
596  Location loc = op->getLoc();
598  rewriter, loc, generateOp.getResult(), options,
599  /*copy=*/false);
600  if (failed(tensorAlloc))
601  return failure();
602 
603  Value result = lowerGenerateLikeOpBody(rewriter, loc, *tensorAlloc,
604  generateOp.getDynamicExtents(),
605  generateOp.getBody());
606  rewriter.replaceOp(generateOp, result);
607 
608  return success();
609  }
610 };
611 
612 /// Bufferization of tensor.insert. Replace with memref.store.
613 ///
614 /// Note: DstBufferizableOpInterfaceExternalModel provides many default method
615 /// implementations for DestinationStyle ops.
616 struct InsertOpInterface
617  : public DstBufferizableOpInterfaceExternalModel<InsertOpInterface,
618  tensor::InsertOp> {
619  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
620  const BufferizationOptions &options) const {
621  auto insertOp = cast<tensor::InsertOp>(op);
622  FailureOr<Value> destMemref =
623  getBuffer(rewriter, insertOp.getDest(), options);
624  if (failed(destMemref))
625  return failure();
626  rewriter.create<memref::StoreOp>(insertOp.getLoc(), insertOp.getScalar(),
627  *destMemref, insertOp.getIndices());
628  replaceOpWithBufferizedValues(rewriter, op, *destMemref);
629  return success();
630  }
631 };
632 
633 /// Bufferization of tensor.insert_slice. Replace with a memory copy. Under
634 /// certain circumstances, this op can also be a no-op.
635 ///
636 /// Note: DstBufferizableOpInterfaceExternalModel provides many default method
637 /// implementations for DestinationStyle ops.
638 struct InsertSliceOpInterface
639  : public DstBufferizableOpInterfaceExternalModel<InsertSliceOpInterface,
640  tensor::InsertSliceOp> {
641  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
642  const AnalysisState &state) const {
643  auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
644  RankedTensorType destType = insertSliceOp.getDestType();
645 
646  // The source is always read.
647  if (opOperand == insertSliceOp.getSourceMutable())
648  return true;
649 
650  // For the destination, it depends...
651  assert(opOperand == insertSliceOp.getDestMutable() && "expected dest");
652 
653  // Dest is not read if it is entirely overwritten. E.g.:
654  // tensor.insert_slice %a into %t[0][10][1] : ... into tensor<10xf32>
655  bool allOffsetsZero =
656  llvm::all_of(insertSliceOp.getMixedOffsets(), [](OpFoldResult ofr) {
657  return isConstantIntValue(ofr, 0);
658  });
659  bool sizesMatchDestSizes = llvm::all_of(
660  llvm::enumerate(insertSliceOp.getMixedSizes()), [&](const auto &it) {
661  return getConstantIntValue(it.value()) ==
662  destType.getDimSize(it.index());
663  });
664  bool allStridesOne =
665  llvm::all_of(insertSliceOp.getMixedStrides(), [](OpFoldResult ofr) {
666  return isConstantIntValue(ofr, 1);
667  });
668  return !(allOffsetsZero && sizesMatchDestSizes && allStridesOne);
669  }
670 
671  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
672  const BufferizationOptions &options) const {
673  // insert_slice ops arise from tiling and bufferizing them out-of-place is
674  // generally a deal breaker. When used with loops, this ends up cloning the
675  // whole tensor on every single iteration and is a symptom of a
676  // catastrophically bad scheduling decision.
677  // TODO: be very loud about it or even consider failing the pass.
678  auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
679  SmallVector<OpFoldResult> mixedOffsets = insertSliceOp.getMixedOffsets();
680  SmallVector<OpFoldResult> mixedSizes = insertSliceOp.getMixedSizes();
681  SmallVector<OpFoldResult> mixedStrides = insertSliceOp.getMixedStrides();
682  Location loc = insertSliceOp.getLoc();
683 
684  // Get destination buffer.
685  FailureOr<Value> dstMemref =
686  getBuffer(rewriter, insertSliceOp.getDest(), options);
687  if (failed(dstMemref))
688  return failure();
689 
690  // Take a subview of the destination buffer.
691  auto dstMemrefType = cast<MemRefType>(dstMemref->getType());
692  auto subviewMemRefType =
693  cast<MemRefType>(memref::SubViewOp::inferRankReducedResultType(
694  insertSliceOp.getSourceType().getShape(), dstMemrefType,
695  mixedOffsets, mixedSizes, mixedStrides));
696  Value subView = rewriter.create<memref::SubViewOp>(
697  loc, subviewMemRefType, *dstMemref, mixedOffsets, mixedSizes,
698  mixedStrides);
699 
700  // Copy tensor. If this tensor.insert_slice has a matching
701  // tensor.extract_slice, the copy operation will eventually fold away.
702  FailureOr<Value> srcMemref =
703  getBuffer(rewriter, insertSliceOp.getSource(), options);
704  if (failed(srcMemref))
705  return failure();
706  if (failed(options.createMemCpy(rewriter, loc, *srcMemref, subView)))
707  return failure();
708 
709  replaceOpWithBufferizedValues(rewriter, op, *dstMemref);
710  return success();
711  }
712 };
713 
714 /// Bufferization of tensor.pad. Replace with bufferization.alloc_tensor +
715 /// linalg.map + insert_slice.
716 /// For best performance, vectorize before bufferization (better performance in
717 /// case of padding with a constant).
718 struct PadOpInterface
719  : public BufferizableOpInterface::ExternalModel<PadOpInterface,
720  tensor::PadOp> {
721  bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
722 
723  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
724  const AnalysisState &state) const {
725  return true;
726  }
727 
728  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
729  const AnalysisState &state) const {
730  return false;
731  }
732 
733  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
734  const AnalysisState &state) const {
735  return {};
736  }
737 
740  SmallVector<Value> &invocationStack) const {
741  // Infer memory space from the source tensor.
742  auto padOp = cast<tensor::PadOp>(op);
743  auto maybeSrcBufferType = bufferization::getBufferType(
744  padOp.getSource(), options, invocationStack);
745  if (failed(maybeSrcBufferType))
746  return failure();
747  MemRefLayoutAttrInterface layout;
748  return MemRefType::get(padOp.getResultType().getShape(),
749  padOp.getResultType().getElementType(), layout,
750  maybeSrcBufferType->getMemorySpace());
751  }
752 
753  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
754  const BufferizationOptions &options) const {
755  auto padOp = cast<tensor::PadOp>(op);
756  Location loc = padOp.getLoc();
757  RankedTensorType resultType = padOp.getResultType();
758  RankedTensorType srcType = padOp.getSourceType();
759 
760  auto toValue = [&](OpFoldResult ofr) {
761  if (ofr.is<Value>())
762  return ofr.get<Value>();
763  return rewriter
764  .create<arith::ConstantIndexOp>(loc, *getConstantIntValue(ofr))
765  .getResult();
766  };
767 
768  // Compute dynamic result dimensions.
769  SmallVector<OpFoldResult> mixedLowPad = padOp.getMixedLowPad();
770  SmallVector<OpFoldResult> mixedHighPad = padOp.getMixedHighPad();
771  SmallVector<Value> dynamicSizes;
772  for (int64_t i = 0; i < resultType.getRank(); ++i) {
773  if (!resultType.isDynamicDim(i))
774  continue;
775  Value srcDim = rewriter.create<tensor::DimOp>(loc, padOp.getSource(), i);
776  Value lowPad = toValue(mixedLowPad[i]);
777  Value highPad = toValue(mixedHighPad[i]);
778  AffineExpr s0, s1, s2;
779  bindSymbols(op->getContext(), s0, s1, s2);
780  AffineExpr sumExpr = s0 + s1 + s2;
781  Value sum = rewriter.create<affine::AffineApplyOp>(
782  loc, sumExpr, ValueRange{srcDim, lowPad, highPad});
783  dynamicSizes.push_back(sum);
784  }
785 
786  // Allocate a buffer for the padded result.
787  FailureOr<Value> tensorAlloc =
788  allocateTensorForShapedValue(rewriter, loc, padOp.getResult(), options,
789  /*copy=*/false);
790  if (failed(tensorAlloc))
791  return failure();
792 
793  // tensor::PadOp is like tensor::GenerateOp: The only difference is that
794  // only a part of the generated tensor is needed. For simplicity, we reuse
795  // the same functionality here.
796  Value filledBuffer = lowerGenerateLikeOpBody(
797  rewriter, loc, *tensorAlloc, dynamicSizes, padOp.getBodyRegion());
798 
799  // Create tensor::InsertSliceOp.
800  SmallVector<OpFoldResult> sliceSizes =
801  getMixedSizes(rewriter, loc, padOp.getSource());
802  SmallVector<OpFoldResult> sliceStrides(srcType.getRank(),
803  rewriter.getIndexAttr(1));
804  rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
805  padOp, padOp.getSource(), filledBuffer,
806  /*offsets=*/padOp.getMixedLowPad(), sliceSizes, sliceStrides);
807 
808  return success();
809  }
810 };
811 
812 /// Bufferization of tensor.rank. Replace with memref.rank.
813 struct RankOpInterface
814  : public BufferizableOpInterface::ExternalModel<RankOpInterface,
815  tensor::RankOp> {
816  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
817  const AnalysisState &state) const {
818  // The op reads the tensor's metadata but not its contents.
819  return false;
820  }
821 
822  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
823  const AnalysisState &state) const {
824  return false;
825  }
826 
827  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
828  const AnalysisState &state) const {
829  return {};
830  }
831 
832  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
833  const BufferizationOptions &options) const {
834  auto rankOp = cast<tensor::RankOp>(op);
835  FailureOr<Value> v = getBuffer(rewriter, rankOp.getTensor(), options);
836  if (failed(v))
837  return failure();
838  replaceOpWithNewBufferizedOp<memref::RankOp>(rewriter, op, rankOp.getType(),
839  *v);
840  return success();
841  }
842 };
843 
844 /// Bufferization of tensor.reshape. Replace with memref.reshape.
845 struct ReshapeOpInterface
846  : public BufferizableOpInterface::ExternalModel<ReshapeOpInterface,
847  tensor::ReshapeOp> {
848  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
849  const AnalysisState &state) const {
850  // Depending on the layout map, the source buffer may have to be copied.
851  auto reshapeOp = cast<tensor::ReshapeOp>(op);
852  return opOperand == reshapeOp.getShapeMutable();
853  }
854 
855  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
856  const AnalysisState &state) const {
857  return false;
858  }
859 
860  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
861  const AnalysisState &state) const {
862  return {{op->getOpResult(0), BufferRelation::Equivalent}};
863  }
864 
865  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
866  const BufferizationOptions &options) const {
867  auto reshapeOp = cast<tensor::ReshapeOp>(op);
868  FailureOr<Value> srcBuffer =
869  getBuffer(rewriter, reshapeOp.getSource(), options);
870  FailureOr<Value> shapeBuffer =
871  getBuffer(rewriter, reshapeOp.getShape(), options);
872  if (failed(srcBuffer) || failed(shapeBuffer))
873  return failure();
874  auto maybeResultMemRefType =
875  bufferization::getBufferType(reshapeOp.getResult(), options);
876  if (failed(maybeResultMemRefType))
877  return failure();
878 
879  // memref.reshape requires the source buffer to have an identity layout.
880  // If the source memref does not have an identity layout, copy the source
881  // into a new buffer with an identity layout.
882  auto srcType = llvm::dyn_cast<MemRefType>(srcBuffer->getType());
883  if (srcType && !srcType.getLayout().isIdentity()) {
885  rewriter, op->getLoc(), reshapeOp.getSource(), options);
886  if (failed(tensorAlloc))
887  return failure();
888  auto memrefType = MemRefType::get(
889  srcType.getShape(), srcType.getElementType(), AffineMap(),
890  cast<BaseMemRefType>(srcBuffer->getType()).getMemorySpace());
891  srcBuffer = rewriter
892  .create<bufferization::ToMemrefOp>(
893  op->getLoc(), memrefType, *tensorAlloc)
894  .getResult();
895  }
896 
897  replaceOpWithNewBufferizedOp<memref::ReshapeOp>(
898  rewriter, op, maybeResultMemRefType.value(), *srcBuffer, *shapeBuffer);
899  return success();
900  }
901 
904  SmallVector<Value> &invocationStack) const {
905  auto reshapeOp = cast<tensor::ReshapeOp>(op);
906  assert(value == reshapeOp.getResult() && "unexpected value provided");
907  auto maybeSourceBufferType = bufferization::getBufferType(
908  reshapeOp.getSource(), options, invocationStack);
909  if (failed(maybeSourceBufferType))
910  return failure();
912  reshapeOp.getResult().getType(),
913  cast<BaseMemRefType>(maybeSourceBufferType.value()).getMemorySpace());
914  }
915 };
916 
917 /// Analysis of ParallelInsertSliceOp.
918 struct ParallelInsertSliceOpInterface
919  : public BufferizableOpInterface::ExternalModel<
920  ParallelInsertSliceOpInterface, ParallelInsertSliceOp> {
921  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
922  const AnalysisState &state) const {
923  return {};
924  }
925 
926  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
927  const AnalysisState &state) const {
928  return true;
929  }
930 
931  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
932  const AnalysisState &state) const {
933  auto parallelInsertSliceOp = cast<ParallelInsertSliceOp>(op);
934  return opOperand == parallelInsertSliceOp.getDestMutable();
935  }
936 
937  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
938  const BufferizationOptions &options) const {
939  OpBuilder::InsertionGuard g(rewriter);
940  auto parallelInsertSliceOp = cast<ParallelInsertSliceOp>(op);
941  ParallelCombiningOpInterface parallelCombiningParent =
942  parallelInsertSliceOp.getParallelCombiningParent();
943 
944  // Bufferize the op outside of the parallel combining terminator.
945  rewriter.setInsertionPoint(parallelCombiningParent);
946 
947  // Get source and destination buffers.
948  FailureOr<Value> destBuffer =
949  getBuffer(rewriter, parallelInsertSliceOp.getDest(), options);
950  if (failed(destBuffer))
951  return failure();
952  FailureOr<Value> srcBuffer =
953  getBuffer(rewriter, parallelInsertSliceOp.getSource(), options);
954  if (failed(srcBuffer))
955  return failure();
956 
957  // Take a subview of the destination buffer.
958  auto destBufferType = cast<MemRefType>(destBuffer->getType());
959  auto subviewMemRefType =
960  cast<MemRefType>(memref::SubViewOp::inferRankReducedResultType(
961  parallelInsertSliceOp.getSourceType().getShape(), destBufferType,
962  parallelInsertSliceOp.getMixedOffsets(),
963  parallelInsertSliceOp.getMixedSizes(),
964  parallelInsertSliceOp.getMixedStrides()));
965  Value subview = rewriter.create<memref::SubViewOp>(
966  parallelInsertSliceOp.getLoc(), subviewMemRefType, *destBuffer,
967  parallelInsertSliceOp.getMixedOffsets(),
968  parallelInsertSliceOp.getMixedSizes(),
969  parallelInsertSliceOp.getMixedStrides());
970 
971  // This memcpy will fold away if everything bufferizes in-place.
972  if (failed(options.createMemCpy(rewriter, parallelInsertSliceOp.getLoc(),
973  *srcBuffer, subview)))
974  return failure();
975 
976  // In case the source was allocated in the same block, make sure that the
977  // deallocation op (if any) appears after the memcpy. By default, deallocs
978  // are placed before the terminator, but this does not work for ForallOp
979  // because the terminator does more than just yielding a value.
980  //
981  // Note: This is not a problem for the destination buffer because these are
982  // assumed to always bufferize in-place.
983  for (Operation *user : srcBuffer->getUsers()) {
984  if (hasEffect<MemoryEffects::Free>(user)) {
985  if (user->getBlock() == parallelCombiningParent->getBlock())
986  user->moveBefore(user->getBlock()->getTerminator());
987  break;
988  }
989  }
990 
991  // Delete the op.
992  rewriter.eraseOp(op);
993  return success();
994  }
995 };
996 
997 /// Bufferization of tensor.splat. Bufferizes to a new allocation that is filled
998 /// with a linalg.map. Similar to tensor.generate.
999 struct SplatOpInterface
1000  : public BufferizableOpInterface::ExternalModel<SplatOpInterface,
1001  tensor::SplatOp> {
1002 
1003  bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
1004 
1005  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
1006  const BufferizationOptions &options) const {
1007  OpBuilder::InsertionGuard g(rewriter);
1008  auto splatOp = cast<tensor::SplatOp>(op);
1009 
1010  // TODO: Implement memory space for this op.
1011  if (options.defaultMemorySpace != Attribute())
1012  return op->emitError("memory space not implemented yet");
1013 
1014  // Allocate memory.
1015  Location loc = op->getLoc();
1017  rewriter, loc, splatOp.getResult(), options,
1018  /*copy=*/false);
1019  if (failed(tensorAlloc))
1020  return failure();
1021 
1022  // Create linalg::MapOp.
1023  auto tensorType = cast<RankedTensorType>(tensorAlloc->getType());
1024  auto linalgOp =
1025  rewriter.create<linalg::MapOp>(loc, tensorType, /*inputs=*/ValueRange(),
1026  /*init=*/*tensorAlloc);
1027  Block &linalgBody = linalgOp.getMapper().emplaceBlock();
1028 
1029  // Create linalg::IndexOps.
1030  rewriter.setInsertionPointToStart(&linalgBody);
1031  rewriter.create<linalg::YieldOp>(loc, splatOp.getInput());
1032  rewriter.replaceOp(splatOp, linalgOp.getResult()[0]);
1033 
1034  return success();
1035  }
1036 };
1037 
1038 } // namespace
1039 } // namespace tensor
1040 } // namespace mlir
1041 
1043  DialectRegistry &registry) {
1044  registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
1045  CastOp::attachInterface<CastOpInterface>(*ctx);
1046  CollapseShapeOp::attachInterface<CollapseShapeOpInterface>(*ctx);
1047  DimOp::attachInterface<DimOpInterface>(*ctx);
1048  EmptyOp::attachInterface<EmptyOpInterface>(*ctx);
1049  ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
1050  ExtractSliceOp::attachInterface<ExtractSliceOpInterface>(*ctx);
1051  ExtractOp::attachInterface<ExtractOpInterface>(*ctx);
1052  FromElementsOp::attachInterface<FromElementsOpInterface>(*ctx);
1053  GenerateOp::attachInterface<GenerateOpInterface>(*ctx);
1054  InsertOp::attachInterface<InsertOpInterface>(*ctx);
1055  InsertSliceOp::attachInterface<InsertSliceOpInterface>(*ctx);
1056  PadOp::attachInterface<PadOpInterface>(*ctx);
1057  ParallelInsertSliceOp::attachInterface<ParallelInsertSliceOpInterface>(
1058  *ctx);
1059  RankOp::attachInterface<RankOpInterface>(*ctx);
1060  ReshapeOp::attachInterface<ReshapeOpInterface>(*ctx);
1061  SplatOp::attachInterface<SplatOpInterface>(*ctx);
1062 
1063  // Load additional dialects of which ops may get created.
1064  ctx->loadDialect<arith::ArithDialect, linalg::LinalgDialect>();
1065  });
1066 
1067  // Bufferization requires SubsetInsertionOpInterface models. Make sure that
1068  // they are registered.
1070 }
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: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:30
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:333
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Definition: Builders.h:416
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:383
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:446
This class represents a single result from folding an operation.
Definition: OpDefinition.h:266
This class represents an operand of an operation.
Definition: Value.h:263
This is a value defined by a result of an operation.
Definition: Value.h:453
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:267
use_range getUses()
Returns a range of all uses, which is useful for iterating over all uses.
Definition: Operation.h:825
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:399
virtual void replaceOp(Operation *op, ValueRange newValues)
This method replaces the results of the operation with the specified list of values.
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.
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replaces the result op with a new op that is created without verification.
Definition: PatternMatch.h:539
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:378
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:125
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:114
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:348
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...