[mlir][tensor][bufferize] Support memory_space on all tensor dialect ops

[springerm]

Summary

This change adds memory space support during bufferization for tensor.from_elements, tensor.generate and tensor.pad.

Depends On D136063

[springerm]

This needs some more iterations. In the mean time, D136265 adds memory support for tensor.pad (but not the other two ops).

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

return rewriter
       .create<bufferization::ToMemrefOp>(value.getLoc(), *memrefType, value)
       .getResult();
 }
 
 FailureOr<BaseMemRefType> bufferization::detail::defaultGetBufferType(
     Value value, const BufferizationOptions &options,
     const DenseMap<Value, BaseMemRefType> &fixedTypes) {
   assert(value.getType().isa<TensorType>() && "expected tensor type");
+  auto tensorType = value.getType().cast<TensorType>();
 
   // No further analysis is possible for a block argument.
   if (value.isa<BlockArgument>())
     return bufferization::getMemRefType(value, options);
 
   // Value is an OpResult.
   Operation *op = getOwnerOfValue(value);
   auto opResult = value.cast<OpResult>();
   auto bufferizableOp = cast<BufferizableOpInterface>(op);
   AnalysisState state(options);
+
+  // Case 1: If the OpResult has an equivalent OpOperand, both OpResult and
+  // OpOperand bufferize to the exact same buffer type.
   auto aliasingOperands = bufferizableOp.getAliasingOpOperand(opResult, state);
   if (!aliasingOperands.empty() &&
       bufferizableOp.bufferRelation(opResult, state) ==
           BufferRelation::Equivalent) {
-    // If the OpResult has an equivalent OpOperand, both OpResult and
-    // OpOperand bufferize to the exact same buffer type.
     Value equivalentOperand = aliasingOperands.front()->get();
     return getBufferType(equivalentOperand, options, fixedTypes);
   }
 
+  // Case 2: If the OpResult bufferizes to a new allocation (and never aliases
+  // any OpOperand), the result buffer type has a static identity layout.
+  if (aliasingOperands.empty() &&
+      bufferizableOp.bufferizesToAllocation(opResult)) {
+    // Compute memory space of this allocation.
+    unsigned memorySpace;
+    if (auto maybeMemorySpace = getMemorySpaceAttr(opResult)) {
+      memorySpace = *maybeMemorySpace;
+    } else if (options.defaultMemorySpace.has_value()) {
+      memorySpace = *options.defaultMemorySpace;
+    } else {
+      return op->emitError("could not infer memory space");
+    }
+    return getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace);
+  }
+
+  // Case 3: Generate a buffer type according to the bufferization options.
   // If we do not know the memory space and there is no default memory space,
   // report a failure.
   if (!options.defaultMemorySpace.has_value())
     return op->emitError("could not infer memory space");
 
   return getMemRefType(value, options, /*layout=*/{},
                        *options.defaultMemorySpace);
 }

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

struct FromElementsOpInterface
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto fromElementsOp = cast<tensor::FromElementsOp>(op);
     // Should the buffer be deallocated?
     bool dealloc = shouldDeallocateOpResult(
         fromElementsOp.getResult().cast<OpResult>(), options);
 
-    // TODO: Implement memory space for this op.
-    if (options.defaultMemorySpace != static_cast<unsigned>(0))
-      return op->emitError("memory space not implemented yet");
-
     // Allocate a buffer for the result.
     Location loc = op->getLoc();
     auto tensorType = fromElementsOp.getType().cast<RankedTensorType>();
     auto shape = tensorType.getShape();
     // TODO: Create alloc_tensor ops during TensorCopyInsertion.
     FailureOr<Value> tensorAlloc =
         allocateTensorForShapedValue(rewriter, loc, fromElementsOp.getResult(),
                                      /*escape=*/!dealloc, options,
                                      /*copy=*/false);
     if (failed(tensorAlloc))
       return failure();
-    auto memrefType =
-        MemRefType::get(tensorType.getShape(), tensorType.getElementType());
+    FailureOr<BaseMemRefType> maybeMemrefType =
+        bufferization::getBufferType(*tensorAlloc, options);
+    if (failed(maybeMemrefType))
+      return failure();
+    auto memrefType = maybeMemrefType->cast<MemRefType>();
     Value buffer = rewriter.create<bufferization::ToMemrefOp>(
         op->getLoc(), memrefType, *tensorAlloc);
 
     // Case: tensor<0xelem_type>.
     if (fromElementsOp.getElements().empty()) {
       replaceOpWithBufferizedValues(rewriter, op, buffer);
       return success();
     }
struct GenerateOpInterface
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto generateOp = cast<tensor::GenerateOp>(op);
     // Should the buffer be deallocated?
     bool dealloc = shouldDeallocateOpResult(
         generateOp.getResult().cast<OpResult>(), options);
 
-    // TODO: Implement memory space for this op.
-    if (options.defaultMemorySpace != static_cast<unsigned>(0))
-      return op->emitError("memory space not implemented yet");
-
-    auto tensorType = generateOp.getType().cast<RankedTensorType>();
     // Allocate memory.
     Location loc = op->getLoc();
     // TODO: Create alloc_tensor ops during TensorCopyInsertion.
     FailureOr<Value> tensorAlloc =
         allocateTensorForShapedValue(rewriter, loc, generateOp.getResult(),
                                      /*escape=*/!dealloc, options,
                                      /*copy=*/false);
     if (failed(tensorAlloc))
       return failure();
-    auto memrefType =
-        MemRefType::get(tensorType.getShape(), tensorType.getElementType());
+    FailureOr<BaseMemRefType> maybeMemrefType =
+        bufferization::getBufferType(*tensorAlloc, options);
+    if (failed(maybeMemrefType))
+      return failure();
+    auto memrefType = maybeMemrefType->cast<MemRefType>();
     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);
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     // 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());
-
     replaceOpWithBufferizedValues(rewriter, op, buffer);
-
     return success();
   }
 };
 
 /// Bufferization of tensor.insert. Replace with memref.store.
 struct InsertOpInterface
     : public BufferizableOpInterface::ExternalModel<InsertOpInterface,
                                                     tensor::InsertOp> {
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) {
       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);
+    setMemorySpaceAttr(generateOp->getOpResult(0),
+                       resultBufferType->getMemorySpaceAsInt());
     // 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());
 
     // Create tensor::InsertSliceOp.
     SmallVector<OpFoldResult> sliceSizes =
         getMixedSizes(rewriter, loc, padOp.getSource());
     SmallVector<OpFoldResult> sliceStrides(srcType.getRank(),
                                            rewriter.getIndexAttr(1));

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 @generate_memory_space
+func.func @generate_memory_space(%sz: index, %idx: index) -> index {
+  // CHECK: memref.alloc{{.*}} : memref<?xindex, 3>
+  // CHECK: scf.parallel
+  // CHECK:  memref.store {{.*}} : memref<?xindex, 3>
+  // CHECK: memref.load {{.*}} : memref<?xindex, 3>
+  // CHECK: memref.dealloc {{.*}} : memref<?xindex, 3>
+  %0 = tensor.generate %sz {
+  ^bb0(%i : index):
+    tensor.yield %sz : index
+  } { bufferization.memory_space = [3] }: tensor<?xindex>
+  %r = tensor.extract %0[%idx] : tensor<?xindex>
+  return %r : index
+}
+
+// -----
+
+// CHECK-LABEL: func @from_elements_memory_space
+func.func @from_elements_memory_space(%val: index, %idx: index) -> index {
+  // CHECK: memref.alloc() {{.*}} : memref<2x3xindex, 3>
+  // CHECK: memref.store {{.*}} : memref<2x3xindex, 3>
+  // CHECK: memref.store {{.*}} : memref<2x3xindex, 3>
+  // CHECK: memref.store {{.*}} : memref<2x3xindex, 3>
+  // CHECK: memref.store {{.*}} : memref<2x3xindex, 3>
+  // CHECK: memref.store {{.*}} : memref<2x3xindex, 3>
+  // CHECK: memref.store {{.*}} : memref<2x3xindex, 3>
+  // CHECK: memref.load {{.*}} : memref<2x3xindex, 3>
+  // CHECK: memref.dealloc {{.*}} : memref<2x3xindex, 3>
+  %0 = tensor.from_elements %val, %val, %val, %val, %val, %val
+      { bufferization.memory_space = [3] }: tensor<2x3xindex>
+  %r = tensor.extract %0[%idx, %idx] : tensor<2x3xindex>
+  return %r : index
+}
+
+// -----
+
+// 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)
+      {bufferization.memory_space = [3]} : 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
+}