91 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -split-input-file | FileCheck %s2 3///----------------------------------------------------------------------------------------4/// Tests for tensor.insert_slice5///----------------------------------------------------------------------------------------6 7// The pad value for xfer-read is neither needed nor available - use the default (0.0).8 9// CHECK-LABEL: func @insert_static_slice_default_pad10// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x2x3xf32>,11// CHECK-SAME: %[[ARG_1:.*]]: tensor<9x8x7x1x2x3xf32>) -> tensor<9x8x7x1x2x3xf32> {12// CHECK: %[[PAD:.*]] = arith.constant 0.000000e+00 : f3213// CHECK: %[[C0:.*]] = arith.constant 0 : index14// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG_0]]{{\[}}%[[C0]], %[[C0]], %[[C0]]], %[[PAD]] {in_bounds = [true, true, true]} : tensor<1x2x3xf32>, vector<1x2x3xf32>15// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[READ]], %[[ARG_1]]{{\[}}%[[C0]], %[[C0]], %[[C0]], %[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x2x3xf32>, tensor<9x8x7x1x2x3xf32>16// CHECK: return %[[WRITE]] : tensor<9x8x7x1x2x3xf32>17func.func @insert_static_slice_default_pad(%arg0: tensor<1x2x3xf32>, %arg1: tensor<9x8x7x1x2x3xf32>) -> tensor<9x8x7x1x2x3xf32> {18 %res = tensor.insert_slice %arg0 into %arg1[0, 0, 0, 0, 0, 0] [1, 1, 1, 1, 2, 3][1, 1, 1, 1, 1, 1] : tensor<1x2x3xf32> into tensor<9x8x7x1x2x3xf32>19 return %res : tensor<9x8x7x1x2x3xf32>20}21 22module attributes {transform.with_named_sequence} {23 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {24 %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op25 %1 = transform.get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op26 %2 = transform.structured.vectorize_children_and_apply_patterns %1 { vectorize_padding } : (!transform.any_op) -> !transform.any_op27 transform.yield28 }29}30 31// -----32 33// Same as above, but there's a pad value available that should be used instead of the default value.34 35// CHECK-LABEL: func.func @insert_static_slice_non_zero_pad36// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x2x3xf32>,37// CHECK-SAME: %[[PAD:.*]]: f32) -> tensor<9x8x7x1x2x3xf32> {38// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<9x8x7x1x2x3xf32>39// CHECK: %[[BC:.*]] = vector.broadcast %[[PAD]] : f32 to vector<9x8x7x1x2x3xf32>40// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[BC]], %[[EMPTY]]{{.*}} {in_bounds = [true, true, true, true, true, true]} : vector<9x8x7x1x2x3xf32>, tensor<9x8x7x1x2x3xf32>41// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG_0]]{{.*}}, %[[PAD]] {in_bounds = [true, true, true]} : tensor<1x2x3xf32>, vector<1x2x3xf32>42// CHECK: %[[RES:.*]] = vector.transfer_write %[[READ]], %[[WRITE]]{{.*}} {in_bounds = [true, true, true]} : vector<1x2x3xf32>, tensor<9x8x7x1x2x3xf32>43// CHECK: return %[[RES]] : tensor<9x8x7x1x2x3xf32>44func.func @insert_static_slice_non_zero_pad(%arg0: tensor<1x2x3xf32>, %pad : f32) -> tensor<9x8x7x1x2x3xf32> {45 %init = tensor.empty() : tensor<9x8x7x1x2x3xf32>46 %fill = linalg.fill ins(%pad : f32) outs(%init : tensor<9x8x7x1x2x3xf32>) -> tensor<9x8x7x1x2x3xf32>47 %res = tensor.insert_slice %arg0 into %fill[0, 0, 0, 0, 0, 0] [1, 1, 1, 1, 2, 3][1, 1, 1, 1, 1, 1] : tensor<1x2x3xf32> into tensor<9x8x7x1x2x3xf32>48 return %res : tensor<9x8x7x1x2x3xf32>49}50 51module attributes {transform.with_named_sequence} {52 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {53 %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op54 %1 = transform.get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op55 %2 = transform.structured.vectorize_children_and_apply_patterns %1 : (!transform.any_op) -> !transform.any_op56 transform.yield57 }58}59 60// -----61 62// Same as above, but the source type has is dynamically shaped. This means63// that the pad value is now required and the vector dim corresponding to the64// dynamic shape has to be inferred from the shape of the destination tensor.65 66// CHECK-LABEL: func.func @insert_dynamic_slice_non_zero_pad(67// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x?x3xf32>,68// CHECK-SAME: %[[PAD:.*]]: f32,69// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<9x8x7x1x2x3xf32> {70// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<9x8x7x1x2x3xf32>71// CHECK: %[[BC:.*]] = vector.broadcast %[[PAD]] : f32 to vector<9x8x7x1x2x3xf32>72// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[BC]], %[[EMPTY]]{{.*}} {in_bounds = [true, true, true, true, true, true]} : vector<9x8x7x1x2x3xf32>, tensor<9x8x7x1x2x3xf32>73// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG_0]]{{.*}}, %[[PAD]] {in_bounds = [true, false, true]} : tensor<1x?x3xf32>, vector<1x2x3xf32>74// CHECK: %[[RES:.*]] = vector.transfer_write %[[READ]], %[[WRITE]]{{.*}} {in_bounds = [true, true, true]} : vector<1x2x3xf32>, tensor<9x8x7x1x2x3xf32>75// CHECK: return %[[RES]] : tensor<9x8x7x1x2x3xf32>76func.func @insert_dynamic_slice_non_zero_pad(%arg0: tensor<1x?x3xf32>, %pad : f32, %size: index) -> tensor<9x8x7x1x2x3xf32> {77 %init = tensor.empty() : tensor<9x8x7x1x2x3xf32>78 %fill = linalg.fill ins(%pad : f32) outs(%init : tensor<9x8x7x1x2x3xf32>) -> tensor<9x8x7x1x2x3xf32>79 %res = tensor.insert_slice %arg0 into %fill[0, 0, 0, 0, 0, 0] [1, 1, 1, 1, %size, 3][1, 1, 1, 1, 1, 1] : tensor<1x?x3xf32> into tensor<9x8x7x1x2x3xf32>80 return %res : tensor<9x8x7x1x2x3xf32>81}82 83module attributes {transform.with_named_sequence} {84 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {85 %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op86 %1 = transform.get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op87 %2 = transform.structured.vectorize_children_and_apply_patterns %1 : (!transform.any_op) -> !transform.any_op88 transform.yield89 }90}91