146 lines · plain
1// RUN: mlir-opt -transform-interpreter %s --split-input-file --allow-unregistered-dialect | FileCheck %s2 3// CHECK-LABEL: func @insert_slice_to_copy4 // CHECK-SAME: %[[I:.*]]: tensor<2x3xf32>5 // CHECK-SAME: %[[O:.*]]: tensor<?x?xf32>,6 // CHECK-SAME: %[[OFF0:[0-9a-zA-Z]+]]: index,7 // CHECK-SAME: %[[OFF1:[0-9a-zA-Z]+]]: index,8 // CHECK-SAME: %[[SZ0:[0-9a-zA-Z]+]]: index,9 // CHECK-SAME: %[[SZ1:[0-9a-zA-Z]+]]: index,10 // CHECK-SAME: %[[ST0:[0-9a-zA-Z]+]]: index,11 // CHECK-SAME: %[[ST1:[0-9a-zA-Z]+]]: index)12func.func @insert_slice_to_copy(13 %I : tensor<2x3xf32>, %O : tensor<?x?xf32>,14 %off0 : index, %off1 : index,15 %sz0 : index, %sz1 : index,16 %st0 : index, %st1 : index) -> tensor<?x?xf32> {17 18 // CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[O]][%[[OFF0]], %[[OFF1]]] [2, 3] [%[[ST0]], %[[ST1]]]19 // CHECK-SAME: : tensor<?x?xf32> to tensor<2x3xf32>20 // CHECK: linalg.copy ins(%[[I]] : tensor<2x3xf32>) outs(%[[EXTRACTED_SLICE]] : tensor<2x3xf32>) -> tensor<2x3xf32>21 // CHECK: tensor.insert_slice %{{.*}} into %[[O]][%[[OFF0]], %[[OFF1]]] [2, 3] [%[[ST0]], %[[ST1]]]22 // CHECK-SAME: : tensor<2x3xf32> into tensor<?x?xf32>23 24 %0 = tensor.insert_slice %I into %O[%off0, %off1] [2, 3] [%st0, %st1]25 : tensor<2x3xf32> into tensor<?x?xf32>26 return %0 : tensor<?x?xf32>27}28 29module attributes {transform.with_named_sequence} {30 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {31 %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op32 %1 = transform.structured.insert_slice_to_copy %0 : (!transform.any_op) -> !transform.any_op33 transform.cast %1 : !transform.any_op to !transform.op<"linalg.copy">34 transform.yield35 }36}37 38// -----39 40// CHECK-LABEL: func @insert_slice_to_copy41 // CHECK-SAME: %[[I:[0-9a-zA-Z]+]]: tensor<?x?xf32>42 // CHECK-SAME: %[[O:[0-9a-zA-Z]+]]: tensor<?x?xf32>,43 // CHECK-SAME: %[[OFF0:[0-9a-zA-Z]+]]: index,44 // CHECK-SAME: %[[OFF1:[0-9a-zA-Z]+]]: index,45 // CHECK-SAME: %[[SZ0:[0-9a-zA-Z]+]]: index,46 // CHECK-SAME: %[[SZ1:[0-9a-zA-Z]+]]: index,47 // CHECK-SAME: %[[ST0:[0-9a-zA-Z]+]]: index,48 // CHECK-SAME: %[[ST1:[0-9a-zA-Z]+]]: index)49func.func @insert_slice_to_copy(50 %I : tensor<?x?xf32>, %O : tensor<?x?xf32>,51 %off0 : index, %off1 : index,52 %sz0 : index, %sz1 : index,53 %st0 : index, %st1 : index) -> tensor<?x?xf32> {54 55 // CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[O]][%[[OFF0]], %[[OFF1]]] [%[[SZ0]], %[[SZ1]]] [1, 1]56 // CHECK-SAME: : tensor<?x?xf32> to tensor<?x?xf32>57 // CHECK: linalg.copy ins(%[[I]] : tensor<?x?xf32>) outs(%[[EXTRACTED_SLICE]] : tensor<?x?xf32>) -> tensor<?x?xf32>58 // CHECK: tensor.insert_slice %{{.*}} into %[[O]][%[[OFF0]], %[[OFF1]]] [%[[SZ0]], %[[SZ1]]] [1, 1]59 // CHECK-SAME: : tensor<?x?xf32> into tensor<?x?xf32>60 61 %0 = tensor.insert_slice %I into %O[%off0, %off1] [%sz0, %sz1] [1, 1]62 : tensor<?x?xf32> into tensor<?x?xf32>63 return %0 : tensor<?x?xf32>64}65 66module attributes {transform.with_named_sequence} {67 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {68 %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op69 %1 = transform.structured.insert_slice_to_copy %0 : (!transform.any_op) -> !transform.any_op70 transform.cast %1 : !transform.any_op to !transform.op<"linalg.copy">71 transform.yield72 }73}74 75// -----76// CHECK-LABEL: func @insert_slice_to_copy77 // CHECK-SAME: %[[I:.*]]: tensor<2x3xf32>78 // CHECK-SAME: %[[O:.*]]: tensor<?x?xf32>,79 // CHECK-SAME: %[[OFF0:[0-9a-zA-Z]+]]: index,80 // CHECK-SAME: %[[OFF1:[0-9a-zA-Z]+]]: index,81 // CHECK-SAME: %[[SZ0:[0-9a-zA-Z]+]]: index,82 // CHECK-SAME: %[[SZ1:[0-9a-zA-Z]+]]: index,83 // CHECK-SAME: %[[ST0:[0-9a-zA-Z]+]]: index,84 // CHECK-SAME: %[[ST1:[0-9a-zA-Z]+]]: index)85func.func @insert_slice_to_copy(86 %I : tensor<2x3xf32>, %O : tensor<?x?xf32>,87 %off0 : index, %off1 : index,88 %sz0 : index, %sz1 : index,89 %st0 : index, %st1 : index) -> tensor<?x?xf32> {90 91 // CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[O]][%[[OFF0]], %[[OFF1]]] [2, 3] [%[[ST0]], %[[ST1]]]92 // CHECK-SAME: : tensor<?x?xf32> to tensor<2x3xf32>93 // CHECK: linalg.copy ins(%[[I]] : tensor<2x3xf32>) outs(%[[EXTRACTED_SLICE]] : tensor<2x3xf32>) -> tensor<2x3xf32>94 // CHECK-NOT: linalg.copy95 // CHECK: tensor.insert_slice %{{.*}} into %[[O]][%[[OFF0]], %[[OFF1]]] [2, 3] [%[[ST0]], %[[ST1]]]96 // CHECK-SAME: : tensor<2x3xf32> into tensor<?x?xf32>97 98 %extracted_slice = tensor.extract_slice %O[%off0, %off1] [2, 3] [%st0, %st1]99 : tensor<?x?xf32> to tensor<2x3xf32>100 %0 = linalg.copy ins(%I : tensor<2x3xf32>) outs(%extracted_slice101 : tensor<2x3xf32>) -> tensor<2x3xf32>102 %inserted_slice = tensor.insert_slice %0 into %O[%off0, %off1] [2, 3] [%st0, %st1]103 : tensor<2x3xf32> into tensor<?x?xf32>104 105 return %inserted_slice : tensor<?x?xf32>106}107 108module attributes {transform.with_named_sequence} {109 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {110 %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op111 %1 = transform.structured.insert_slice_to_copy %0 : (!transform.any_op) -> !transform.any_op112 transform.cast %1 : !transform.any_op to !transform.op<"linalg.copy">113 transform.yield114 }115}116 117// -----118 119// CHECK-LABEL: func @parallel_insert_slice_to_copy120func.func @parallel_insert_slice_to_copy(%out : tensor<?x?xf32>, %sz0: index, %sz1: index) {121 %0 = scf.forall (%arg0, %arg1) in (27, 8) shared_outs(%arg2 = %out) -> (tensor<?x?xf32>) {122 %t = "make_me_a_tensor"() : () -> (tensor<?x?xf32> )123 124 // CHECK: tensor.extract_slice125 // CHECK: linalg.copy126 // CHECK: scf.forall.in_parallel127 // CHECK: tensor.parallel_insert_slice128 scf.forall.in_parallel {129 tensor.parallel_insert_slice %t into %arg2[0, 0] [%sz0, %sz1] [1, 1]130 : tensor<?x?xf32> into tensor<?x?xf32>131 }132 }133 return134}135 136module attributes {transform.with_named_sequence} {137 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {138 %0 = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg1139 : (!transform.any_op) -> !transform.any_op140 %1 = transform.structured.insert_slice_to_copy %0141 : (!transform.any_op) -> !transform.any_op142 transform.cast %1 : !transform.any_op to !transform.op<"linalg.copy">143 transform.yield144 }145}146