add tests for non const pad val

Signed-off-by: Max Dawkins <[email protected]>

Author: Max191

mlir/test/Dialect/Linalg/fuse-with-reshape-by-collapsing.mlir

@@ -594,6 +594,24 @@ func.func @fuse_by_collapsing_pad(%arg0 : tensor<2x12x5x336x9xi32>) -> tensor<8x
 
 // -----
 
+func.func @no_fuse_by_collapsing_pad_non_constant_padding(%arg0 : tensor<2x12xi32>) -> tensor<8x3x4xi32> {
+  %expand = tensor.expand_shape %arg0 [[0], [1, 2]] output_shape [2, 3, 4] : tensor<2x12xi32> into tensor<2x3x4xi32>
+  %cst = arith.constant 0 : i32
+  %padded_0 = tensor.pad %expand low[1, 0, 0] high[5, 0, 0] {
+  ^bb0(%arg1: index, %arg2: index, %arg3: index):
+    %pad_val = arith.index_cast %arg1 : index to i32
+    tensor.yield %pad_val : i32
+  } : tensor<2x3x4xi32> to tensor<8x3x4xi32>
+  return %padded_0 : tensor<8x3x4xi32>
+}
+//      CHECK: func @no_fuse_by_collapsing_pad_non_constant_padding(
+// CHECK-SAME:   %[[ARG0:.+]]: tensor<2x12xi32>)
+//      CHECK:   %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]]
+//      CHECK:   %[[PAD:.+]] = tensor.pad %[[EXPAND]]
+//      CHECK:   return %[[PAD]]
+
+// -----
+
 func.func @no_fuse_by_collapsing_pad(%arg0 : tensor<2x12x5x336x9xi32>) -> tensor<8x5x4x17x6x7x8x14xi32> {
   %expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] output_shape [2, 3, 4, 5, 6, 7, 8, 9] : tensor<2x12x5x336x9xi32> into tensor<2x3x4x5x6x7x8x9xi32>
   %cst = arith.constant 0 : i32
@@ -678,6 +696,24 @@ func.func @collapse_shape_with_producer_pad_dynamic(%arg0: tensor<?x?x?x?x?x?xf3
 //      CHECK:   %[[PAD:.+]] = tensor.pad %[[COLLAPSE]] low[%[[L0]], 0, %[[L1]], 0] high[%[[H0]], 0, %[[H1]], 0]
 //      CHECK:   return %[[PAD]]
 
+// -----
+
+func.func @collapse_shape_with_producer_pad_non_constant_padding(%arg0 : tensor<2x3x4xi32>) -> tensor<8x12xi32> {
+  %cst = arith.constant 0 : i32
+  %padded_0 = tensor.pad %arg0 low[1, 0, 0] high[5, 0, 0] {
+  ^bb0(%arg1: index, %arg2: index, %arg3: index):
+    %pad_val = arith.index_cast %arg1 : index to i32
+    tensor.yield %pad_val : i32
+  } : tensor<2x3x4xi32> to tensor<8x3x4xi32>
+  %collapsed = tensor.collapse_shape %padded_0 [[0], [1, 2]] : tensor<8x3x4xi32> into tensor<8x12xi32>
+  return %collapsed : tensor<8x12xi32>
+}
+//      CHECK: func @collapse_shape_with_producer_pad_non_constant_padding(
+// CHECK-SAME:   %[[ARG0:.+]]: tensor<2x3x4xi32>)
+//      CHECK:   %[[PAD:.+]] = tensor.pad %[[ARG0]]
+//      CHECK:   %[[COLLAPSED:.+]] = tensor.collapse_shape %[[PAD]]
+//      CHECK:   return %[[COLLAPSED]]
+
 // -----
 // Static problem sizes. Checks all aspects of fusion by collapsing with bubbling up collapse shapes.
 #map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>

mlir/test/Dialect/Linalg/reshape_fusion.mlir

@@ -822,6 +822,23 @@ func.func @fuse_by_expanding_pad(%arg0 : tensor<2x3x4x5x6x7x8x9xi32>) -> tensor<
 
 // -----
 
+func.func @no_fuse_by_expanding_pad_non_constant_padding(%arg0 : tensor<2x3x4xi32>) -> tensor<8x12xi32> {
+  %collapse = tensor.collapse_shape %arg0 [[0], [1, 2]] : tensor<2x3x4xi32> into tensor<2x12xi32>
+  %padded_0 = tensor.pad %collapse low[1, 0] high[5, 0] {
+  ^bb0(%arg1: index, %arg2: index):
+    %pad_val = arith.index_cast %arg1 : index to i32
+    tensor.yield %pad_val : i32
+  } : tensor<2x12xi32> to tensor<8x12xi32>
+  return %padded_0 : tensor<8x12xi32> 
+}
+//      CHECK: func @no_fuse_by_expanding_pad_non_constant_padding(
+// CHECK-SAME:   %[[ARG0:.+]]: tensor<2x3x4xi32>)
+//      CHECK:   %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG0]]
+//      CHECK:   %[[PAD:.+]] = tensor.pad %[[COLLAPSE]]
+//      CHECK:   return %[[PAD]]
+
+// -----
+
 func.func @no_fuse_by_expanding_pad(%arg0 : tensor<2x3x4x5x6x7x8x9xi32>) -> tensor<8x12x17x339x14xi32> {
   %collapse = tensor.collapse_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] : tensor<2x3x4x5x6x7x8x9xi32> into tensor<2x12x5x336x9xi32>
   %cst = arith.constant 0 : i32
@@ -904,6 +921,23 @@ func.func @expand_shape_with_producer_pad_dynamic(%arg0: tensor<?x?x?x?xf32>,
 
 // -----
 
+func.func @expand_shape_with_producer_pad_non_constant_padding(%arg0 : tensor<2x12xi32>) -> tensor<8x3x4xi32> {
+  %padded_0 = tensor.pad %arg0 low[1, 0] high[5, 0] {
+  ^bb0(%arg1: index, %arg2: index):
+    %pad_val = arith.index_cast %arg1 : index to i32
+    tensor.yield %pad_val : i32
+  } : tensor<2x12xi32> to tensor<8x12xi32>
+  %expand = tensor.expand_shape %padded_0 [[0], [1, 2]] output_shape [8, 3, 4] : tensor<8x12xi32> into tensor<8x3x4xi32>
+  return %expand : tensor<8x3x4xi32> 
+}
+//      CHECK: func @expand_shape_with_producer_pad_non_constant_padding(
+// CHECK-SAME:   %[[ARG0:.+]]: tensor<2x12xi32>)
+//      CHECK:   %[[PAD:.+]] = tensor.pad %[[ARG0]]
+//      CHECK:   %[[EXPAND:.+]] = tensor.expand_shape %[[PAD]]
+//      CHECK:   return %[[EXPAND]]
+
+// -----
+
 func.func @move_operand_deps(%arg0 : tensor<?x128xf16>,
     %arg1 : tensor<4x?x32x128xf16>, %empty : tensor<4x?x32x128xf16>) -> tensor<4x?x32x8x16xf16> {
   %c0 = arith.constant 0 : index