[mlir][linalg] Canonicalize PadTensorOp with zero source dimension

[springerm]

Summary

If at least one source dimension has a static zero dimension, rewrite the PadTensorOp into a tensor::GenerateOp.

[springerm]

Added a few ppl from the discussion at https://llvm.discourse.group/t/what-is-the-semantics-of-memref-0xf32-and-tensor-0xf32/3557 for better visibility.

[ThomasRaoux]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

1187–1189

Can we inline the original region instead of cloning it?

I think you should be able to do something like:
rewriter.inlineRegionBefore(padOp.region(), generateOp.getRegion(), generateOp.getRegion().begin());

[springerm]

Thanks for the review.
After talking with Nicolas, we decided to not generate tensors with 0 dims in the first place. Abandoning this commit.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

LogicalResult matchAndRewrite(PadTensorOp padTensorOp,
     if (!padTensorOp.hasZeroLowPad() || !padTensorOp.hasZeroHighPad())
       return failure();
     rewriter.replaceOpWithNewOp<tensor::CastOp>(
         padTensorOp, padTensorOp.result().getType(), padTensorOp.source());
     return success();
   }
 };
 
+/// Rewrite a PadTensorOp on a tensor of size zero into a GenerateOp. E.g.:
+/// ```
+/// %0 = linalg.pad_tensor %src low[2, 3] high[4, 5]
+///     : tensor<0x0xf32> to tensor<6x8xf32>
+/// ```
+/// is rewritten into:
+/// ```
+/// %0 = tensor.generate : tensor<6x8xf32>
+/// ```
+/// Note: Basic blocks are omitted in the above example. Dynamic indices and
+/// mixed static/dynamic indices are supported.
+struct PadZeroSizeTensor : public OpRewritePattern<PadTensorOp> {
+  using OpRewritePattern<PadTensorOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(PadTensorOp padOp,
+                                PatternRewriter &rewriter) const override {
+    auto srcType = padOp.getSourceType();
+    auto resType = padOp.getResultType();
+    auto loc = padOp.getLoc();
+
+    // Pattern applies if at least one dimension has static size 0.
+    if (!llvm::any_of(srcType.getShape(), [](int64_t dim) { return dim == 0; }))
+      return failure();
+
+    // Helper function for converting an OpFoldResult to a Value.
+    auto toValue = [&](OpFoldResult ofr) -> Value {
+      if (auto val = ofr.dyn_cast<Value>())
+        return val;
+      auto intVal = getConstantIntValue(ofr);
+      assert(intVal && "expected Value or IntegerAttr");
+      return rewriter.create<ConstantIndexOp>(loc, *intVal);
+    };
+
+    // Compute dynamic dimensions of the new GenerateOp.
+    SmallVector<Value> dynDims;
+    for (int64_t i = 0; i < resType.getRank(); ++i) {
+      if (resType.isDynamicDim(i)) {
+        auto dimSize =
+            rewriter.createOrFold<memref::DimOp>(loc, padOp.source(), i);
+        auto s1 = rewriter.create<AddIOp>(loc, dimSize,
+                                          toValue(padOp.getMixedLowPad()[i]));
+        auto s2 = rewriter.create<AddIOp>(loc, s1,
+                                          toValue(padOp.getMixedHighPad()[i]));
+        dynDims.push_back(s2);
+      }
+    }
+
+    // Emit GenerateOp.
+    auto generateOp = rewriter.create<tensor::GenerateOp>(
+        loc, padOp.getResultType(), dynDims);
+    // Copy region to new op.
+    BlockAndValueMapping bvm;
+    padOp.region().cloneInto(&generateOp.getRegion(), bvm);
+    // Rewrite linalg::YieldOp to tensor::YieldOp.
+    {
+      OpBuilder::InsertionGuard guard(rewriter);
+      auto yieldOp = dyn_cast<linalg::YieldOp>(
+          generateOp.getRegion().front().getTerminator());
+      assert(yieldOp && "malformed PadTensorOp: expected YieldOp terminator");
+      assert(yieldOp.values().size() == 1);
+      rewriter.setInsertionPoint(yieldOp);
+      rewriter.replaceOpWithNewOp<tensor::YieldOp>(yieldOp,
+                                                   yieldOp.values()[0]);
+    }
+    // Replace PadTensorOp with new GenerateOp.
+    rewriter.replaceOp(padOp, generateOp.getResult());
+
+    return success();
+  }
+};
 } // namespace
 
 void PadTensorOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
-  results.add<FoldStaticZeroPadding>(context);
+  results.add<FoldStaticZeroPadding, PadZeroSizeTensor>(context);
 }
 
 /// Return the padding value of the PadTensorOp if it constant. In this context,
 /// "constant" means an actual constant or "defined outside of the block".
 ///
 /// Values are considered constant in three cases:
 ///  - A ConstantLike value.
 ///  - A basic block argument from a different block.

mlir/test/Dialect/Linalg/canonicalize.mlir

%1 = linalg.pad_tensor %0 low[%c0, %c0] high[%c0, %c0]  {
     ^bb0(%arg1: index, %arg2: index):  // no predecessors
       linalg.yield %cst : f32
   } : tensor<?x?xf32> to tensor<4x4xf32>
   return %1 : tensor<4x4xf32>
 }
 // CHECK-LABEL: @tensor_pad_cast
 // CHECK-SAME: %[[ARG0:.+]]: tensor<4x4xf32>
 // CHECK: return %[[ARG0]]
+
+// -----
+
+// CHECK-LABEL: func @pad_zero_tensor_to_generate(
+//  CHECK-SAME:     %[[ARG0:.*]]: tensor<?x0x8xf32>, %[[L1:.*]]: index, %[[PAD:.*]]: f32
+//   CHECK-NOT:   linalg.pad_tensor
+//   CHECK-DAG:   %[[C0:.*]] = constant 0 : index
+//   CHECK-DAG:   %[[C2:.*]] = constant 2 : index
+//   CHECK-DAG:   %[[C5:.*]] = constant 5 : index
+//       CHECK:   %[[DIM:.*]] = memref.dim %[[ARG0]], %[[C0]]
+//       CHECK:   %[[S1:.*]] = addi %[[DIM]], %[[C5]]
+//       CHECK:   %[[S2:.*]] = addi %[[L1]], %[[C2]]
+//       CHECK:   %[[RESULT:.*]] = tensor.generate %[[S1]], %[[S2]] {
+//       CHECK:     tensor.yield %[[PAD]] : f32
+//       CHECK:   } : tensor<?x?x12xf32>
+//       CHECK:   return %[[RESULT]]
+func @pad_zero_tensor_to_generate(
+    %arg0: tensor<?x0x8xf32>, %l1: index, %pad: f32) -> tensor<?x?x12xf32> {
+  %0 = linalg.pad_tensor %arg0 low[1, %l1, 3] high[4, 2, 1] {
+    ^bb0(%arg1: index, %arg2: index, %arg3: index):
+      linalg.yield %pad : f32
+  } : tensor<?x0x8xf32> to tensor<?x?x12xf32>
+  return %0 : tensor<?x?x12xf32>
+}