[Draft][MLIR] Add reshape propagation through tensor.pad

[ita9naiwa]

iree-org/iree#17492 (comment)

I’ve implemented fusion for tensor.expand_shape → tensor.pad, but two gaps remain:

1. Missing collapse‑side pattern.
   I haven’t yet added the mirror case for tensor.collapse_shape → tensor.pad.
2. Static‑only support
   The current pattern only handles fully static shapes and padding.

Before (expand then pad):

func.func @fold_tensor_pad_with_expand(%arg0: tensor<512x256x256xf32>) -> tensor<32x16x258x258xf32> {
  %c0 = arith.constant 0.0 : f32
  %producer = linalg.fill ins(%c0 : f32) outs(%arg0 : tensor<512x256x256xf32>) -> tensor<512x256x256xf32>

  %pad = tensor.pad %producer low[0, 1, 1] high[0, 1, 1] {
    ^bb0(%i: index, %j: index, %k: index):
      tensor.yield %c0 : f32
  } : tensor<512x256x256xf32> to tensor<512x258x258xf32>
  %reshape = tensor.expand_shape %pad [[0, 1], [2], [3]]
      output_shape [32, 16, 258, 258] : tensor<512x258x258xf32> into tensor<32x16x258x258xf32>

  return %reshape : tensor<32x16x258x258xf32>
}

After (reshape then pad):

func.func @fold_tensor_pad_with_expand(%arg0: tensor<512x256x256xf32>) -> tensor<32x16x258x258xf32> {
  %c0 = arith.constant 0.0 : f32
  %producer = linalg.fill ins(%c0 : f32) outs(%arg0 : tensor<512x256x256xf32>) -> tensor<512x256x256xf32>

  %reshape = tensor.expand_shape %producer  [[0, 1], [2], [3]]
      output_shape [32, 16, 258, 258] : tensor<512x256x256xf32> into tensor<32x16x256x256xf32>
  %pad = tensor.pad %reshape low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%i0: index, %i1: index, %i2: index, %i3: index):
      tensor.yield %c0 : f32
  } : tensor<32x16x256x256xf32> to tensor<32x16x258x258xf32>

  return %pad : tensor<32x16x258x258xf32>
}

Next steps
• Add a CollapseShapeOp→PadOp pattern to cover the missing collapse‑side fusion.
• Lift the “static‑only” guard so both patterns handle dynamic shapes and pads.

CC @Max191 for awareness—would love any pointers on the collapse‑side implementation or dynamic‑shape handling!

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.