[mlir][tensor][bufferize] Bufferize tensor.pad

[springerm]

Summary

tensor.pad is lowered to tensor.generate + tensor.insert_slice during bufferization. For best performance with constant padding values, users should vectorize the IR before bufferizing it.

This change also relaxes tje restriction that no new ops that bufferize to a memory write should be added during bufferization. Since bufferization has been split into two steps a while ago (tensor copy insertion + bufferization), it is reasonable to allow this now.

[tpopp]

Nice!

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

789–790

These methods recompute a vector of values every time, so ideally move them out of the loop. Alternatively, you could only obtain the dynamic values and keep a separate variable for indexing into them, but I like this more.

800

Would it be better to move over all attributes except the few that are PadOp specific? I've seen some place in MLIR provide a filter function for this.

[springerm]

address comments

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

800

Not sure, it's not clear if these attributes would have any meaning on the GenerateOp. (Should they be added to the GenerateOp or the InsertSliceOp, or both?) The escape attribute is very specific to bufferization and only makes sense for the GenerateOp.

[nicolasvasilache]

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

746

Some comments and warnings that this will bufferize to tensor.generate which in turn will jump the abstraction gap and bufferize to loops.

Also a // TODO: Reevaluate when we have a higher-level representation on buffers for generate we can avoid that jump.

775

This feels like it should be a first-class citizen of StaticValueUtils.

793

AffineApply of 2 symbols would generally be better and more composable.

811

Can we hoist this into a properly named helper in a some util file close to a notional BuiltinTypeUtils.h ?

[springerm]

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

775

That would require a dependency of StaticValueUtils (part of DialectUtils) on a dialect (Arith in this case).

[tpopp]

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

811

I feel like there is a missing method that should exist on the Op similar to getMixedLowPad()

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

void notifyOperationInserted(Operation *op) override {
     if (!hasTensorSemantics(op))
       return;
 
     // Skip ops that are not allowed to be bufferized.
     auto const &options = analysisState.getOptions();
     if (!options.isOpAllowed(op) || (opFilter && !opFilter->isOpAllowed(op)))
       return;
 
-#ifndef NDEBUG
-    // Read-only tensor ops may be created during bufferization. Ops that are
-    // writing should not be created because such ops were never analyzed.
-    // Bufferizing such ops could introduce a RaW conflict.
-    for (OpOperand &operand : op->getOpOperands())
-      if (operand.get().getType().isa<TensorType>())
-        assert(!analysisState.bufferizesToMemoryWrite(operand) &&
-               "creating tensor ops that bufferize to a memory write is not "
-               "allowed during bufferization");
-#endif // NDEBUG
-
     // Add op to worklist.
     worklist.push_back(op);
   }
 
 private:
   /// A set of all erased ops.
   DenseSet<Operation *> &erasedOps;
 

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

 //===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 using namespace mlir::tensor;
 
 namespace mlir {
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     if (failed(options.createMemCpy(rewriter, loc, *srcMemref, subView)))
       return failure();
 
     replaceOpWithBufferizedValues(rewriter, op, *dstMemref);
     return success();
   }
 };
 
+/// Bufferization of tensor.pad. Replace with tensor.generate + insert_slice.
+/// For best performance, vectorize before bufferization (better performance in
+/// case of padding with a constant).
+struct PadOpInterface
+    : public BufferizableOpInterface::ExternalModel<PadOpInterface,
+                                                    tensor::PadOp> {
+  bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    return false;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    return {};
+  }
+
+  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 toValue = [&](OpFoldResult ofr) {
+      if (ofr.is<Value>())
+        return ofr.get<Value>();
+      return rewriter
+          .create<arith::ConstantIndexOp>(loc, *getConstantIntValue(ofr))
+          .getResult();
+    };
+
+    // Compute dynamic result dimensions.
+    SmallVector<Value> dynamicSizes;
+    for (int64_t i = 0; i < resultType.getRank(); ++i) {
+      if (!resultType.isDynamicDim(i))
+        continue;
+      Value srcDim = rewriter.create<tensor::DimOp>(loc, padOp.getSource(), i);
+      Value lowPad = toValue(padOp.getMixedLowPad()[i]);
+      Value highPad = toValue(padOp.getMixedHighPad()[i]);
+      Value s1 = rewriter.create<arith::AddIOp>(loc, lowPad, highPad);
+      Value s2 = rewriter.create<arith::AddIOp>(loc, s1, srcDim);
+      dynamicSizes.push_back(s2);
+    }
+
+    // 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());
+
+    // Create tensor::InsertSliceOp.
+    SmallVector<OpFoldResult> sliceSizes, sliceStrides;
+    for (int64_t i = 0; i < resultType.getRank(); ++i) {
+      sliceStrides.push_back(rewriter.getIndexAttr(1));
+      if (srcType.isDynamicDim(i)) {
+        Value size = rewriter.create<tensor::DimOp>(loc, padOp.getSource(), i);
+        sliceSizes.push_back(size);
+      } else {
+        sliceSizes.push_back(rewriter.getIndexAttr(srcType.getDimSize(i)));
+      }
+    }
+    rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
+        padOp, padOp.getSource(), generateOp.getResult(),
+        /*offsets=*/padOp.getMixedLowPad(), sliceSizes, sliceStrides);
+
+    return success();
+  }
+};
+
 /// Bufferization of tensor.rank. Replace with memref.rank.
 struct RankOpInterface
     : public BufferizableOpInterface::ExternalModel<RankOpInterface,
                                                     tensor::RankOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
     DimOp::attachInterface<DimOpInterface>(*ctx);
     ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
     ExtractSliceOp::attachInterface<ExtractSliceOpInterface>(*ctx);
     ExtractOp::attachInterface<ExtractOpInterface>(*ctx);
     FromElementsOp::attachInterface<FromElementsOpInterface>(*ctx);
     GenerateOp::attachInterface<GenerateOpInterface>(*ctx);
     InsertOp::attachInterface<InsertOpInterface>(*ctx);
     InsertSliceOp::attachInterface<InsertSliceOpInterface>(*ctx);
+    PadOp::attachInterface<PadOpInterface>(*ctx);
     ParallelInsertSliceOp::attachInterface<ParallelInsertSliceOpInterface>(
         *ctx);
     RankOp::attachInterface<RankOpInterface>(*ctx);
     ReshapeOp::attachInterface<ReshapeOpInterface>(*ctx);
 
     // Load additional dialects of which ops may get created.
     ctx->loadDialect<arith::ArithmeticDialect, scf::SCFDialect>();
   });
 }

mlir/test/Dialect/Tensor/bufferize.mlir

func.func @tensor.reshape(%t1: tensor<?x10xf32>) -> tensor<2x2x5xf32> {
 
   // CHECK: %[[reshaped:.*]] = memref.reshape %[[m1]](%[[alloc]]) : (memref<?x10xf32>, memref<3xi64>) -> memref<2x2x5xf32>
   %reshaped = tensor.reshape %t1(%shape) : (tensor<?x10xf32>, tensor<3xi64>) -> tensor<2x2x5xf32>
 
   // CHECK: %[[r:.*]] = bufferization.to_tensor %[[reshaped]]
   // CHECK: return %[[r]]
   return %reshaped : tensor<2x2x5xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @tensor.pad(
+//  CHECK-SAME:   %[[t1:.*]]: tensor<?x10xindex>, %[[l2:.*]]: index, %[[h1:.*]]: index, %[[h2:.*]]: index
+func.func @tensor.pad(%t1: tensor<?x10xindex>, %l2: index, %h1: index,
+                      %h2: index) -> tensor<?x?xindex> {
+  // CHECK-DAG: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<?x10xindex>
+  // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
+  // CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
+  // CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index
+  // CHECK-DAG: %[[dim0:.*]] = memref.dim %[[m1]], %[[c0]]
+  // CHECK-DAG: %[[dim1:.*]] = memref.dim %[[m1]], %[[c1]]
+  // CHECK-DAG: %[[pad0:.*]] = arith.addi %[[c5]], %[[h1]]
+  // CHECK-DAG: %[[size0:.*]] = arith.addi %[[pad0]], %[[dim0]]
+  // CHECK-DAG: %[[pad1:.*]] = arith.addi %[[l2]], %[[h2]]
+  // CHECK-DAG: %[[size1:.*]] = arith.addi %[[pad1]], %[[dim1]]
+  // CHECK:     %[[alloc:.*]] = memref.alloc(%[[size0]], %[[size1]]) {{.*}} : memref<?x?xindex>
+  // CHECK:     scf.parallel ({{.*}}) = (%[[c0]], %[[c0]]) to (%[[size0]], %[[size1]]) step (%[[c1]], %[[c1]]) {
+  // CHECK:       memref.store
+  // CHECK:     }
+  // CHECK:     %[[subview:.*]] = memref.subview %[[alloc]][5, %[[l2]]] [%[[dim0]], 10] [1, 1]
+  // CHECK:     memref.copy %[[m1]], %[[subview]]
+  %0 = tensor.pad %t1 low[5, %l2] high[%h1, %h2] {
+  ^bb0(%arg0: index, %arg1: index):
+    %m = arith.muli %arg0, %arg1 : index
+    tensor.yield %m : index
+  } : tensor<?x10xindex> to tensor<?x?xindex>
+
+  // CHECK:     %[[r:.*]] = bufferization.to_tensor %[[alloc]]
+  // CHECK:     return %[[r]] : tensor<?x?xindex>
+  return %0 : tensor<?x?xindex>
+}

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

func.func @dealloc_generate_buffer(%arg: tensor<*xf32>, %sz: index, %idx: index)
   %0 = tensor.generate %sz {
   ^bb0(%i : index):
     %elem = tensor.dim %arg, %i : tensor<*xf32>
     tensor.yield %elem : index
   } : tensor<?xindex>
   %r = tensor.extract %0[%idx] : tensor<?xindex>
   return %r : index
 }
+
+// -----
+
+// CHECK-LABEL: func @dealloc_pad_buffer
+func.func @dealloc_pad_buffer(%t1: tensor<?x10xindex>, %l2: index, %h1: index,
+                              %h2: index, %idx: index) -> index {
+  // CHECK: memref.alloc
+  // CHECK: scf.parallel
+  // CHECK: memref.load
+  // CHECK: memref.dealloc
+  %0 = tensor.pad %t1 low[5, %l2] high[%h1, %h2] {
+  ^bb0(%arg0: index, %arg1: index):
+    %m = arith.muli %arg0, %arg1 : index
+    tensor.yield %m : index
+  } : tensor<?x10xindex> to tensor<?x?xindex>
+  %r = tensor.extract %0[%idx, %idx] : tensor<?x?xindex>
+  return %r : index
+}