166 lines · plain
1// RUN: mlir-opt %s -split-input-file -test-eliminate-vector-masks --split-input-file | FileCheck %s2 3// This tests a general pattern the vectorizer tends to emit.4 5// CHECK-LABEL: @eliminate_redundant_masks_through_insert_and_extracts6// 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.outerproduct10// CHECK: vector.transfer_write {{.*}} %[[ALL_TRUE_MASK]]11#map = affine_map<()[s0] -> (-(1080 mod s0) + 1080)>12 13func.func @eliminate_redundant_masks_through_insert_and_extracts(%tensor: tensor<1x1000xf32>, %rhs : f32) {14 %c4 = arith.constant 4 : index15 %vscale = vector.vscale16 %c4_vscale = arith.muli %vscale, %c4 : index17 %ub = affine.apply #map()[%c4_vscale]18 19 %c0 = arith.constant 0 : index20 %c1000 = arith.constant 1000 : index21 %c0_f32 = arith.constant 0.0 : f3222 %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 return44}45 46// -----47 48// CHECK-LABEL: @negative_extract_slice_size_shrink49// CHECK-NOT: vector.constant_mask50// CHECK: %[[MASK:.*]] = vector.create_mask51// CHECK: "test.some_use"(%[[MASK]]) : (vector<[4]xi1>) -> ()52func.func @negative_extract_slice_size_shrink(%tensor: tensor<1000xf32>) {53 %c0 = arith.constant 0 : index54 %c4 = arith.constant 4 : index55 %c1000 = arith.constant 1000 : index56 %vscale = vector.vscale57 %c4_vscale = arith.muli %vscale, %c4 : index58 %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 above61 // (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 return72}73 74// -----75 76// CHECK-LABEL: @trivially_all_true_case77// 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>) -> ()79func.func @trivially_all_true_case(%tensor: tensor<2x?xf32>)80{81 %c2 = arith.constant 2 : index82 %c4 = arith.constant 4 : index83 %vscale = vector.vscale84 %c4_vscale = arith.muli %vscale, %c4 : index85 // 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 return89}90 91// -----92 93// CHECK-LABEL: @negative_constant_dim_not_all_true94// CHECK-NOT: vector.constant_mask95// CHECK: %[[MASK:.*]] = vector.create_mask96// CHECK: "test.some_use"(%[[MASK]]) : (vector<2x[4]xi1>) -> ()97func.func @negative_constant_dim_not_all_true()98{99 %c1 = arith.constant 1 : index100 %c4 = arith.constant 4 : index101 %vscale = vector.vscale102 %c4_vscale = arith.muli %vscale, %c4 : index103 // Since %c1 is a constant, this will be found not to be all-true via simple104 // pattern matching.105 %mask = vector.create_mask %c1, %c4_vscale : vector<2x[4]xi1>106 "test.some_use"(%mask) : (vector<2x[4]xi1>) -> ()107 return108}109 110// -----111 112// CHECK-LABEL: @negative_constant_vscale_multiple_not_all_true113// CHECK-NOT: vector.constant_mask114// CHECK: %[[MASK:.*]] = vector.create_mask115// CHECK: "test.some_use"(%[[MASK]]) : (vector<2x[4]xi1>) -> ()116func.func @negative_constant_vscale_multiple_not_all_true() {117 %c2 = arith.constant 2 : index118 %c3 = arith.constant 3 : index119 %vscale = vector.vscale120 %c3_vscale = arith.muli %vscale, %c3 : index121 // Since %c3_vscale is a constant vscale multiple, this will be found not to122 // 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 return126}127 128// -----129 130// CHECK-LABEL: @negative_value_bounds_fixed_dim_not_all_true131// CHECK-NOT: vector.constant_mask132// CHECK: %[[MASK:.*]] = vector.create_mask133// CHECK: "test.some_use"(%[[MASK]]) : (vector<3x[4]xi1>) -> ()134func.func @negative_value_bounds_fixed_dim_not_all_true(%tensor: tensor<2x?xf32>)135{136 %c0 = arith.constant 0 : index137 %c4 = arith.constant 4 : index138 %vscale = vector.vscale139 %c4_vscale = arith.muli %vscale, %c4 : index140 // This is _very_ simple, but since tensor.dim is not a constant, value bounds141 // 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 return146}147 148// -----149 150// CHECK-LABEL: @negative_value_bounds_scalable_dim_not_all_true151// CHECK-NOT: vector.constant_mask152// CHECK: %[[MASK:.*]] = vector.create_mask153// CHECK: "test.some_use"(%[[MASK]]) : (vector<3x[4]xi1>) -> ()154func.func @negative_value_bounds_scalable_dim_not_all_true(%tensor: tensor<2x100xf32>) {155 %c1 = arith.constant 1 : index156 %c3 = arith.constant 3 : index157 %vscale = vector.vscale158 %c3_vscale = arith.muli %vscale, %c3 : index159 %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 return165}166