[mlir][tensor][bufferize] Support memory_space for tensor.pad

[springerm]

Summary

This change adds memory space support to tensor.pad. (tensor.generate and tensor.from_elements do not support memory spaces yet.)

The memory space is inferred from the buffer of the source tensor.

Instead of lowering tensor.pad to tensor.generate + tensor.insert_slice, it is now lowered to bufferization.alloc_tensor (with the correct memory space) + scf.parallel + tensor.insert_slice.

Memory space support for the remaining two tensor ops is left for a later point, as this requires some more design discussions.

Depends On D136767

[springerm]

update

[KoolJBlack]

Awesome, thanks for setting this up.

[nicolasvasilache]

This triggers all sorts of red flags for me (lowering to loops directly during bufferization, increasing the number of uses of ToMemref/ToTensor, having to choose a loop type) and I don't see this as a good path forward.

[springerm]

address comments

[nicolasvasilache]

Awesome, thanks for improving this!

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

createStores(rewriter, loc, /*dim=*/0, buffer, shape, constants, elementIt,
                  indices);
 
     replaceOpWithBufferizedValues(rewriter, op, buffer);
 
     return success();
   }
 };
 
+/// Bufferize a tensor.generate like op body (one index-typed bbArg per dim).
+/// It is lowered to an scf.parallel loop and scalar stores inside of the loop.
+static void bufferizeGenerateLikeOpBody(RewriterBase &rewriter, Location loc,
+                                        Value buffer, ValueRange dynamicSizes,
+                                        Region &body) {
+  assert(body.hasOneBlock() && "expected body with single block");
+  auto memrefType = buffer.getType().cast<MemRefType>();
+  assert(body.getNumArguments() == memrefType.getRank() && "rank mismatch");
+  assert(llvm::all_of(
+             body.getArguments(),
+             [](BlockArgument bbArg) { return bbArg.getType().isIndex(); }) &&
+         "expected only index-typed bbArgs");
+  OpBuilder::InsertionGuard g(rewriter);
+
+  // Collect loop bounds.
+  int64_t rank = memrefType.getRank();
+  Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+  Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
+  SmallVector<Value, 4> lowerBounds(rank, zero);
+  SmallVector<Value, 4> steps(rank, one);
+  SmallVector<Value, 4> upperBounds;
+  int nextDynamicIndex = 0;
+  for (int i = 0; i < rank; i++) {
+    Value upperBound = memrefType.isDynamicDim(i)
+                           ? dynamicSizes[nextDynamicIndex++]
+                           : rewriter.create<arith::ConstantIndexOp>(
+                                 loc, memrefType.getDimSize(i));
+    upperBounds.push_back(upperBound);
+  }
+
+  // Generate tensor elements with a parallel loop that stores into
+  // each element of the resulting memref. We use mergeBlockBefore to "move"
+  // this op's body into the scf.parallel's body.
+  auto parallel =
+      rewriter.create<scf::ParallelOp>(loc, lowerBounds, upperBounds, steps);
+  Block *parallelBody = parallel.getBody();
+  rewriter.mergeBlockBefore(&body.front(), parallelBody->getTerminator(),
+                            parallelBody->getArguments());
+  // Replace the yielding op with a store op. The scf.parallel's builder
+  // already populated an scf.yield at the end, so we don't need to worry
+  // about creating that.
+  Operation *yieldOp = parallelBody->getTerminator()->getPrevNode();
+  Value yieldedValue = yieldOp->getOperands()[0];
+  assert(yieldedValue.getType() == memrefType.getElementType() &&
+         "type mismatch");
+  rewriter.setInsertionPointAfter(yieldOp);
+  rewriter.replaceOpWithNewOp<memref::StoreOp>(yieldOp, yieldedValue, buffer,
+                                               parallelBody->getArguments());
+}
+
 /// Bufferization of tensor.generate.
 struct GenerateOpInterface
     : public BufferizableOpInterface::ExternalModel<GenerateOpInterface,
                                                     tensor::GenerateOp> {
 
   bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
     return true;
   }
FailureOr<Value> tensorAlloc =
                                      /*copy=*/false);
     if (failed(tensorAlloc))
       return failure();
     auto memrefType =
         MemRefType::get(tensorType.getShape(), tensorType.getElementType());
     Value buffer = rewriter.create<bufferization::ToMemrefOp>(
         op->getLoc(), memrefType, *tensorAlloc);
 
-    // Collect loop bounds.
-    int64_t rank = memrefType.getRank();
-    Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
-    Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
-    SmallVector<Value, 4> lowerBounds(rank, zero);
-    SmallVector<Value, 4> steps(rank, one);
-    SmallVector<Value, 4> upperBounds;
-    int nextDynamicIndex = 0;
-    for (int i = 0; i < rank; i++) {
-      Value upperBound =
-          memrefType.isDynamicDim(i)
-              ? generateOp.getDynamicExtents()[nextDynamicIndex++]
-              : rewriter.create<arith::ConstantIndexOp>(
-                    loc, memrefType.getDimSize(i));
-      upperBounds.push_back(upperBound);
-    }
-
-    // Generate tensor elements with a parallel loop that stores into
-    // each element of the resulting memref. We use mergeBlockBefore to "move"
-    // this op's body into the scf.parallel's body.
-    auto parallel =
-        rewriter.create<scf::ParallelOp>(loc, lowerBounds, upperBounds, steps);
-    Block *parallelBody = parallel.getBody();
-    rewriter.mergeBlockBefore(&generateOp.getBody().front(),
-                              parallelBody->getTerminator(),
-                              parallelBody->getArguments());
-    // Replace the inlined yield op with a store op. The scf.parallel's builder
-    // already populated an scf.yield at the end, so we don't need to worry
-    // about creating that.
-    Operation *elementYield = parallelBody->getTerminator()->getPrevNode();
-    rewriter.setInsertionPointAfter(elementYield);
-    rewriter.replaceOpWithNewOp<memref::StoreOp>(
-        elementYield, elementYield->getOperands()[0], buffer,
-        parallelBody->getArguments());
-
+    bufferizeGenerateLikeOpBody(rewriter, loc, buffer,
+                                generateOp.getDynamicExtents(),
+                                generateOp.getBody());
     replaceOpWithBufferizedValues(rewriter, op, buffer);
 
     return success();
   }
 };
 
 /// Bufferization of tensor.insert. Replace with memref.store.
 struct InsertOpInterface
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
     return false;
   }
 
   SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
+  FailureOr<BaseMemRefType>
+  getBufferType(Operation *op, Value value, const BufferizationOptions &options,
+                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+    // Infer memory space from the source tensor.
+    auto padOp = cast<tensor::PadOp>(op);
+    auto maybeSrcBufferType =
+        bufferization::getBufferType(padOp.getSource(), options, fixedTypes);
+    if (failed(maybeSrcBufferType))
+      return failure();
+    MemRefLayoutAttrInterface layout;
+    return MemRefType::get(padOp.getResultType().getShape(),
+                           padOp.getResultType().getElementType(), layout,
+                           maybeSrcBufferType->getMemorySpace());
+  }
+
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto padOp = cast<tensor::PadOp>(op);
     Location loc = padOp.getLoc();
     RankedTensorType resultType = padOp.getResultType();
     RankedTensorType srcType = padOp.getSourceType();
+    auto resultBufferType =
+        bufferization::getBufferType(padOp.getResult(), options);
+    if (failed(resultBufferType))
+      return failure();
 
     auto toValue = [&](OpFoldResult ofr) {
       if (ofr.is<Value>())
         return ofr.get<Value>();
       return rewriter
           .create<arith::ConstantIndexOp>(loc, *getConstantIntValue(ofr))
           .getResult();
     };
for (int64_t i = 0; i < resultType.getRank(); ++i) {
       AffineExpr s0, s1, s2;
       bindSymbols(op->getContext(), s0, s1, s2);
       AffineExpr sumExpr = s0 + s1 + s2;
       Value sum = rewriter.create<AffineApplyOp>(
           loc, sumExpr, ValueRange{srcDim, lowPad, highPad});
       dynamicSizes.push_back(sum);
     }
 
-    // Create tensor::GenerateOp.
-    auto generateOp =
-        rewriter.create<tensor::GenerateOp>(loc, resultType, dynamicSizes);
-    // Move over "escape" attribute if present.
-    if (padOp->hasAttr(BufferizationDialect::kEscapeAttrName))
-      generateOp->setAttr(
-          BufferizationDialect::kEscapeAttrName,
-          padOp->getAttr(BufferizationDialect::kEscapeAttrName));
-    // TODO: Memory space
-    rewriter.inlineRegionBefore(padOp.getRegion(), generateOp.getBody(),
-                                generateOp.getBody().begin());
+    // Should the buffer be deallocated?
+    bool dealloc =
+        shouldDeallocateOpResult(padOp.getResult().cast<OpResult>(), options);
+    // Allocate a buffer for the padded result.
+    FailureOr<Value> tensorAlloc =
+        allocateTensorForShapedValue(rewriter, loc, padOp.getResult(),
+                                     /*escape=*/!dealloc, options,
+                                     /*copy=*/false);
+    if (failed(tensorAlloc))
+      return failure();
+    Value buffer = rewriter.create<bufferization::ToMemrefOp>(
+        op->getLoc(), *resultBufferType, *tensorAlloc);
+
+    // tensor::PadOp is like tensor::GenerateOp: The only difference is that
+    // only a part of the generated tensor is needed. For simplicity, we reuse
+    // the same functionality here.
+    bufferizeGenerateLikeOpBody(rewriter, loc, buffer, dynamicSizes,
+                                padOp.getBodyRegion());
 
     // Create tensor::InsertSliceOp.
+    Value filledBuffer =
+        rewriter.create<bufferization::ToTensorOp>(loc, buffer);
     SmallVector<OpFoldResult> sliceSizes =
         getMixedSizes(rewriter, loc, padOp.getSource());
     SmallVector<OpFoldResult> sliceStrides(srcType.getRank(),
                                            rewriter.getIndexAttr(1));
     rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
-        padOp, padOp.getSource(), generateOp.getResult(),
+        padOp, padOp.getSource(), filledBuffer,
         /*offsets=*/padOp.getMixedLowPad(), sliceSizes, sliceStrides);
 
     return success();
   }
 };
 
 /// Bufferization of tensor.rank. Replace with memref.rank.
 struct RankOpInterface

mlir/test/Dialect/Tensor/one-shot-bufferize.mlir

func.func @insert_equivalent_tensor(%t: tensor<10xf32>) -> tensor<10xf32> {
   // CHECK-NOT: memref.alloc
   %cst = arith.constant 4.200000e+01 : f32
   // CHECK: linalg.fill
   %0 = linalg.fill ins(%cst : f32) outs(%t : tensor<10xf32>) -> tensor<10xf32>
   // CHECK-NOT: memref.copy
   %1 = tensor.insert_slice %0 into %t[0][10][1] : tensor<10xf32> into tensor<10xf32>
   return %1 : tensor<10xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @pad_memory_space(
+//  CHECK-SAME:     %[[t:.*]]: memref<?xf32, strided<[?], offset: ?>>
+func.func @pad_memory_space(%t: tensor<?xf32>, %h1: index, %f: f32, %pos: index) -> f32
+{
+  // CHECK: %[[alloc_tensor:.*]] = memref.alloc{{.*}} : memref<?xf32, 3>
+  // CHECK: memref.copy %[[t]], %[[alloc_tensor]]
+  %0 = bufferization.alloc_tensor() copy(%t)
+      {memory_space = 3 : ui64} : tensor<?xf32>
+  // CHECK: %[[padded_alloc:.*]] = memref.alloc() {{.*}} : memref<15xf32, 3>
+  // CHECK: scf.parallel
+  // CHECK:  memref.store {{.*}} : memref<15xf32, 3>
+  // CHECK: %[[subview:.*]] = memref.subview {{.*}} : memref<15xf32, 3> to memref<?xf32, strided<[1], offset: 2>, 3>
+  // CHECK: memref.copy %[[alloc_tensor]], %[[subview]]
+  %1 = tensor.pad %0 low[2] high[%h1] {
+  ^bb0(%arg0: index):
+    tensor.yield %f : f32
+  } : tensor<?xf32> to tensor<15xf32>
+  // CHECK: memref.load {{.*}} : memref<15xf32, 3>
+  %2 = tensor.extract %1[%pos] : tensor<15xf32>
+  // CHECK-DAG: memref.dealloc %[[alloc_tensor]]
+  // CHECK-DAG: memref.dealloc %[[padded_alloc]]
+  return %2 : f32
+}