mlir/test/Dialect/Tensor/fold-tensor-subset-ops-into-vector-transfers.mlir

   1 // RUN: mlir-opt -split-input-file -transform-interpreter %s | FileCheck %s
   2
   3 module 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]]
  23 func.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]]
  38 func.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]]
  54 func.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]]
  69 func.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>
  83 func.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]]
  98 func.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]]
 110 func.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>
 122 func.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]]
 135 func.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 }