mlir/test/Dialect/Vector/eliminate-masks.mlir

   1 // RUN: mlir-opt %s -split-input-file -test-eliminate-vector-masks --split-input-file | FileCheck %s
   2
   3 // This tests a general pattern the vectorizer tends to emit.
   4
   5 // CHECK-LABEL: @eliminate_redundant_masks_through_insert_and_extracts
   6 // CHECK: %[[ALL_TRUE_MASK:.*]] = vector.constant_mask [4] : vector<[4]xi1>
   7 // CHECK: vector.transfer_read {{.*}} %[[ALL_TRUE_MASK]]
   8 // CHECK: vector.mask %[[ALL_TRUE_MASK:.*]] {
   9 // CHECK-SAME:  vector.outerproduct
  10 // CHECK: vector.transfer_write {{.*}} %[[ALL_TRUE_MASK]]
  11 #map = affine_map<()[s0] -> (-(1080 mod s0) + 1080)>
  12
  13 func.func @eliminate_redundant_masks_through_insert_and_extracts(%tensor: tensor<1x1000xf32>, %rhs : f32) {
  14   %c4 = arith.constant 4 : index
  15   %vscale = vector.vscale
  16   %c4_vscale = arith.muli %vscale, %c4 : index
  17   %ub = affine.apply #map()[%c4_vscale]
  18
  19   %c0 = arith.constant 0 : index
  20   %c1000 = arith.constant 1000 : index
  21   %c0_f32 = arith.constant 0.0 : f32
  22   %extracted_slice_0 = tensor.extract_slice %tensor[0, 0] [1, %c4_vscale] [1, 1] : tensor<1x1000xf32> to tensor<1x?xf32>
  23   %output_tensor = scf.for %i = %c0 to %ub step %c4_vscale iter_args(%arg = %extracted_slice_0) -> tensor<1x?xf32> {
  24     // 1. Extract a slice.
  25     %extracted_slice_1 = tensor.extract_slice %arg[0, %i] [1, %c4_vscale] [1, 1] : tensor<1x?xf32> to tensor<?xf32>
  26
  27     // 2. Create a mask for the slice.
  28     %dim_1 = tensor.dim %extracted_slice_1, %c0 : tensor<?xf32>
  29     %mask = vector.create_mask %dim_1 : vector<[4]xi1>
  30
  31     // 3. Read the slice and do some computation.
  32     %lhs = vector.transfer_read %extracted_slice_1[%c0], %c0_f32, %mask {in_bounds = [true]} : tensor<?xf32>, vector<[4]xf32>
  33     %new_vec = vector.mask %mask { vector.outerproduct %lhs, %rhs {kind = #vector.kind<add>} : vector<[4]xf32>, f32 } : vector<[4]xi1> -> vector<[4]xf32>
  34
  35     // 4. Write the new value.
  36     %write = vector.transfer_write %new_vec, %extracted_slice_1[%c0], %mask {in_bounds = [true]} : vector<[4]xf32>, tensor<?xf32>
  37
  38     // 5. Insert and yield the new tensor value.
  39     %result = tensor.insert_slice %write into %arg[0, %i] [1, %c4_vscale] [1, 1] : tensor<?xf32> into tensor<1x?xf32>
  40     scf.yield %result : tensor<1x?xf32>
  41   }
  42   "test.some_use"(%output_tensor) : (tensor<1x?xf32>) -> ()
  43   return
  44 }
  45
  46 // -----
  47
  48 // CHECK-LABEL: @negative_extract_slice_size_shrink
  49 // CHECK-NOT: vector.constant_mask
  50 // CHECK: %[[MASK:.*]] = vector.create_mask
  51 // CHECK: "test.some_use"(%[[MASK]]) : (vector<[4]xi1>) -> ()
  52 func.func @negative_extract_slice_size_shrink(%tensor: tensor<1000xf32>) {
  53   %c0 = arith.constant 0 : index
  54   %c4 = arith.constant 4 : index
  55   %c1000 = arith.constant 1000 : index
  56   %vscale = vector.vscale
  57   %c4_vscale = arith.muli %vscale, %c4 : index
  58   %extracted_slice = tensor.extract_slice %tensor[0] [%c4_vscale] [1] : tensor<1000xf32> to tensor<?xf32>
  59   %slice = scf.for %i = %c0 to %c1000 step %c4_vscale iter_args(%arg = %extracted_slice) -> tensor<?xf32> {
  60     // This mask cannot be eliminated even though looking at the operations above
  61     // (this comment) it appears `tensor.dim` will always be c4_vscale (so the mask all-true).
  62     %dim = tensor.dim %arg, %c0 : tensor<?xf32>
  63     %mask = vector.create_mask %dim : vector<[4]xi1>
  64     "test.some_use"(%mask) : (vector<[4]xi1>) -> ()
  65     // !!! Here the size of the mask could shrink in the next iteration.
  66     %next_num_elts = affine.min  affine_map<(d0)[s0] -> (-d0 + 1000, s0)>(%i)[%c4_vscale]
  67     %new_extracted_slice = tensor.extract_slice %tensor[%c4_vscale] [%next_num_elts] [1] : tensor<1000xf32> to tensor<?xf32>
  68     scf.yield %new_extracted_slice : tensor<?xf32>
  69   }
  70   "test.some_use"(%slice) : (tensor<?xf32>) -> ()
  71   return
  72 }
  73
  74 // -----
  75
  76 // CHECK-LABEL: @trivially_all_true_case
  77 // CHECK: %[[ALL_TRUE_MASK:.*]] = vector.constant_mask [2, 4] : vector<2x[4]xi1>
  78 // CHECK: "test.some_use"(%[[ALL_TRUE_MASK]]) : (vector<2x[4]xi1>) -> ()
  79 func.func @trivially_all_true_case(%tensor: tensor<2x?xf32>)
  80 {
  81   %c2 = arith.constant 2 : index
  82   %c4 = arith.constant 4 : index
  83   %vscale = vector.vscale
  84   %c4_vscale = arith.muli %vscale, %c4 : index
  85   // Is found to be all true _without_ value bounds analysis.
  86   %mask = vector.create_mask %c2, %c4_vscale : vector<2x[4]xi1>
  87   "test.some_use"(%mask) : (vector<2x[4]xi1>) -> ()
  88   return
  89 }
  90
  91 // -----
  92
  93 // CHECK-LABEL: @negative_constant_dim_not_all_true
  94 // CHECK-NOT: vector.constant_mask
  95 // CHECK: %[[MASK:.*]] = vector.create_mask
  96 // CHECK: "test.some_use"(%[[MASK]]) : (vector<2x[4]xi1>) -> ()
  97 func.func @negative_constant_dim_not_all_true()
  98 {
  99   %c1 = arith.constant 1 : index
 100   %c4 = arith.constant 4 : index
 101   %vscale = vector.vscale
 102   %c4_vscale = arith.muli %vscale, %c4 : index
 103   // Since %c1 is a constant, this will be found not to be all-true via simple
 104   // pattern matching.
 105   %mask = vector.create_mask %c1, %c4_vscale : vector<2x[4]xi1>
 106   "test.some_use"(%mask) : (vector<2x[4]xi1>) -> ()
 107   return
 108 }
 109
 110 // -----
 111
 112 // CHECK-LABEL: @negative_constant_vscale_multiple_not_all_true
 113 // CHECK-NOT: vector.constant_mask
 114 // CHECK: %[[MASK:.*]] = vector.create_mask
 115 // CHECK: "test.some_use"(%[[MASK]]) : (vector<2x[4]xi1>) -> ()
 116 func.func @negative_constant_vscale_multiple_not_all_true() {
 117   %c2 = arith.constant 2 : index
 118   %c3 = arith.constant 3 : index
 119   %vscale = vector.vscale
 120   %c3_vscale = arith.muli %vscale, %c3 : index
 121   // Since %c3_vscale is a constant vscale multiple, this will be found not to
 122   // be all-true via simple pattern matching.
 123   %mask = vector.create_mask %c2, %c3_vscale : vector<2x[4]xi1>
 124   "test.some_use"(%mask) : (vector<2x[4]xi1>) -> ()
 125   return
 126 }
 127
 128 // -----
 129
 130 // CHECK-LABEL: @negative_value_bounds_fixed_dim_not_all_true
 131 // CHECK-NOT: vector.constant_mask
 132 // CHECK: %[[MASK:.*]] = vector.create_mask
 133 // CHECK: "test.some_use"(%[[MASK]]) : (vector<3x[4]xi1>) -> ()
 134 func.func @negative_value_bounds_fixed_dim_not_all_true(%tensor: tensor<2x?xf32>)
 135 {
 136   %c0 = arith.constant 0 : index
 137   %c4 = arith.constant 4 : index
 138   %vscale = vector.vscale
 139   %c4_vscale = arith.muli %vscale, %c4 : index
 140   // This is _very_ simple, but since tensor.dim is not a constant, value bounds
 141   // will be used to resolve it.
 142   %dim = tensor.dim %tensor, %c0 : tensor<2x?xf32>
 143   %mask = vector.create_mask %dim, %c4_vscale : vector<3x[4]xi1>
 144   "test.some_use"(%mask) : (vector<3x[4]xi1>) -> ()
 145   return
 146 }
 147
 148 // -----
 149
 150 // CHECK-LABEL: @negative_value_bounds_scalable_dim_not_all_true
 151 // CHECK-NOT: vector.constant_mask
 152 // CHECK: %[[MASK:.*]] = vector.create_mask
 153 // CHECK: "test.some_use"(%[[MASK]]) : (vector<3x[4]xi1>) -> ()
 154 func.func @negative_value_bounds_scalable_dim_not_all_true(%tensor: tensor<2x100xf32>) {
 155   %c1 = arith.constant 1 : index
 156   %c3 = arith.constant 3 : index
 157   %vscale = vector.vscale
 158   %c3_vscale = arith.muli %vscale, %c3 : index
 159   %slice = tensor.extract_slice %tensor[0, 0] [2, %c3_vscale] [1, 1] : tensor<2x100xf32> to tensor<2x?xf32>
 160   // Another simple example, but value bounds will be used to resolve the tensor.dim.
 161   %dim = tensor.dim %slice, %c1 : tensor<2x?xf32>
 162   %mask = vector.create_mask %c3, %dim : vector<3x[4]xi1>
 163   "test.some_use"(%mask) : (vector<3x[4]xi1>) -> ()
 164   return
 165 }