45 lines
1.8 KiB
MLIR
45 lines
1.8 KiB
MLIR
// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-split-padding-patterns %s | FileCheck %s
|
|
|
|
// CHECK-LABEL: func @pad_all_zero_sizes
|
|
func.func @pad_all_zero_sizes(%input: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
|
|
%f0 = arith.constant 0.0 : f32
|
|
%c0 = arith.constant 0 : index
|
|
%0 = tensor.pad %input low[0, %c0, 0] high[%c0, 0, 0] {
|
|
^bb0(%dim0: index, %dim1: index, %dim2: index):
|
|
tensor.yield %f0 : f32
|
|
} : tensor<?x?x?xf32> to tensor<?x?x?xf32>
|
|
return %0 : tensor<?x?x?xf32>
|
|
}
|
|
|
|
// CHECK-NOT: scf.if
|
|
// CHECK: tensor.pad
|
|
|
|
// -----
|
|
|
|
// CHECK-LABEL: func @pad_non_zero_sizes
|
|
// CHECK-SAME: (%[[INPUT:.+]]: tensor<?x?x8xf32>, %[[LOW0:.+]]: index, %[[HIGH1:.+]]: index)
|
|
func.func @pad_non_zero_sizes(%input: tensor<?x?x8xf32>, %low0: index, %high1: index) -> tensor<?x?x8xf32> {
|
|
%f0 = arith.constant 0.0 : f32
|
|
%0 = tensor.pad %input low[%low0, 0, 0] high[0, %high1, 0] {
|
|
^bb0(%dim0: index, %dim1: index, %dim2: index):
|
|
tensor.yield %f0 : f32
|
|
} : tensor<?x?x8xf32> to tensor<?x?x8xf32>
|
|
return %0 : tensor<?x?x8xf32>
|
|
}
|
|
|
|
// CHECK-DAG: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
|
|
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
|
|
// CHECK: %[[EQ0:.+]] = arith.cmpi eq, %[[LOW0]], %[[C0]] : index
|
|
// CHECK: %[[EQ1:.+]] = arith.cmpi eq, %[[HIGH1]], %[[C0]] : index
|
|
// CHECK: %[[AND:.+]] = arith.andi %[[EQ0]], %[[EQ1]] : i1
|
|
// CHECK: %[[IF:.+]] = scf.if %[[AND]] -> (tensor<?x?x8xf32>) {
|
|
// CHECK: scf.yield %[[INPUT]] : tensor<?x?x8xf32>
|
|
// CHECK: } else {
|
|
// CHECK: %[[PAD:.+]] = tensor.pad %[[INPUT]] low[%[[LOW0]], 0, 0] high[0, %[[HIGH1]], 0] {
|
|
// CHECK: ^bb0(%{{.+}}: index, %{{.+}}: index, %{{.+}}: index):
|
|
// CHECK: tensor.yield %[[F0]] : f32
|
|
// CHECK: } : tensor<?x?x8xf32> to tensor<?x?x8xf32>
|
|
// CHECK: scf.yield %[[PAD]] : tensor<?x?x8xf32>
|
|
// CHECK: }
|
|
// CHECK: return %[[IF]] : tensor<?x?x8xf32>
|