[mlir][linalg] Do not emit FillOp for tensor.pad with zero padding

[springerm]

Summary

No need to fill the buffer if no padding is added. I.e., the tensor.pad is packing only.

Also improve tensor::createPadHighOp to use attributes instead of SSA values for high padding sizes when possible.

Depends On: D153555

[springerm]

rebase

[mravishankar]

mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp

184–189

This seems strange. Why is there even a pad op if we don't have low pad or high pad. That seems like that pad should be folded away.

[springerm]

mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp

184–189

This is used for "packing" tensor.pad ops. I.e., tensor.pad ops that do not necessarily add any padding, but where a new buffer allocation is required (e.g., in shared memory on GPU). Such tensor.pad ops have a nofold attribute.

If `nofold` is set, the padding operation will not be folded away even
if the source type and the padded type have the same static shape. This can
be used, e.g., for packing or promotion to faster memory.

I'd like to remove the nofold attribute and this kind of use of tensor.pad, but it's used in quite a few places, so this will still take a while. The infrastructure that I've been adding in ConvertToDestinationStyle.cpp (in particular bufferizeToAllocation) over the last weeks will bring us closer to that. (The idea is that we generate a memref allocation right away and do not even have a tensor.pad if there's nothing to pad.)

[springerm]

mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp

184–189

To clarify: We currently use tensor.pad with nofold to force a shared memory allocation (or any allocation in general). I'm working on using a different abstraction instead (memref.alloc). We could not do this until recently because bufferization of mixed tensor/memref IR was not supported.

[nicolasvasilache]

thanks!

mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp

Value linalg::bufferizeToAllocation(RewriterBase &rewriter, PadOp padOp,
   rewriter.setInsertionPoint(insertionPoint ? insertionPoint : padOp);
   Location loc = padOp.getLoc();
 
   // Create buffer allocation.
   Value alloc =
       createAllocationForTensor(rewriter, loc, padOp.getResult(), memorySpace);
   rewriter.setInsertionPoint(padOp);
 
-  // Create linalg.fill or linalg.generic.
-  Operation *fillOp = movePaddingToFillOrGenericOp(rewriter, loc, padOp, alloc);
+  if (!padOp.hasZeroLowPad() || !padOp.hasZeroHighPad()) {
+    // Create linalg.fill or linalg.generic. Not needed if there is no padding.
+    Operation *fillOp =
+        movePaddingToFillOrGenericOp(rewriter, loc, padOp, alloc);
     rewriter.setInsertionPointAfter(fillOp);
+  }
 
   // Create memref.tensor_store.
   SmallVector<OpFoldResult> sizes =
       getMixedSizes(rewriter, loc, padOp.getSource());
   SmallVector<OpFoldResult> strides(padOp.getResultType().getRank(),
                                     rewriter.getIndexAttr(1));
   Value subview = rewriter.create<memref::SubViewOp>(
       loc, alloc, /*offsets=*/padOp.getMixedLowPad(), sizes, strides);

mlir/test/Dialect/Linalg/pad-to-specific-memory-space.mlir

func.func @pad_to_memory_space(%arg0: tensor<24x12xf32>,
   // CHECK: memref.copy %[[s0]], %[[alloc0_view]]
 
   // CHECK: %[[alloc1:.*]] = memref.alloc() : memref<7x5xf32, 3>
   // CHECK: linalg.fill {{.*}} outs(%[[alloc1]]
   // CHECK: %[[alloc1_view:.*]] = memref.subview %[[alloc1]][0, 0] [%{{.*}}, 5] [1, 1]
   // CHECK: memref.copy %[[s1]], %[[alloc1_view]]
 
   // CHECK: %[[alloc2:.*]] = memref.alloc() : memref<4x5xf32, 3>
-  // CHECK: linalg.fill {{.*}} outs(%[[alloc2]]
+  // CHECK-NOT: linalg.fill {{.*}} outs(%[[alloc2]]
   // No subview because there is 0 padding
   // CHECK: memref.copy %[[s2]], %[[alloc2]]
 
   // CHECK: linalg.matmul ins(%[[alloc0]], %[[alloc1]] : {{.*}}) outs(%[[alloc2]] : {{.*}})
   // Copy back result.
   // CHECK: memref.copy %[[alloc2]], %[[s2]]
   %4 = linalg.matmul ins(%1, %2 : tensor<4x?xf32>, tensor<?x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>