xref: /llvm-project/mlir/test/Dialect/Vector/vector-mask-lowering-transforms.mlir (revision bb6d5c220004a5d7e466a669324001285a688918)
1// RUN: mlir-opt %s --transform-interpreter --split-input-file | FileCheck %s
2
3// CHECK-LABEL: func @genbool_1d
4// CHECK: %[[T0:.*]] = arith.constant dense<[true, true, true, true, false, false, false, false]> : vector<8xi1>
5// CHECK: return %[[T0]] : vector<8xi1>
6
7func.func @genbool_1d() -> vector<8xi1> {
8  %0 = vector.constant_mask [4] : vector<8xi1>
9  return %0 : vector<8xi1>
10}
11
12// CHECK-LABEL: func @genbool_2d
13// CHECK: %[[C1:.*]] = arith.constant dense<[true, true, false, false]> : vector<4xi1>
14// CHECK: %[[C2:.*]] = arith.constant dense<false> : vector<4x4xi1>
15// CHECK: %[[T0:.*]] = vector.insert %[[C1]], %[[C2]] [0] : vector<4xi1> into vector<4x4xi1>
16// CHECK: %[[T1:.*]] = vector.insert %[[C1]], %[[T0]] [1] : vector<4xi1> into vector<4x4xi1>
17// CHECK: return %[[T1]] : vector<4x4xi1>
18
19func.func @genbool_2d() -> vector<4x4xi1> {
20  %v = vector.constant_mask [2, 2] : vector<4x4xi1>
21  return %v: vector<4x4xi1>
22}
23
24// CHECK-LABEL: func @genbool_3d
25// CHECK-DAG: %[[C1:.*]] = arith.constant dense<[true, true, true, false]> : vector<4xi1>
26// CHECK-DAG: %[[C2:.*]] = arith.constant dense<false> : vector<3x4xi1>
27// CHECK-DAG: %[[C3:.*]] = arith.constant dense<false> : vector<2x3x4xi1>
28// CHECK: %[[T0:.*]] = vector.insert %[[C1]], %[[C2]] [0] : vector<4xi1> into vector<3x4xi1>
29// CHECK: %[[T1:.*]] = vector.insert %[[T0]], %[[C3]] [0] : vector<3x4xi1> into vector<2x3x4xi1>
30// CHECK: return %[[T1]] : vector<2x3x4xi1>
31
32func.func @genbool_3d() -> vector<2x3x4xi1> {
33  %v = vector.constant_mask [1, 1, 3] : vector<2x3x4xi1>
34  return %v: vector<2x3x4xi1>
35}
36
37// CHECK-LABEL: func @genbool_var_1d(
38// CHECK-SAME: %[[A:.*]]: index)
39// CHECK:      %[[T0:.*]] = vector.create_mask %[[A]] : vector<3xi1>
40// CHECK:      return %[[T0]] : vector<3xi1>
41
42func.func @genbool_var_1d(%arg0: index) -> vector<3xi1> {
43  %0 = vector.create_mask %arg0 : vector<3xi1>
44  return %0 : vector<3xi1>
45}
46
47// CHECK-LABEL: func @genbool_var_2d(
48// CHECK-SAME: %[[A:.*0]]: index,
49// CHECK-SAME: %[[B:.*1]]: index)
50// CHECK-DAG:  %[[C1:.*]] = arith.constant dense<false> : vector<3xi1>
51// CHECK-DAG:  %[[C2:.*]] = arith.constant dense<false> : vector<2x3xi1>
52// CHECK-DAG:  %[[c0:.*]] = arith.constant 0 : index
53// CHECK-DAG:  %[[c1:.*]] = arith.constant 1 : index
54// CHECK:      %[[T0:.*]] = vector.create_mask %[[B]] : vector<3xi1>
55// CHECK:      %[[T1:.*]] = arith.cmpi sgt, %[[A]], %[[c0]] : index
56// CHECK:      %[[T2:.*]] = arith.select %[[T1]], %[[T0]], %[[C1]] : vector<3xi1>
57// CHECK:      %[[T3:.*]] = vector.insert %[[T2]], %[[C2]] [0] : vector<3xi1> into vector<2x3xi1>
58// CHECK:      %[[T4:.*]] = arith.cmpi sgt, %[[A]], %[[c1]] : index
59// CHECK:      %[[T5:.*]] = arith.select %[[T4]], %[[T0]], %[[C1]] : vector<3xi1>
60// CHECK:      %[[T6:.*]] = vector.insert %[[T5]], %[[T3]] [1] : vector<3xi1> into vector<2x3xi1>
61// CHECK:      return %[[T6]] : vector<2x3xi1>
62
63func.func @genbool_var_2d(%arg0: index, %arg1: index) -> vector<2x3xi1> {
64  %0 = vector.create_mask %arg0, %arg1 : vector<2x3xi1>
65  return %0 : vector<2x3xi1>
66}
67
68// CHECK-LABEL: func @genbool_var_3d(
69// CHECK-SAME: %[[A:.*0]]: index,
70// CHECK-SAME: %[[B:.*1]]: index,
71// CHECK-SAME: %[[C:.*2]]: index)
72// CHECK-DAG:  %[[C1:.*]] = arith.constant dense<false> : vector<7xi1>
73// CHECK-DAG:  %[[C2:.*]] = arith.constant dense<false> : vector<1x7xi1>
74// CHECK-DAG:  %[[C3:.*]] = arith.constant dense<false> : vector<2x1x7xi1>
75// CHECK-DAG:  %[[c0:.*]] = arith.constant 0 : index
76// CHECK-DAG:  %[[c1:.*]] = arith.constant 1 : index
77// CHECK:      %[[T0:.*]] = vector.create_mask %[[C]] : vector<7xi1>
78// CHECK:      %[[T1:.*]] = arith.cmpi sgt, %[[B]], %[[c0]] : index
79// CHECK:      %[[T2:.*]] = arith.select %[[T1]], %[[T0]], %[[C1]] : vector<7xi1>
80// CHECK:      %[[T3:.*]] = vector.insert %[[T2]], %[[C2]] [0] : vector<7xi1> into vector<1x7xi1>
81// CHECK:      %[[T4:.*]] = arith.cmpi sgt, %[[A]], %[[c0]] : index
82// CHECK:      %[[T5:.*]] = arith.select %[[T4]], %[[T3]], %[[C2]] : vector<1x7xi1>
83// CHECK:      %[[T6:.*]] = vector.insert %[[T5]], %[[C3]] [0] : vector<1x7xi1> into vector<2x1x7xi1>
84// CHECK:      %[[T7:.*]] = arith.cmpi sgt, %[[A]], %[[c1]] : index
85// CHECK:      %[[T8:.*]] = arith.select %[[T7]], %[[T3]], %[[C2]] : vector<1x7xi1>
86// CHECK:      %[[T9:.*]] = vector.insert %[[T8]], %[[T6]] [1] : vector<1x7xi1> into vector<2x1x7xi1>
87// CHECK:      return %[[T9]] : vector<2x1x7xi1>
88
89func.func @genbool_var_3d(%arg0: index, %arg1: index, %arg2: index) -> vector<2x1x7xi1> {
90  %0 = vector.create_mask %arg0, %arg1, %arg2 : vector<2x1x7xi1>
91  return %0 : vector<2x1x7xi1>
92}
93
94module attributes {transform.with_named_sequence} {
95  transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
96    %f = transform.structured.match ops{["func.func"]} in %module_op
97      : (!transform.any_op) -> !transform.any_op
98
99    transform.apply_patterns to %f {
100      transform.apply_patterns.vector.lower_masks
101    } : !transform.any_op
102    transform.yield
103  }
104}
105
106// -----
107
108// CHECK-LABEL: func @transfer_read_3d(
109func.func @transfer_read_3d(
110    %t: tensor<?x?x?xf32>, %arg0: index, %arg1: index, %arg2: index)
111  -> vector<2x1x7xf32> {
112  %c0 = arith.constant 0 : index
113  %f0 = arith.constant 0.0 : f32
114  //      CHECK: %[[mask:.*]] = vector.create_mask
115  //  CHECK-NOT: vector.mask
116  //      CHECK: vector.transfer_read {{.*}}, %[[mask]] {in_bounds = [true, true, true]}
117  // CHECK-SAME:   : tensor<?x?x?xf32>, vector<2x1x7xf32>
118  %0 = vector.create_mask %arg0, %arg1, %arg2 : vector<2x1x7xi1>
119  %1 = vector.mask %0 {
120    vector.transfer_read %t[%c0, %c0, %c0], %f0 {in_bounds = [true, true, true]}
121      : tensor<?x?x?xf32>, vector<2x1x7xf32>
122  } : vector<2x1x7xi1> -> vector<2x1x7xf32>
123
124  return %1: vector<2x1x7xf32>
125}
126
127module attributes {transform.with_named_sequence} {
128  transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
129    %f = transform.structured.match ops{["func.func"]} in %module_op
130      : (!transform.any_op) -> !transform.any_op
131
132    transform.apply_patterns to %f {
133      transform.apply_patterns.vector.lower_masked_transfers
134    } : !transform.any_op
135    transform.yield
136  }
137}
138