227 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -canonicalize -cse -split-input-file | FileCheck %s2 3// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 + 8)>4// CHECK-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 7)>5// CHECK: func @dynamic_pad_tensor_3_4(6// CHECK-SAME: %[[IN:.*]]: tensor<?x?xf32>7// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index8// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index9// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index10// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index11// CHECK-DAG: %[[DIM_IN0:.*]] = tensor.dim %[[IN]], %[[C0]]12// CHECK-DAG: %[[DIM_IN1:.*]] = tensor.dim %[[IN]], %[[C1]]13// CHECK-DAG: %[[DIM0:.*]] = affine.apply #[[MAP0]]()[%[[DIM_IN0]]]14// CHECK-DAG: %[[DIM1:.*]] = affine.apply #[[MAP1]]()[%[[DIM_IN1]]]15// CHECK: %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[DIM0]] step %[[C2]]16// CHECK: scf.for {{.*}} = %[[C0]] to %[[DIM1]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =17// CHECK: %[[SWAP_RESULT:.*]] = scf.if18// CHECK: tensor.generate19// CHECK: else20// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]21// CHECK: %[[PAD:.*]] = tensor.pad %[[SLICE]]22// CHECK: tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]23// CHECK: return %[[RESULT]]24 25func.func @dynamic_pad_tensor_3_4(%input_tensor: tensor<?x?xf32>,26 %pad_value: f32) -> tensor<?x?xf32> {27 %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {28 ^bb0(%arg1: index, %arg2: index):29 tensor.yield %pad_value : f3230 } : tensor<?x?xf32> to tensor<?x?xf32>31 return %0 : tensor<?x?xf32>32}33 34module attributes {transform.with_named_sequence} {35 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {36 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op37 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 3] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)38 transform.yield39 }40}41 42// -----43 44// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 + 7)>45// CHECK-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 8)>46// CHECK: func @dynamic_pad_tensor_0_3(47// CHECK-SAME: %[[IN:.*]]: tensor<?x?xf32>48// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index49// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index50// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index51// CHECK-DAG: %[[DIM_IN1:.*]] = tensor.dim %[[IN]], %[[C1]]52// CHECK-DAG: %[[DIM1:.*]] = affine.apply #[[MAP0]]()[%[[DIM_IN1]]]53// CHECK-DAG: %[[DIM_IN0:.*]] = tensor.dim %[[IN]], %[[C0]]54// CHECK-DAG: %[[DIM0:.*]] = affine.apply #[[MAP1]]()[%[[DIM_IN0]]]55// CHECK: %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[DIM1]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =56// CHECK: %[[SWAP_RESULT:.*]] = scf.if57// CHECK: tensor.generate58// CHECK: else59// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]60// CHECK: %[[PAD:.*]] = tensor.pad %[[SLICE]] low[3, %{{.*}}] high[{{.*}}, {{.*}}]61// CHECK: tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][0, {{.*}}] [%[[DIM0]], {{.*}}] [1, 1]62// CHECK: return %[[RESULT]]63 64func.func @dynamic_pad_tensor_0_3(%input_tensor: tensor<?x?xf32>,65 %pad_value: f32) -> tensor<?x?xf32> {66 %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {67 ^bb0(%arg1: index, %arg2: index):68 tensor.yield %pad_value : f3269 } : tensor<?x?xf32> to tensor<?x?xf32>70 return %0 : tensor<?x?xf32>71}72 73module attributes {transform.with_named_sequence} {74 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {75 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op76 %1, %loop = transform.structured.tile_using_for %0 tile_sizes [0, 3] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)77 transform.yield78 }79}80 81// -----82 83// CHECK-LABEL: func @static_pad_tensor_3_4(84// CHECK-SAME: %[[IN:.*]]: tensor<7x9xf32>85// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index86// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index87// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index88// CHECK-DAG: %[[C15:.*]] = arith.constant 15 : index89// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index90// CHECK: %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[C15]] step %[[C2]]91// CHECK: scf.for {{.*}} = %[[C0]] to %[[C16]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =92// CHECK: %[[SWAP_RESULT:.*]] = scf.if93// CHECK: tensor.generate94// CHECK: else95// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]96// CHECK: %[[PAD:.*]] = tensor.pad %[[SLICE]]97// CHECK: tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]98// CHECK: return %[[RESULT]]99 100func.func @static_pad_tensor_3_4(%input_tensor: tensor<7x9xf32>,101 %pad_value: f32) -> tensor<15x16xf32> {102 %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {103 ^bb0(%arg1: index, %arg2: index):104 tensor.yield %pad_value : f32105 } : tensor<7x9xf32> to tensor<15x16xf32>106 return %0 : tensor<15x16xf32>107}108 109module attributes {transform.with_named_sequence} {110 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {111 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op112 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 3] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)113 transform.yield114 }115}116 117// -----118 119// CHECK-LABEL: func @fuse_static_pad_tensor_3_4(120// CHECK-SAME: %[[IN:.*]]: tensor<7x9xf32>121// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index122// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index123// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index124// CHECK-DAG: %[[C15:.*]] = arith.constant 15 : index125// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index126// CHECK: %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[C15]] step %[[C2]]127// CHECK: scf.for {{.*}} = %[[C0]] to %[[C16]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =128// CHECK: %[[SWAP_RESULT:.*]] = scf.if129// CHECK: tensor.generate130// CHECK: else131// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]132// CHECK: %[[PAD:.*]] = tensor.pad %[[SLICE]]133// CHECK: %[[COPY:.*]] = linalg.copy ins(%[[SWAP_RESULT:.*]]134// CHECK: tensor.insert_slice %[[COPY]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]135// CHECK: return %[[RESULT]]136 137func.func @fuse_static_pad_tensor_3_4(%input_tensor: tensor<7x9xf32>,138 %pad_value: f32) -> tensor<15x16xf32> {139 %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {140 ^bb0(%arg1: index, %arg2: index):141 tensor.yield %pad_value : f32142 } : tensor<7x9xf32> to tensor<15x16xf32>143 %empty = tensor.empty() : tensor<15x16xf32>144 %1 = linalg.copy ins(%0 : tensor<15x16xf32>) outs(%empty : tensor<15x16xf32>) -> tensor<15x16xf32>145 return %1 : tensor<15x16xf32>146}147 148module attributes {transform.with_named_sequence} {149 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {150 %copy = transform.structured.match ops{["linalg.copy"]} in %arg1151 : (!transform.any_op) -> !transform.any_op152 %a, %b, %c = transform.structured.fuse %copy tile_sizes [2, 3]153 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)154 transform.yield155 }156}157 158// -----159 160// CHECK-LABEL: func @static_pad_tensor_0_3(161// CHECK-SAME: %[[IN:.*]]: tensor<7x9xf32>162// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index163// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index164// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index165// CHECK: %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[C16]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =166// CHECK: %[[SWAP_RESULT:.*]] = scf.if167// CHECK: tensor.generate168// CHECK: else169// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[IN]][0, {{.*}}] [7, {{.*}}] [1, 1]170// CHECK: %[[PAD:.*]] = tensor.pad %[[SLICE]] low[3, %{{.*}}] high[5, {{.*}}]171// CHECK: tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][0, {{.*}}] [15, {{.*}}] [1, 1]172// CHECK: return %[[RESULT]]173 174func.func @static_pad_tensor_0_3(%input_tensor: tensor<7x9xf32>,175 %pad_value: f32) -> tensor<15x16xf32> {176 %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {177 ^bb0(%arg1: index, %arg2: index):178 tensor.yield %pad_value : f32179 } : tensor<7x9xf32> to tensor<15x16xf32>180 return %0 : tensor<15x16xf32>181}182 183module attributes {transform.with_named_sequence} {184 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {185 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op186 %1, %loop = transform.structured.tile_using_for %0 tile_sizes [0, 3] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)187 transform.yield188 }189}190 191// -----192 193// CHECK-LABEL: func @static_pad_tile_evenly_0_3(194// CHECK-SAME: %[[IN:.*]]: tensor<7x9xf32>, %[[OUT:.*]]: tensor<14x15xf32>195// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index196// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index197// CHECK-DAG: %[[C15:.*]] = arith.constant 15 : index198// CHECK: %[[RESULT:.*]] = scf.for %[[IV:.*]] = %[[C0]] to %[[C15]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =199// CHECK: %[[R2:.*]] = scf.if200// CHECK: %[[GEN:.*]] = tensor.generate201// CHECK: scf.yield %[[GEN]] : tensor<14x3xf32>202// CHECK: else203// CHECK: %[[SLICE:.*]] = tensor.extract_slice %arg0[0, %{{.*}}] [7, %{{.*}}] [1, 1] : tensor<7x9xf32> to tensor<7x?xf32>204// CHECK: %[[PAD:.*]] = tensor.pad %[[SLICE]] low[0, 0] high[7, %{{.*}}]205// CHECK: scf.yield %[[PAD]] : tensor<14x3xf32>206// CHECK: %[[R3:.*]] = tensor.insert_slice %[[R2]] into %[[INNER_OUT]][0, %[[IV]]] [14, 3] [1, 1] : tensor<14x3xf32> into tensor<14x15xf32>207// CHECK: scf.yield %[[R3]] : tensor<14x15xf32>208// CHECK: return %[[RESULT]] : tensor<14x15xf32>209 210func.func @static_pad_tile_evenly_0_3(%input_tensor: tensor<7x9xf32>,211 %output_tensor: tensor<14x15xf32>,212 %pad_value: f32) -> tensor<14x15xf32> {213 %0 = tensor.pad %input_tensor low[0, 0] high[7, 6] {214 ^bb0(%arg1: index, %arg2: index):215 tensor.yield %pad_value : f32216 } : tensor<7x9xf32> to tensor<14x15xf32>217 return %0 : tensor<14x15xf32>218}219 220module attributes {transform.with_named_sequence} {221 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {222 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op223 %1, %loop = transform.structured.tile_using_for %0 tile_sizes [0, 3] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)224 transform.yield225 }226}227