172 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -split-input-file | FileCheck %s2 3///----------------------------------------------------------------------------------------4/// Tests for tensor.insert_slice5///----------------------------------------------------------------------------------------6 7func.func private @insert_slice_static_sizes(%source: tensor<?x3x?x1xi32>) -> tensor<5x3xi32> {8 %c2 = arith.constant 2 : index9 %init = tensor.empty() : tensor<5x3xi32>10 11 %source_slice = tensor.extract_slice %source[0, %c2, 0, 0] [1, 1, 5, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<5x1xi32>12 %res = tensor.insert_slice %source_slice into %init[0, %c2] [5, 1] [1, 1] : tensor<5x1xi32> into tensor<5x3xi32>13 14 return %res : tensor<5x3xi32>15}16 17// CHECK-LABEL: func.func private @insert_slice_static_sizes(18// CHECK-SAME: %[[SEC:.*]]: tensor<?x3x?x1xi32>) -> tensor<5x3xi32> {19// CHECK: %[[C_2:.*]] = arith.constant 2 : index20// CHECK: %[[INIT:.*]] = tensor.empty() : tensor<5x3xi32>21// CHECK: %[[SRC_SLICE:.*]] = tensor.extract_slice %[[SEC]][0, %[[C_2]], 0, 0] [1, 1, 5, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<5x1xi32>22// CHECK-DAG: %[[PAD:.*]] = arith.constant 0 : i3223// CHECK-DAG: %[[C_0_1:.*]] = arith.constant 0 : index24// CHECK-DAG: %[[C_0:.*]] = arith.constant 0 : index25// CHECK-DAG: %[[C_5:.*]] = arith.constant 5 : index26// CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index27// CHECK: %[[MASK_READ:.*]] = vector.create_mask %[[C_5]], %[[C_1]] : vector<8x1xi1>28// CHECK: %[[READ:.*]] = vector.mask %[[MASK_READ]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0]], %[[C_0]]], %[[PAD]] {{.*}} : tensor<5x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>29// CHECK: %[[C_0_1:.*]] = arith.constant 0 : index30// CHECK: %[[C_5_1:.*]] = arith.constant 5 : index31// CHECK: %[[C_3:.*]] = arith.constant 3 : index32// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[C_5_1]], %[[C_3]] : vector<8x1xi1>33// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_1]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> tensor<5x3xi32>34// CHECK: return %[[RES]] : tensor<5x3xi32>35 36module attributes {transform.with_named_sequence} {37 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {38 %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op39 transform.structured.vectorize %0 vector_sizes [8, 1] : !transform.any_op40 transform.yield41 }42}43 44// -----45 46// One of the _source_ dimensions is dynamic (but _destination_ dimensions are static).47 48func.func private @insert_slice_dynamic_src_dim(%source: tensor<?x3x?x1xi32>, %size: index) -> tensor<5x3xi32> {49 %c2 = arith.constant 2 : index50 %init = tensor.empty() : tensor<5x3xi32>51 52 %source_slice = tensor.extract_slice %source[0, %c2, 0, 0] [1, 1, %size, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<?x1xi32>53 %res = tensor.insert_slice %source_slice into %init[0, %c2] [%size, 1] [1, 1] : tensor<?x1xi32> into tensor<5x3xi32>54 55 return %res : tensor<5x3xi32>56}57 58// CHECK-LABEL: func.func private @insert_slice_dynamic_src_dim(59// CHECK-SAME: %[[SRC:.*]]: tensor<?x3x?x1xi32>,60// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<5x3xi32> {61// CHECK: %[[C_2:.*]] = arith.constant 2 : index62// CHECK: %[[INIT:.*]] = tensor.empty() : tensor<5x3xi32>63// CHECK: %[[SRC_SLICE:.*]] = tensor.extract_slice %[[SRC]][0, %[[C_2]], 0, 0] [1, 1, %[[SIZE]], 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<?x1xi32>64// CHECK-DAG: %[[PAD:.*]] = arith.constant 0 : i3265// CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index66// CHECK-DAG: %[[C_0:.*]] = arith.constant 0 : index67// CHECK-DAG: %[[C_0_1:.*]] = arith.constant 0 : index68// CHECK-DAG: %[[C_0_2:.*]] = arith.constant 0 : index69// CHECK-DAG: %[[D0:.*]] = tensor.dim %[[SRC_SLICE]], %[[C_0_2]] : tensor<?x1xi32>70// CHECK: %[[MASK:.*]] = vector.create_mask %[[D0]], %[[C_1]] : vector<8x1xi1>71// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0_1]], %[[C_0_1]]], %[[PAD]] {{.*}} : tensor<?x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>72// CHECK: %[[C_0_1:.*]] = arith.constant 0 : index73// CHECK: %[[C_5_1:.*]] = arith.constant 5 : index74// CHECK: %[[C_3:.*]] = arith.constant 3 : index75// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[C_5_1]], %[[C_3]] : vector<8x1xi1>76// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_1]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> tensor<5x3xi32>77// CHECK: return %[[RES]] : tensor<5x3xi32>78 79 module attributes {transform.with_named_sequence} {80 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {81 %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op82 transform.structured.vectorize %0 vector_sizes [8, 1] : !transform.any_op83 transform.yield84 }85 }86 87// -----88 89// One of the _destination_ dimensions is dynamic (but _source_ dimensions are static).90 91func.func private @insert_slice_dynamic_dest_dim(%source: tensor<?x3x?x1xi32>, %size: index) -> tensor<?x3xi32> {92 %c2 = arith.constant 2 : index93 %init = tensor.empty(%size) : tensor<?x3xi32>94 95 %source_slice = tensor.extract_slice %source[0, %c2, 0, 0] [1, 1, 5, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<5x1xi32>96 %res = tensor.insert_slice %source_slice into %init[0, %c2] [5, 1] [1, 1] : tensor<5x1xi32> into tensor<?x3xi32>97 98 return %res : tensor<?x3xi32>99}100 101// CHECK-LABEL: func.func private @insert_slice_dynamic_dest_dim(102// CHECK-SAME: %[[SRC:.*]]: tensor<?x3x?x1xi32>,103// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?x3xi32> {104// CHECK: %[[C_2:.*]] = arith.constant 2 : index105// CHECK: %[[INIT:.*]] = tensor.empty(%[[SIZE]]) : tensor<?x3xi32>106// CHECK: %[[SRC_SLICE:.*]] = tensor.extract_slice %[[SRC]][0, %[[C_2]], 0, 0] [1, 1, 5, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<5x1xi32>107// CHECK-DAG: %[[PAD:.*]] = arith.constant 0 : i32108// CHECK-DAG: %[[C_5:.*]] = arith.constant 5 : index109// CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index110// CHECK-DAG: %[[MASK:.*]] = vector.create_mask %[[C_5]], %[[C_1]] : vector<8x1xi1>111// CHECK-DAG: %[[C_0:.*]] = arith.constant 0 : index112// CHECK-DAG: %[[C_0_1:.*]] = arith.constant 0 : index113// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0_1]], %[[C_0_1]]], %[[PAD]] {{.*}} : tensor<5x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>114// CHECK: %[[C_0_1:.*]] = arith.constant 0 : index115// CHECK: %[[C_0_2:.*]] = arith.constant 0 : index116// CHECK: %[[DIM:.*]] = tensor.dim %[[INIT]], %[[C_0_2]] : tensor<?x3xi32>117// CHECK: %[[C_3:.*]] = arith.constant 3 : index118// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[DIM]], %[[C_3]] : vector<8x1xi1>119// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_1]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<?x3xi32> } : vector<8x1xi1> -> tensor<?x3xi32>120// CHECK: return %[[RES]] : tensor<?x3xi32>121 122 module attributes {transform.with_named_sequence} {123 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {124 %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op125 transform.structured.vectorize %0 vector_sizes [8, 1] : !transform.any_op126 transform.yield127 }128 }129 130// -----131 132// At least one _source_ and one _destination_ dimensions are dynamic.133 134func.func private @insert_slice_dynamic_source_and_dest_dim(%source: tensor<?x3x?x1xi32>, %size: index) -> tensor<?x3xi32> {135 %c2 = arith.constant 2 : index136 %init = tensor.empty(%size) : tensor<?x3xi32>137 138 %source_slice = tensor.extract_slice %source[0, %c2, 0, 0] [1, 1, %size, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<?x1xi32>139 %res = tensor.insert_slice %source_slice into %init[0, %c2] [%size, 1] [1, 1] : tensor<?x1xi32> into tensor<?x3xi32>140 141 return %res : tensor<?x3xi32>142}143 144// CHECK-LABEL: func.func private @insert_slice_dynamic_source_and_dest_dim(145// CHECK-SAME: %[[SRC:.*]]: tensor<?x3x?x1xi32>,146// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?x3xi32> {147// CHECK: %[[C_2:.*]] = arith.constant 2 : index148// CHECK: %[[INIT:.*]] = tensor.empty(%[[SIZE]]) : tensor<?x3xi32>149// CHECK: %[[SRC_SLICE:.*]] = tensor.extract_slice %[[SRC]][0, %[[C_2]], 0, 0] [1, 1, %[[SIZE]], 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<?x1xi32>150// CHECK-DAG: %[[PAD:.*]] = arith.constant 0 : i32151// CHECK-DAG: %[[C0_0:.*]] = arith.constant 0 : index152// CHECK-DAG: %[[C0_1:.*]] = arith.constant 0 : index153// CHECK-DAG: %[[C0_2:.*]] = arith.constant 0 : index154// CHECK: %[[D0:.*]] = tensor.dim %[[SRC_SLICE]], %[[C0_2]] : tensor<?x1xi32>155// CHECK: %[[C1:.*]] = arith.constant 1 : index156// CHECK: %[[MASK:.*]] = vector.create_mask %[[D0]], %[[C1]] : vector<8x1xi1>157// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C0_1]], %[[C0_1]]], %[[PAD]] {{.*}} : tensor<?x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>158// CHECK: %[[C_0_1:.*]] = arith.constant 0 : index159// CHECK: %[[C_0_2:.*]] = arith.constant 0 : index160// CHECK: %[[DIM:.*]] = tensor.dim %[[INIT]], %[[C_0_2]] : tensor<?x3xi32>161// CHECK: %[[C_3:.*]] = arith.constant 3 : index162// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[DIM]], %[[C_3]] : vector<8x1xi1>163// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_1]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<?x3xi32> } : vector<8x1xi1> -> tensor<?x3xi32>164 165 module attributes {transform.with_named_sequence} {166 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {167 %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op168 transform.structured.vectorize %0 vector_sizes [8, 1] : !transform.any_op169 transform.yield170 }171 }172