227 lines · plain
1// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-rewrite-extract-slice-from-collapse-shape %s | FileCheck %s2// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns="test-rewrite-extract-slice-from-collapse-shape use-foreach" %s | FileCheck %s --check-prefix=FOREACH3 4func.func @extract_slice_static(%input: tensor<3x5x7x11xf32>) -> tensor<20x11xf32> {5 %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3]] : tensor<3x5x7x11xf32> into tensor<105x11xf32>6 %slice = tensor.extract_slice %collapsed [0, 0] [20, 11] [1, 1] : tensor<105x11xf32> to tensor<20x11xf32>7 return %slice : tensor<20x11xf32>8}9 10// CHECK: func.func @extract_slice_static(%[[arg0:.+]]:11// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index12// CHECK-DAG: %[[c20:.+]] = arith.constant 20 : index13// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index14// CHECK-DAG: %[[init:.+]] = tensor.empty() : tensor<20x11xf32>15// CHECK-DAG: %[[tile:.+]] = scf.for %[[iv:.+]] = %[[c0]] to %[[c20]] step %[[c1]] iter_args(%[[iterArg:.+]] = %[[init]])16// CHECK: %[[multiIndex:.+]]:3 = affine.delinearize_index %[[iv]] into (3, 5, 717// CHECK: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 11] [1, 1, 1, 1] :18// CHECK: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} :19// CHECK: %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg]][%[[iv]], 0] [1, 11] [1, 1] :20// CHECK: scf.yield %[[update]] :21// CHECK: return %[[tile]]22 23// FOREACH: func.func @extract_slice_static(%[[arg0:.+]]:24// FOREACH-DAG: %[[init:.+]] = tensor.empty() : tensor<20x11xf32>25// FOREACH: %[[tile:.+]] = scf.forall (%[[iv:.+]]) in (20) shared_outs(%[[dest:.+]] = %[[init]])26// FOREACH: %[[multiIndex:.+]]:3 = affine.delinearize_index %[[iv]] into (3, 5, 727// FOREACH: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 11] [1, 1, 1, 1] :28// FOREACH: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} :29// FOREACH: in_parallel30// FOREACH-NEXT: tensor.parallel_insert_slice %[[sliceFlat]] into %[[dest]][%[[iv]], 0] [1, 11] [1, 1] :31// FOREACH: return %[[tile]]32 33// -----34 35 36func.func @extract_slice_static_strided(%input: tensor<3x5x7x11xf32>) -> tensor<10x5xf32> {37 %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3]] : tensor<3x5x7x11xf32> into tensor<105x11xf32>38 %slice = tensor.extract_slice %collapsed [13, 0] [10, 5] [2, 2] : tensor<105x11xf32> to tensor<10x5xf32>39 return %slice : tensor<10x5xf32>40}41 42// CHECK: #[[$map0:.+]] = affine_map<(d0) -> (d0 * 2 + 13)>43// CHECK: func.func @extract_slice_static_strided(%[[arg0:.+]]:44// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index45// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index46// CHECK-DAG: %[[c10:.+]] = arith.constant 10 : index47// CHECK: %[[init:.+]] = tensor.empty() : tensor<10x5xf32>48// CHECK: %[[tile:.+]] = scf.for %[[iv:.+]] = %[[c0]] to %[[c10]] step %[[c1]] iter_args(%[[iterArg:.+]] = %[[init]])49// CHECK: %[[inputIv:.+]] = affine.apply #[[$map0]](%[[iv]])50// CHECK: %[[multiIndex:.+]]:3 = affine.delinearize_index %[[inputIv]] into (3, 5, 751// CHECK: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 5] [1, 1, 1, 2] :52// CHECK: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} :53// CHECK: %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg]][%[[iv]], 0] [1, 5] [1, 1] :54// CHECK: scf.yield %[[update]] :55// CHECK: return %[[tile]]56 57 58// -----59 60 61func.func @extract_slice_dynamic(%input: tensor<3x?x?x11xf32>, %offt: index, %size: index) -> tensor<?x5xf32> {62 %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3]] : tensor<3x?x?x11xf32> into tensor<?x11xf32>63 %slice = tensor.extract_slice %collapsed [%offt, 0] [%size, 5] [2, 2] : tensor<?x11xf32> to tensor<?x5xf32>64 return %slice : tensor<?x5xf32>65}66 67// CHECK: #[[map0:.+]] = affine_map<(d0)[s0] -> (d0 * 2 + s0)>68// CHECK: func.func @extract_slice_dynamic(%[[arg0:.+]]: tensor<{{.*}}>, %[[lb:.+]]: index, %[[sz:.+]]: index)69// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index70// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index71// CHECK-DAG: %[[c2:.+]] = arith.constant 2 : index72// CHECK: %[[init:.+]] = tensor.empty(%[[sz]]) : tensor<?x5xf32>73// CHECK-DAG: %[[d1:.+]] = tensor.dim %arg0, %[[c1]] : tensor<3x?x?x11xf32>74// CHECK-DAG: %[[d2:.+]] = tensor.dim %arg0, %[[c2]] : tensor<3x?x?x11xf32>75// CHECK: %[[tile:.+]] = scf.for %[[iv:.+]] = %[[c0]] to %[[sz]] step %[[c1]] iter_args(%[[iterArg:.+]] = %[[init]])76// CHECK: %[[inputIv:.+]] = affine.apply #[[map0]](%[[iv]])[%[[lb]]]77// CHECK: %[[multiIndex:.+]]:3 = affine.delinearize_index %[[inputIv]] into (3, %[[d1]], %[[d2]]) :78// CHECK: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 5] [1, 1, 1, 2] :79// CHECK: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} :80// CHECK: %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg]][%[[iv]], 0] [1, 5] [1, 1] :81// CHECK: scf.yield %[[update]] :82// CHECK: return %[[tile]] :83 84// -----85 86 87func.func @extract_slice_dynamic_multidim(%input: tensor<3x?x?x11x?xf32>, %offt0: index, %size0: index, %offt1: index, %size1: index) -> tensor<?x?xf32> {88 %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3, 4]] : tensor<3x?x?x11x?xf32> into tensor<?x?xf32>89 %slice = tensor.extract_slice %collapsed [%offt0, %offt1] [%size0, %size1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>90 return %slice : tensor<?x?xf32>91}92 93// CHECK: #[[map0:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>94// CHECK: func.func @extract_slice_dynamic_multidim(%[[arg0:.+]]: tensor<3x?x?x11x?xf32>, %[[lb1:.+]]: index, %[[sz1:.+]]: index, %[[lb2:.+]]: index, %[[sz2:.+]]: index)95// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index96// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index97// CHECK-DAG: %[[c2:.+]] = arith.constant 2 : index98// CHECK-DAG: %[[c4:.+]] = arith.constant 4 : index99// CHECK: %[[init:.+]] = tensor.empty(%[[sz1]], %[[sz2]]) : tensor<?x?xf32>100// CHECK-DAG: %[[d1:.+]] = tensor.dim %[[arg0]], %[[c1]] :101// CHECK-DAG: %[[d2:.+]] = tensor.dim %[[arg0]], %[[c2]] :102// CHECK-DAG: %[[d4:.+]] = tensor.dim %[[arg0]], %[[c4]] :103// CHECK: %[[tile1:.+]] = scf.for %[[iv1:.+]] = %[[c0]] to %[[sz1]] step %[[c1]] iter_args(%[[iterArg1:.+]] = %[[init]])104// CHECK: %[[tile2:.+]] = scf.for %[[iv2:.+]] = %[[c0]] to %[[sz2]] step %[[c1]] iter_args(%[[iterArg2:.+]] = %[[iterArg1]])105// CHECK: %[[inputIv1:.+]] = affine.apply #[[map0:.+]](%[[iv1]])[%[[lb1]]]106// CHECK: %[[multiIndex1:.+]]:3 = affine.delinearize_index %[[inputIv1]] into (3, %[[d1]], %[[d2]]) :107// CHECK: %[[inputIv2:.+]] = affine.apply #[[map0:.+]](%[[iv2]])[%[[lb2]]]108// CHECK: %[[multiIndex2:.+]]:2 = affine.delinearize_index %[[inputIv2]] into (11, %[[d4]]) :109// CHECK: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex1]]#0, %[[multiIndex1]]#1, %[[multiIndex1]]#2, %[[multiIndex2]]#0, %[[multiIndex2]]#1] [1, 1, 1, 1, 1] [1, 1, 1, 1, 1] :110// CHECK: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3, 4]{{\]}} :111// CHECK: %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg2]][%[[iv1]], %[[iv2]]] [1, 1] [1, 1] :112// CHECK: scf.yield %[[update]] :113// CHECK: scf.yield %[[tile2]] :114// CHECK: return %[[tile1]] :115 116// FOREACH: #[[map1:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>117// FOREACH: func.func @extract_slice_dynamic_multidim(%[[arg0:.+]]: tensor<3x?x?x11x?xf32>, %[[lb1:.+]]: index, %[[sz1:.+]]: index, %[[lb2:.+]]: index, %[[sz2:.+]]: index)118// FOREACH-DAG: %[[c1:.+]] = arith.constant 1 : index119// FOREACH-DAG: %[[c2:.+]] = arith.constant 2 : index120// FOREACH-DAG: %[[c4:.+]] = arith.constant 4 : index121// FOREACH: %[[init:.+]] = tensor.empty(%[[sz1]], %[[sz2]]) : tensor<?x?xf32>122// FOREACH-DAG: %[[d1:.+]] = tensor.dim %[[arg0]], %[[c1]] :123// FOREACH-DAG: %[[d2:.+]] = tensor.dim %[[arg0]], %[[c2]] :124// FOREACH-DAG: %[[d4:.+]] = tensor.dim %[[arg0]], %[[c4]] :125// FOREACH: %[[tile1:.+]] = scf.forall (%[[tid1:.+]], %[[tid2:.+]]) in (%[[sz1]], %[[sz2]]) shared_outs(%[[dest:.+]] = %[[init]])126// FOREACH-DAG: %[[iv1:.+]] = affine.apply #[[map1]](%[[tid1]])[%[[lb1]]]127// FOREACH: %[[multiIndex1:.+]]:3 = affine.delinearize_index %[[iv1]] into (3, %[[d1]], %[[d2]]) :128// FOREACH-DAG: %[[iv2:.+]] = affine.apply #[[map1]](%[[tid2]])[%[[lb2]]]129// FOREACH: %[[multiIndex2:.+]]:2 = affine.delinearize_index %[[iv2]] into (11, %[[d4]]) :130// FOREACH: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex1]]#0, %[[multiIndex1]]#1, %[[multiIndex1]]#2, %[[multiIndex2]]#0, %[[multiIndex2]]#1] [1, 1, 1, 1, 1] [1, 1, 1, 1, 1] :131// FOREACH: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3, 4]{{\]}} :132// FOREACH: in_parallel133//FOREACH-NEXT: tensor.parallel_insert_slice %[[sliceFlat]] into %[[dest]][%[[tid1]], %[[tid2]]] [1, 1] [1, 1] :134 135// -----136 137// Verifies that a linearized dimension that is not sliced does not generate a loop. Note that this138// only works for static shapes.139 140// CHECK: @extract_slice_non_sliced_linearized_dim(%[[arg0:.+]]: tensor<{{.*}}>,141func.func @extract_slice_non_sliced_linearized_dim(%input: tensor<3x?x?x11x2xf32>, %offt: index, %size: index) -> tensor<?x22xf32> {142 %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3, 4]] : tensor<3x?x?x11x2xf32> into tensor<?x22xf32>143 %slice = tensor.extract_slice %collapsed [%offt, 0] [%size, 22] [1, 1] : tensor<?x22xf32> to tensor<?x22xf32>144 // CHECK: scf.for145 // CHECK-NOT: scf.for146 // CHECK: %[[multiIndex:.+]]:3 = affine.delinearize_index147 // CHECK: tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0, 0] [1, 1, 1, 11, 2] [1, 1, 1, 1, 1]148 return %slice : tensor<?x22xf32>149}150 151// -----152 153// CHECK: @no_sliced_linearized_dims(%[[arg0:.+]]: tensor<{{.*}}>, %[[arg1:.+]]: index, %[[arg2:.+]]: index154func.func @no_sliced_linearized_dims(%input: tensor<30x11x100xf32>, %offt: index, %size: index) -> tensor<330x?xf32> {155 %collapsed = tensor.collapse_shape %input [[0, 1], [2]] : tensor<30x11x100xf32> into tensor<330x100xf32>156 %slice = tensor.extract_slice %collapsed [0, %offt] [330, %size] [1, 1] : tensor<330x100xf32> to tensor<330x?xf32>157 // CHECK-NOT: scf.for158 // CHECK: %[[init:.+]] = tensor.empty(%[[arg2]])159 // CHECK: %[[e:.+]] = tensor.extract_slice %[[arg0]][0, 0, %[[arg1]]] [30, 11, %[[arg2]]] [1, 1, 1]160 // CHECK: %[[c:.+]] = tensor.collapse_shape %[[e]] {{\[}}[0, 1], [2]]161 // CHECK: %[[res:.+]] = tensor.insert_slice %[[c]] into %[[init]]162 // CHECK: return %[[res]]163 return %slice : tensor<330x?xf32>164}165 166// -----167 168// The below tests verify that a dimension which is the result of collapsing at169// most one non-unit dim is handled properly.170 171// CHECK: @collapse_and_slice_unit_dim(%[[arg0:.+]]: tensor<{{.*}}>, %[[arg1:.+]]: index, %[[arg2:.+]]: index172func.func @collapse_and_slice_unit_dim(%input: tensor<1x11x100xf32>, %offt: index, %size: index) -> tensor<?x100xf32> {173 %collapsed = tensor.collapse_shape %input [[0, 1], [2]] : tensor<1x11x100xf32> into tensor<11x100xf32>174 %slice = tensor.extract_slice %collapsed [%offt, 0] [%size, 100] [1, 1] : tensor<11x100xf32> to tensor<?x100xf32>175 // CHECK-NOT: scf.for176 // CHECK: %[[e:.+]] = tensor.extract_slice %[[arg0]][0, 0, 0] [1, 11, 100] [1, 1, 1]177 // CHECK-SAME: tensor<1x11x100xf32> to tensor<11x100xf32>178 // CHECK: %[[e1:.+]] = tensor.extract_slice %[[e]][%[[arg1]], 0] [%[[arg2]], 100] [1, 1]179 // CHECK-SAME: tensor<11x100xf32> to tensor<?x100xf32>180 return %slice : tensor<?x100xf32>181}182 183// CHECK: @collapse_and_slice_multiple_unit_dim_dynamic(%[[arg0:.+]]: tensor<{{.*}}>, %[[arg1:.+]]: index, %[[arg2:.+]]: index184func.func @collapse_and_slice_multiple_unit_dim_dynamic(%input: tensor<1x?x1x100xf32>, %offt: index, %size: index) -> tensor<?x100xf32> {185 %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3]] : tensor<1x?x1x100xf32> into tensor<?x100xf32>186 %slice = tensor.extract_slice %collapsed [%offt, 0] [%size, 100] [1, 1] : tensor<?x100xf32> to tensor<?x100xf32>187 // CHECK-NOT: scf.for188 // CHECK: %[[c1:.+]] = arith.constant 1 : index189 // CHECK: %[[dim:.+]] = tensor.dim %[[arg0]], %[[c1]] :190 // CHECK: %[[e:.+]] = tensor.extract_slice %[[arg0]][0, 0, 0, 0] [1, %[[dim]], 1, 100] [1, 1, 1, 1]191 // CHECK-SAME: tensor<1x?x1x100xf32> to tensor<?x100xf32>192 // CHECK: %[[e1:.+]] = tensor.extract_slice %[[e]][%[[arg1]], 0] [%[[arg2]], 100] [1, 1]193 // CHECK-SAME: tensor<?x100xf32> to tensor<?x100xf32>194 return %slice : tensor<?x100xf32>195}196 197// CHECK: @collapse_and_slice_multiple_unit_dim_mixed(%[[arg0:.+]]: tensor<{{.*}}>, %[[arg1:.+]]: index, %[[arg2:.+]]: index198func.func @collapse_and_slice_multiple_unit_dim_mixed(%input: tensor<1x?x1x100x10xf32>, %offt: index, %size: index) -> tensor<?x?xf32> {199 %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3, 4]] : tensor<1x?x1x100x10xf32> into tensor<?x1000xf32>200 %slice = tensor.extract_slice %collapsed [%offt, %offt] [%size, %size] [1, 1] : tensor<?x1000xf32> to tensor<?x?xf32>201 // CHECK-DAG: %[[c0]] = arith.constant 0 : index202 // CHECK-DAG: %[[c1]] = arith.constant 1 : index203 // CHECK: %[[dim:.+]] = tensor.dim %[[arg0]], %[[c1]]204 // CHECK: %[[rank_reduced:.+]] = tensor.extract_slice %[[arg0]][0, 0, 0, 0, 0] [1, %[[dim]], 1, 100, 10] [1, 1, 1, 1, 1]205 // CHECK: %[[empty:.+]] = tensor.empty206 // CHECK: %[[result:.+]] = scf.for %[[iv:.+]] = %[[c0]] to %[[arg2]] step %[[c1]] iter_args(%[[ia:.+]] = %[[empty]])207 // CHECK: %[[idx:.+]] = affine.apply208 // CHECK: %[[multi_index:.+]] = affine.delinearize_index %[[idx]] into209 // CHECK: %[[collapsed:.+]] = tensor.collapse_shape210 // CHECK: %[[updated:.+]] = tensor.insert_slice211 // CHECK: scf.yield %[[updated]]212 // CHECK: return %[[result]]213 return %slice : tensor<?x?xf32>214}215 216// Edge case where all collapsed dims are unit dims. This pattern can't eliminate the collapse shape,217// that should be handled by `linalg-fold-unit-extent-dims`.218 219// CHECK: @collapse_and_slice_multiple_all_unit_dim(%[[arg0:.+]]: tensor<{{.*}}>)220func.func @collapse_and_slice_multiple_all_unit_dim(%input: tensor<1x1x1x100xf32>) -> tensor<1x100xf32> {221 %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3]] : tensor<1x1x1x100xf32> into tensor<1x100xf32>222 %slice = tensor.extract_slice %collapsed [0, 0] [1, 100] [1, 1] : tensor<1x100xf32> to tensor<1x100xf32>223 return %slice : tensor<1x100xf32>224 // CHECK: %[[collapse:.+]] = tensor.collapse_shape %[[arg0]] {{\[}}[0, 1, 2], [3]] : tensor<1x1x1x100xf32> into tensor<1x100xf32>225 // CHECK: return %[[collapse]]226}227