1// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-split-padding-patterns %s | FileCheck %s
2
3// CHECK-LABEL: func @pad_all_zero_sizes
4func.func @pad_all_zero_sizes(%input: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
5  %f0 = arith.constant 0.0 : f32
6  %c0 = arith.constant 0 : index
7  %0 = tensor.pad %input low[0, %c0, 0] high[%c0, 0, 0] {
8  ^bb0(%dim0: index, %dim1: index, %dim2: index):
9    tensor.yield %f0 : f32
10  } : tensor<?x?x?xf32> to tensor<?x?x?xf32>
11  return %0 : tensor<?x?x?xf32>
12}
13
14// CHECK-NOT: scf.if
15//     CHECK: tensor.pad
16
17// -----
18
19// CHECK-LABEL: func @pad_non_zero_sizes
20//  CHECK-SAME: (%[[INPUT:.+]]: tensor<?x?x8xf32>, %[[LOW0:.+]]: index, %[[HIGH1:.+]]: index)
21func.func @pad_non_zero_sizes(%input: tensor<?x?x8xf32>, %low0: index, %high1: index) -> tensor<?x?x8xf32> {
22  %f0 = arith.constant 0.0 : f32
23  %0 = tensor.pad %input low[%low0, 0, 0] high[0, %high1, 0] {
24  ^bb0(%dim0: index, %dim1: index, %dim2: index):
25    tensor.yield %f0 : f32
26  } : tensor<?x?x8xf32> to tensor<?x?x8xf32>
27  return %0 : tensor<?x?x8xf32>
28}
29
30// CHECK-DAG: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
31// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
32// CHECK: %[[EQ0:.+]] = arith.cmpi eq, %[[LOW0]], %[[C0]] : index
33// CHECK: %[[EQ1:.+]] = arith.cmpi eq, %[[HIGH1]], %[[C0]] : index
34// CHECK: %[[AND:.+]] = arith.andi %[[EQ0]], %[[EQ1]] : i1
35// CHECK: %[[IF:.+]] = scf.if %[[AND]] -> (tensor<?x?x8xf32>) {
36// CHECK:   scf.yield %[[INPUT]] : tensor<?x?x8xf32>
37// CHECK: } else {
38// CHECK:   %[[PAD:.+]] = tensor.pad %[[INPUT]] low[%[[LOW0]], 0, 0] high[0, %[[HIGH1]], 0]  {
39// CHECK:   ^bb0(%{{.+}}: index, %{{.+}}: index, %{{.+}}: index):
40// CHECK:     tensor.yield %[[F0]] : f32
41// CHECK:   } : tensor<?x?x8xf32> to tensor<?x?x8xf32>
42// CHECK:   scf.yield %[[PAD]] : tensor<?x?x8xf32>
43// CHECK: }
44// CHECK: return %[[IF]] : tensor<?x?x8xf32>
45