1// RUN: mlir-opt -split-input-file -transform-interpreter %s | FileCheck %s 2 3module attributes {transform.with_named_sequence} { 4 transform.named_sequence @__transform_main(%root : !transform.any_op {transform.readonly}) { 5 %func_op = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.op<"func.func"> 6 transform.apply_patterns to %func_op { 7 transform.apply_patterns.tensor.fold_tensor_subset_ops_into_vector_transfers 8 } : !transform.op<"func.func"> 9 transform.yield 10 } 11} 12 13// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 + 4)> 14// CHECK: #[[$map1:.*]] = affine_map<()[s0] -> (s0 + 3)> 15// CHECK: #[[$map2:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)> 16 17// CHECK-LABEL: func @transfer_read_of_extract_slice( 18// CHECK-SAME: %[[t:.*]]: tensor<?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index 19// CHECK-DAG: %[[c8:.*]] = arith.constant 8 : index 20// CHECK: %[[add:.*]] = affine.apply #[[$map]]()[%[[s1]]] 21// CHECK: %[[r:.*]] = vector.transfer_read %[[t]][%[[c8]], %[[add]]], %{{.*}} {in_bounds = [true, true]} : tensor<?x?xf32>, vector<5x6xf32> 22// CHECK: return %[[r]] 23func.func @transfer_read_of_extract_slice(%t : tensor<?x?xf32>, %s1 : index, %s2 : index) -> vector<5x6xf32> { 24 %c3 = arith.constant 3 : index 25 %c4 = arith.constant 4 : index 26 %cst = arith.constant 0.0 : f32 27 %0 = tensor.extract_slice %t[5, %s1] [10, %s2] [1, 1] : tensor<?x?xf32> to tensor<10x?xf32> 28 %1 = vector.transfer_read %0[%c3, %c4], %cst {in_bounds = [true, true]} : tensor<10x?xf32>, vector<5x6xf32> 29 return %1 : vector<5x6xf32> 30} 31 32// CHECK-LABEL: func @transfer_read_of_extract_slice_1d( 33// CHECK-SAME: %[[t:.*]]: tensor<?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index 34// CHECK-DAG: %[[c8:.*]] = arith.constant 8 : index 35// CHECK: %[[add:.*]] = affine.apply #[[$map]]()[%[[s1]]] 36// CHECK: %[[r:.*]] = vector.transfer_read %[[t]][%[[c8]], %[[add]]], %{{.*}} {in_bounds = [true]} : tensor<?x?xf32>, vector<6xf32> 37// CHECK: return %[[r]] 38func.func @transfer_read_of_extract_slice_1d(%t : tensor<?x?xf32>, %s1 : index, %s2 : index) -> vector<6xf32> { 39 %c3 = arith.constant 3 : index 40 %c4 = arith.constant 4 : index 41 %cst = arith.constant 0.0 : f32 42 %0 = tensor.extract_slice %t[5, %s1] [10, %s2] [1, 1] : tensor<?x?xf32> to tensor<10x?xf32> 43 %1 = vector.transfer_read %0[%c3, %c4], %cst {in_bounds = [true]} : tensor<10x?xf32>, vector<6xf32> 44 return %1 : vector<6xf32> 45} 46 47// CHECK-LABEL: func @transfer_read_of_extract_slice_rank_reducing( 48// CHECK-SAME: %[[t:.*]]: tensor<?x?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index 49// CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index 50// CHECK-DAG: %[[c10:.*]] = arith.constant 10 : index 51// CHECK: %[[add:.*]] = affine.apply #[[$map1]]()[%[[s1]]] 52// CHECK: %[[r:.*]] = vector.transfer_read %[[t]][%[[c5]], %[[add]], %[[c10]]], %{{.*}} {in_bounds = [true, true]} : tensor<?x?x?xf32>, vector<5x6xf32> 53// CHECK: return %[[r]] 54func.func @transfer_read_of_extract_slice_rank_reducing(%t : tensor<?x?x?xf32>, %s1 : index, %s2 : index) -> vector<5x6xf32> { 55 %c3 = arith.constant 3 : index 56 %c4 = arith.constant 4 : index 57 %cst = arith.constant 0.0 : f32 58 %0 = tensor.extract_slice %t[5, %s1, 6] [1, %s2, 12] [1, 1, 1] : tensor<?x?x?xf32> to tensor<?x12xf32> 59 %1 = vector.transfer_read %0[%c3, %c4], %cst {in_bounds = [true, true]} : tensor<?x12xf32>, vector<5x6xf32> 60 return %1 : vector<5x6xf32> 61} 62 63// CHECK-LABEL: func @transfer_read_of_extract_slice_non_leading_rank_reduction( 64// CHECK-SAME: %[[t:.*]]: tensor<?x?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index 65// CHECK-DAG: %[[c8:.*]] = arith.constant 8 : index 66// CHECK-DAG: %[[c10:.*]] = arith.constant 10 : index 67// CHECK: %[[r:.*]] = vector.transfer_read %[[t]][%[[c8]], %[[s1]], %[[c10]]], %{{.*}} {in_bounds = [true, true], permutation_map = #[[$map2]]} : tensor<?x?x?xf32>, vector<5x6xf32> 68// CHECK: return %[[r]] 69func.func @transfer_read_of_extract_slice_non_leading_rank_reduction(%t : tensor<?x?x?xf32>, %s1 : index, %s2 : index) -> vector<5x6xf32> { 70 %c3 = arith.constant 3 : index 71 %c4 = arith.constant 4 : index 72 %cst = arith.constant 0.0 : f32 73 %0 = tensor.extract_slice %t[5, %s1, 6] [%s2, 1, 12] [1, 1, 1] : tensor<?x?x?xf32> to tensor<?x12xf32> 74 %1 = vector.transfer_read %0[%c3, %c4], %cst {in_bounds = [true, true]} : tensor<?x12xf32>, vector<5x6xf32> 75 return %1 : vector<5x6xf32> 76} 77 78// CHECK-LABEL: func @masked_transfer_read_of_extract_slice 79// CHECK-SAME: %[[t:.*]]: tensor<?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index 80// CHECK-DAG: %[[m:.*]] = vector.create_mask{{.*}} : vector<5x6xi1> 81// CHECK-DAG: %[[a:.*]] = affine.apply {{.*}}[[s1]] 82// CHECK: vector.mask %[[m]] { vector.transfer_read %[[t]]{{.*}}: tensor<?x?xf32>, vector<5x6xf32> } : vector<5x6xi1> -> vector<5x6xf32> 83func.func @masked_transfer_read_of_extract_slice(%t : tensor<?x?xf32>, %s1 : index, %s2 : index) -> vector<5x6xf32> { 84 %c3 = arith.constant 3 : index 85 %c4 = arith.constant 4 : index 86 %cst = arith.constant 0.0 : f32 87 %0 = tensor.extract_slice %t[5, %s1] [10, %s2] [1, 1] : tensor<?x?xf32> to tensor<10x?xf32> 88 %mask = vector.create_mask %c3, %c4 : vector<5x6xi1> 89 %1 = vector.mask %mask {vector.transfer_read %0[%c3, %c4], %cst {in_bounds = [true, true]} : tensor<10x?xf32>, vector<5x6xf32>} : vector<5x6xi1> -> vector<5x6xf32> 90 return %1 : vector<5x6xf32> 91} 92 93// CHECK-LABEL: func @insert_slice_of_transfer_write( 94// CHECK-SAME: %[[t1:.*]]: tensor<?x12xf32>, %[[v:.*]]: vector<5x6xf32>, %[[s:.*]]: index 95// CHECK: %[[c3:.*]] = arith.constant 3 : index 96// CHECK: %[[r:.*]] = vector.transfer_write %[[v]], %[[t1]][%[[c3]], %[[s]]] {in_bounds = [true, true]} : vector<5x6xf32>, tensor<?x12xf32> 97// CHECK: return %[[r]] 98func.func @insert_slice_of_transfer_write(%t1 : tensor<?x12xf32>, %v : vector<5x6xf32>, %s : index, %t2 : tensor<5x6xf32>) -> tensor<?x12xf32> { 99 %c0 = arith.constant 0 : index 100 %0 = vector.transfer_write %v, %t2[%c0, %c0] {in_bounds = [true, true]} : vector<5x6xf32>, tensor<5x6xf32> 101 %1 = tensor.insert_slice %0 into %t1[3, %s] [5, 6] [1, 1] : tensor<5x6xf32> into tensor<?x12xf32> 102 return %1 : tensor<?x12xf32> 103} 104 105// CHECK-LABEL: func @unit_insert_slice_of_unit_transfer_write( 106// CHECK-SAME: %[[t1:.*]]: tensor<1x1x12xf32>, %[[v:.*]]: vector<1x6xf32>, %[[s:.*]]: index 107// CHECK: %[[c0:.*]] = arith.constant 0 : index 108// CHECK: %[[r:.*]] = vector.transfer_write %[[v]], %[[t1]][%[[c0]], %[[c0]], %[[s]]] {in_bounds = [true, true]} : vector<1x6xf32>, tensor<1x1x12xf32> 109// CHECK: return %[[r]] 110func.func @unit_insert_slice_of_unit_transfer_write(%t1 : tensor<1x1x12xf32>, %v : vector<1x6xf32>, %s : index, %t2 : tensor<1x6xf32>) -> tensor<1x1x12xf32> { 111 %c0 = arith.constant 0 : index 112 %0 = vector.transfer_write %v, %t2[%c0, %c0] {in_bounds = [true, true]} : vector<1x6xf32>, tensor<1x6xf32> 113 %1 = tensor.insert_slice %0 into %t1[0, 0, %s] [1, 1, 6] [1, 1, 1] : tensor<1x6xf32> into tensor<1x1x12xf32> 114 return %1 : tensor<1x1x12xf32> 115} 116 117// CHECK-LABEL: func @insert_slice_of_transfer_write_non_leading_rank_reduction( 118// CHECK-SAME: %[[t1:.*]]: tensor<?x?x12xf32>, %[[v:.*]]: vector<5x6xf32>, %[[s:.*]]: index 119// CHECK-DAG: %[[c3:.*]] = arith.constant 3 : index 120// CHECK-DAG: %[[c4:.*]] = arith.constant 4 : index 121// CHECK: %[[r:.*]] = vector.transfer_write %[[v]], %[[t1]][%[[c4]], %[[c3]], %[[s]]] {in_bounds = [true, true], permutation_map = #[[$map2]]} : vector<5x6xf32>, tensor<?x?x12xf32> 122func.func @insert_slice_of_transfer_write_non_leading_rank_reduction(%t1 : tensor<?x?x12xf32>, %v : vector<5x6xf32>, %s : index, %t2 : tensor<5x6xf32>) -> tensor<?x?x12xf32> { 123 %c0 = arith.constant 0 : index 124 %0 = vector.transfer_write %v, %t2[%c0, %c0] {in_bounds = [true, true]} : vector<5x6xf32>, tensor<5x6xf32> 125 %1 = tensor.insert_slice %0 into %t1[4, 3, %s] [5, 1, 6] [1, 1, 1] : tensor<5x6xf32> into tensor<?x?x12xf32> 126 return %1 : tensor<?x?x12xf32> 127} 128 129// CHECK-LABEL: func @insert_slice_of_transfer_write_rank_extending( 130// CHECK-SAME: %[[t1:.*]]: tensor<?x?x12xf32>, %[[v:.*]]: vector<5x6xf32>, %[[s:.*]]: index 131// CHECK-DAG: %[[c3:.*]] = arith.constant 3 : index 132// CHECK-DAG: %[[c4:.*]] = arith.constant 4 : index 133// CHECK: %[[r:.*]] = vector.transfer_write %[[v]], %[[t1]][%[[c4]], %[[c3]], %[[s]]] {in_bounds = [true, true]} : vector<5x6xf32>, tensor<?x?x12xf32> 134// CHECK: return %[[r]] 135func.func @insert_slice_of_transfer_write_rank_extending(%t1 : tensor<?x?x12xf32>, %v : vector<5x6xf32>, %s : index, %t2 : tensor<5x6xf32>) -> tensor<?x?x12xf32> { 136 %c0 = arith.constant 0 : index 137 %0 = vector.transfer_write %v, %t2[%c0, %c0] {in_bounds = [true, true]} : vector<5x6xf32>, tensor<5x6xf32> 138 %1 = tensor.insert_slice %0 into %t1[4, 3, %s] [1, 5, 6] [1, 1, 1] : tensor<5x6xf32> into tensor<?x?x12xf32> 139 return %1 : tensor<?x?x12xf32> 140} 141