201 lines · plain
1// RUN: mlir-opt -transform-interpreter -cse -split-input-file %s | FileCheck %s2 3// 2D tiling of dynamic 2D pad tensor op.4func.func @dynamic_2d_pad_tensor(%input_tensor: tensor<?x?xf32>,5 %pad_value: f32) -> tensor<?x?xf32> {6 %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {7 ^bb0(%arg1: index, %arg2: index):8 tensor.yield %pad_value : f329 } : tensor<?x?xf32> to tensor<?x?xf32>10 return %0 : tensor<?x?xf32>11}12 13module attributes {transform.with_named_sequence} {14 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {15 %pad = transform.structured.match ops{["tensor.pad"]} in %arg116 : (!transform.any_op) -> !transform.any_op17 %a, %b, %c = transform.structured.tile_using_for %pad tile_sizes [2, 3]18 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)19 transform.yield20 }21}22 23// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 + 8)>24// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0] -> (s0 + 7)>25// CHECK: func @dynamic_2d_pad_tensor(26// CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: tensor<?x?xf32>27// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index28// CHECK-DAG: %[[DIM_IN0:.+]] = tensor.dim %[[IN]], %[[C0]]29// CHECK-DAG: %[[DIM0:.+]] = affine.apply #[[MAP0]]()[%[[DIM_IN0]]]30// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index31// CHECK-DAG: %[[DIM_IN1:.+]] = tensor.dim %[[IN]], %[[C1]]32// CHECK-DAG: %[[DIM1:.+]] = affine.apply #[[MAP1]]()[%[[DIM_IN1]]]33// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index34// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index35// CHECK: %[[RESULT:[a-zA-Z0-9]+]] = scf.for %[[IV0:[a-zA-Z0-9]+]] = %[[C0]] to %[[DIM0]] step %[[C2]]36// CHECK: scf.for {{.*}} = %[[C0]] to %[[DIM1]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =37// CHECK: %[[SWAP_RESULT:.*]] = scf.if38// CHECK: tensor.generate39// CHECK: else40// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]41// CHECK: %[[PAD:.*]] = tensor.pad %[[SLICE]]42// CHECK: tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]43// CHECK: return %[[RESULT]]44 45// -----46 47func.func @dynamic_2d_pad_tensor_inner_tiling(%input_tensor: tensor<?x?xf32>,48 %pad_value: f32) -> tensor<?x?xf32> {49 %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {50 ^bb0(%arg1: index, %arg2: index):51 tensor.yield %pad_value : f3252 } : tensor<?x?xf32> to tensor<?x?xf32>53 return %0 : tensor<?x?xf32>54}55 56module attributes {transform.with_named_sequence} {57 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {58 %pad = transform.structured.match ops{["tensor.pad"]} in %arg159 : (!transform.any_op) -> !transform.any_op60 %a, %b = transform.structured.tile_using_for %pad tile_sizes [0, 3]61 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)62 transform.yield63 }64}65// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 + 8)>66// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0] -> (s0 + 7)>67// CHECK: func @dynamic_2d_pad_tensor_inner_tiling(68// CHECK-SAME: %[[IN:.*]]: tensor<?x?xf32>69// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index70// CHECK-DAG: %[[DIM_IN0:.*]] = tensor.dim %[[IN]], %[[C0]]71// CHECK-DAG: %[[DIM0:.*]] = affine.apply #[[MAP0]]()[%[[DIM_IN0]]]72// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index73// CHECK-DAG: %[[DIM_IN1:.*]] = tensor.dim %[[IN]], %[[C1]]74// CHECK-DAG: %[[DIM1:.*]] = affine.apply #[[MAP1]]()[%[[DIM_IN1]]]75// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index76// CHECK: %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[DIM1]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =77// CHECK: %[[SWAP_RESULT:.*]] = scf.if78// CHECK: tensor.generate79// CHECK: else80// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]81// CHECK: %[[PAD:.*]] = tensor.pad %[[SLICE]] low[3, %{{.*}}] high[{{.*}}, {{.*}}]82// CHECK: tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][0, {{.*}}] [%[[DIM0]], {{.*}}] [1, 1]83// CHECK: return %[[RESULT]]84 85// -----86 87func.func @static_pad_tensor(%input_tensor: tensor<7x9xf32>,88 %pad_value: f32) -> tensor<15x16xf32> {89 %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {90 ^bb0(%arg1: index, %arg2: index):91 tensor.yield %pad_value : f3292 } : tensor<7x9xf32> to tensor<15x16xf32>93 return %0 : tensor<15x16xf32>94}95 96module attributes {transform.with_named_sequence} {97 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {98 %pad = transform.structured.match ops{["tensor.pad"]} in %arg199 : (!transform.any_op) -> !transform.any_op100 %a, %b, %c = transform.structured.tile_using_for %pad tile_sizes [2, 3]101 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)102 transform.yield103 }104}105// CHECK-LABEL: func @static_pad_tensor(106// CHECK-SAME: %[[IN:.*]]: tensor<7x9xf32>107// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index108// CHECK-DAG: %[[C15:.*]] = arith.constant 15 : index109// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index110// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index111// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index112// CHECK: %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[C15]] step %[[C2]]113// CHECK: scf.for {{.*}} = %[[C0]] to %[[C16]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =114// CHECK: %[[SWAP_RESULT:.*]] = scf.if115// CHECK: tensor.generate116// CHECK: else117// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]118// CHECK: %[[PAD:.*]] = tensor.pad %[[SLICE]]119// CHECK: tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]120// CHECK: return %[[RESULT]]121 122// -----123 124func.func @static_pad_tensor_inner_tiling(%input_tensor: tensor<7x9xf32>,125 %pad_value: f32) -> tensor<15x16xf32> {126 %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {127 ^bb0(%arg1: index, %arg2: index):128 tensor.yield %pad_value : f32129 } : tensor<7x9xf32> to tensor<15x16xf32>130 return %0 : tensor<15x16xf32>131}132 133module attributes {transform.with_named_sequence} {134 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {135 %pad = transform.structured.match ops{["tensor.pad"]} in %arg1136 : (!transform.any_op) -> !transform.any_op137 %a, %b = transform.structured.tile_using_for %pad tile_sizes [0, 3]138 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)139 transform.yield140 }141}142// CHECK-LABEL: func @static_pad_tensor_inner_tiling(143// CHECK-SAME: %[[IN:.*]]: tensor<7x9xf32>144// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index145// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index146// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index147// CHECK: %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[C16]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =148// CHECK: %[[SWAP_RESULT:.*]] = scf.if149// CHECK: tensor.generate150// CHECK: else151// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[IN]][0, {{.*}}] [7, {{.*}}] [1, 1]152// CHECK: %[[PAD:.*]] = tensor.pad %[[SLICE]] low[3, %{{.*}}] high[5, {{.*}}]153// CHECK: tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][0, {{.*}}] [15, {{.*}}] [1, 1]154// CHECK: return %[[RESULT]]155 156/// Rest of the tests only check that they dont fail.157 158// -----159 160func.func @dynamic_2d_pad_tensor_outer_tiling(%input_tensor: tensor<?x?xf32>,161 %pad_value: f32) -> tensor<?x?xf32> {162 %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {163 ^bb0(%arg1: index, %arg2: index):164 tensor.yield %pad_value : f32165 } : tensor<?x?xf32> to tensor<?x?xf32>166 return %0 : tensor<?x?xf32>167}168 169module attributes {transform.with_named_sequence} {170 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {171 %pad = transform.structured.match ops{["tensor.pad"]} in %arg1172 : (!transform.any_op) -> !transform.any_op173 %a, %b, %c = transform.structured.tile_using_for %pad tile_sizes [2, 3]174 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)175 transform.yield176 }177}178// CHECK-LABEL: func @dynamic_2d_pad_tensor_outer_tiling179 180// -----181 182func.func @static_pad_tensor_outer_tiling(%input_tensor: tensor<7x9xf32>,183 %pad_value: f32) -> tensor<15x16xf32> {184 %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {185 ^bb0(%arg1: index, %arg2: index):186 tensor.yield %pad_value : f32187 } : tensor<7x9xf32> to tensor<15x16xf32>188 return %0 : tensor<15x16xf32>189}190 191module attributes {transform.with_named_sequence} {192 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {193 %pad = transform.structured.match ops{["tensor.pad"]} in %arg1194 : (!transform.any_op) -> !transform.any_op195 %a, %b = transform.structured.tile_using_for %pad tile_sizes [0, 3]196 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)197 transform.yield198 }199}200// CHECK-LABEL: func @static_pad_tensor_outer_tiling201