xref: /llvm-project/mlir/test/Dialect/SPIRV/Transforms/webgpu-prepare.mlir (revision d61ec513c42005bb071eb15386deb5de585ff267)
1// RUN: mlir-opt --split-input-file --verify-diagnostics \
2// RUN:   --spirv-webgpu-prepare --cse %s | FileCheck %s
3
4//===----------------------------------------------------------------------===//
5// spirv.UMulExtended
6//===----------------------------------------------------------------------===//
7
8spirv.module Logical GLSL450 {
9
10// CHECK-LABEL: func @umul_extended_i32
11// CHECK-SAME:       ([[ARG0:%.+]]: i32, [[ARG1:%.+]]: i32)
12// CHECK-DAG:        [[CSTMASK:%.+]] = spirv.Constant 65535 : i32
13// CHECK-DAG:        [[CST16:%.+]]   = spirv.Constant 16 : i32
14// CHECK-NEXT:       [[LHSLOW:%.+]]  = spirv.BitwiseAnd [[ARG0]], [[CSTMASK]] : i32
15// CHECK-NEXT:       [[LHSHI:%.+]]   = spirv.ShiftRightLogical [[ARG0]], [[CST16]] : i32
16// CHECK-NEXT:       [[RHSLOW:%.+]]  = spirv.BitwiseAnd [[ARG1]], [[CSTMASK]] : i32
17// CHECK-NEXT:       [[RHSHI:%.+]]   = spirv.ShiftRightLogical [[ARG1]], [[CST16]] : i32
18// CHECK-DAG:                          spirv.IMul [[LHSLOW]], [[RHSLOW]]
19// CHECK-DAG:                          spirv.IMul [[LHSLOW]], [[RHSHI]]
20// CHECK-DAG:                          spirv.IMul [[LHSHI]],  [[RHSLOW]]
21// CHECK-DAG:                          spirv.IMul [[LHSHI]],  [[RHSHI]]
22// CHECK-DAG:                          spirv.IAdd
23// CHECK-DAG:                          spirv.IAdd
24// CHECK-DAG:                          spirv.IAdd
25// CHECK-DAG:                          spirv.IAdd
26// CHECK:                              spirv.ShiftLeftLogical {{%.+}}, [[CST16]] : i32
27// CHECK:                              spirv.BitwiseOr
28// CHECK:                              spirv.ShiftLeftLogical {{%.+}}, [[CST16]] : i32
29// CHECK:                              spirv.BitwiseOr
30// CHECK:            [[RES:%.+]]     = spirv.CompositeConstruct [[RESLO:%.+]], [[RESHI:%.+]] : (i32, i32) -> !spirv.struct<(i32, i32)>
31// CHECK-NEXT:       spirv.ReturnValue [[RES]] : !spirv.struct<(i32, i32)>
32spirv.func @umul_extended_i32(%arg0 : i32, %arg1 : i32) -> !spirv.struct<(i32, i32)> "None" {
33  %0 = spirv.UMulExtended %arg0, %arg1 : !spirv.struct<(i32, i32)>
34  spirv.ReturnValue %0 : !spirv.struct<(i32, i32)>
35}
36
37// CHECK-LABEL: func @umul_extended_vector_i32
38// CHECK-SAME:       ([[ARG0:%.+]]: vector<3xi32>, [[ARG1:%.+]]: vector<3xi32>)
39// CHECK-DAG:        [[CSTMASK:%.+]] = spirv.Constant dense<65535> : vector<3xi32>
40// CHECK-DAG:        [[CST16:%.+]]   = spirv.Constant dense<16> : vector<3xi32>
41// CHECK-NEXT:       [[LHSLOW:%.+]]  = spirv.BitwiseAnd [[ARG0]], [[CSTMASK]] : vector<3xi32>
42// CHECK-NEXT:       [[LHSHI:%.+]]   = spirv.ShiftRightLogical [[ARG0]], [[CST16]] : vector<3xi32>
43// CHECK-NEXT:       [[RHSLOW:%.+]]  = spirv.BitwiseAnd [[ARG1]], [[CSTMASK]] : vector<3xi32>
44// CHECK-NEXT:       [[RHSHI:%.+]]   = spirv.ShiftRightLogical [[ARG1]], [[CST16]] : vector<3xi32>
45// CHECK-DAG:                          spirv.IMul [[LHSLOW]], [[RHSLOW]]
46// CHECK-DAG:                          spirv.IMul [[LHSLOW]], [[RHSHI]]
47// CHECK-DAG:                          spirv.IMul [[LHSHI]],  [[RHSLOW]]
48// CHECK-DAG:                          spirv.IMul [[LHSHI]],  [[RHSHI]]
49// CHECK-DAG:                          spirv.IAdd
50// CHECK-DAG:                          spirv.IAdd
51// CHECK-DAG:                          spirv.IAdd
52// CHECK-DAG:                          spirv.IAdd
53// CHECK:                              spirv.ShiftLeftLogical {{%.+}}, [[CST16]]
54// CHECK:                              spirv.BitwiseOr
55// CHECK:                              spirv.ShiftLeftLogical {{%.+}}, [[CST16]]
56// CHECK:                              spirv.BitwiseOr
57// CHECK-NEXT:       [[RES:%.+]]     = spirv.CompositeConstruct [[RESLOW:%.+]], [[RESHI:%.+]]
58// CHECK-NEXT:       spirv.ReturnValue [[RES]] : !spirv.struct<(vector<3xi32>, vector<3xi32>)>
59spirv.func @umul_extended_vector_i32(%arg0 : vector<3xi32>, %arg1 : vector<3xi32>)
60  -> !spirv.struct<(vector<3xi32>, vector<3xi32>)> "None" {
61  %0 = spirv.UMulExtended %arg0, %arg1 : !spirv.struct<(vector<3xi32>, vector<3xi32>)>
62  spirv.ReturnValue %0 : !spirv.struct<(vector<3xi32>, vector<3xi32>)>
63}
64
65// CHECK-LABEL: func @umul_extended_i16
66// CHECK-NEXT:       spirv.UMulExtended
67// CHECK-NEXT:       spirv.ReturnValue
68spirv.func @umul_extended_i16(%arg : i16) -> !spirv.struct<(i16, i16)> "None" {
69  %0 = spirv.UMulExtended %arg, %arg : !spirv.struct<(i16, i16)>
70  spirv.ReturnValue %0 : !spirv.struct<(i16, i16)>
71}
72
73//===----------------------------------------------------------------------===//
74// spirv.SMulExtended
75//===----------------------------------------------------------------------===//
76
77// CHECK-LABEL: func @smul_extended_i32
78// CHECK-SAME:       ([[ARG0:%.+]]: i32, [[ARG1:%.+]]: i32)
79// CHECK-DAG:        [[CSTMASK:%.+]] = spirv.Constant 65535 : i32
80// CHECK-DAG:        [[CST16:%.+]]   = spirv.Constant 16 : i32
81// CHECK-NEXT:       [[LHSLOW:%.+]]  = spirv.BitwiseAnd [[ARG0]], [[CSTMASK]] : i32
82// CHECK-NEXT:       [[LHSHI:%.+]]   = spirv.ShiftRightLogical [[ARG0]], [[CST16]] : i32
83// CHECK-NEXT:       [[LHSSIGN:%.+]] = spirv.ShiftRightArithmetic [[ARG0]], [[CST16]] : i32
84// CHECK-NEXT:       [[LHSEXT:%.+]]  = spirv.ShiftRightLogical [[LHSSIGN]], [[CST16]] : i32
85// CHECK-NEXT:       [[RHSLOW:%.+]]  = spirv.BitwiseAnd [[ARG1]], [[CSTMASK]] : i32
86// CHECK-NEXT:       [[RHSHI:%.+]]   = spirv.ShiftRightLogical [[ARG1]], [[CST16]] : i32
87// CHECK-NEXT:       [[RHSSIGN:%.+]] = spirv.ShiftRightArithmetic [[ARG1]], [[CST16]] : i32
88// CHECK-NEXT:       [[RHSEXT:%.+]]  = spirv.ShiftRightLogical [[RHSSIGN]], [[CST16]] : i32
89// CHECK-DAG:                          spirv.IMul [[LHSLOW]], [[RHSLOW]]
90// CHECK-DAG:                          spirv.IMul [[LHSLOW]], [[RHSHI]]
91// CHECK-DAG:                          spirv.IMul [[LHSLOW]], [[RHSEXT]]
92// CHECK-DAG:                          spirv.IMul [[LHSHI]],  [[RHSLOW]]
93// CHECK-DAG:                          spirv.IMul [[LHSHI]],  [[RHSHI]]
94// CHECK-DAG:                          spirv.IMul [[LHSHI]],  [[RHSEXT]]
95// CHECK-DAG:                          spirv.IMul [[LHSEXT]], [[RHSLOW]]
96// CHECK-DAG:                          spirv.IMul [[LHSEXT]], [[RHSHI]]
97// CHECK:                              spirv.ShiftLeftLogical {{%.+}}, [[CST16]] : i32
98// CHECK:                              spirv.BitwiseOr
99// CHECK:                              spirv.ShiftLeftLogical {{%.+}}, [[CST16]] : i32
100// CHECK:                              spirv.BitwiseOr
101// CHECK:            [[RES:%.+]]     = spirv.CompositeConstruct [[RESLO:%.+]], [[RESHI:%.+]] : (i32, i32) -> !spirv.struct<(i32, i32)>
102// CHECK-NEXT:       spirv.ReturnValue [[RES]] : !spirv.struct<(i32, i32)>
103spirv.func @smul_extended_i32(%arg0 : i32, %arg1 : i32) -> !spirv.struct<(i32, i32)> "None" {
104  %0 = spirv.SMulExtended %arg0, %arg1 : !spirv.struct<(i32, i32)>
105  spirv.ReturnValue %0 : !spirv.struct<(i32, i32)>
106}
107
108// CHECK-LABEL: func @smul_extended_vector_i32
109// CHECK-SAME:       ([[ARG0:%.+]]: vector<3xi32>, [[ARG1:%.+]]: vector<3xi32>)
110// CHECK-DAG:        [[CSTMASK:%.+]] = spirv.Constant dense<65535> : vector<3xi32>
111// CHECK-DAG:        [[CST16:%.+]]   = spirv.Constant dense<16> : vector<3xi32>
112// CHECK-NEXT:       [[LHSLOW:%.+]]  = spirv.BitwiseAnd [[ARG0]], [[CSTMASK]] : vector<3xi32>
113// CHECK-NEXT:       [[LHSHI:%.+]]   = spirv.ShiftRightLogical [[ARG0]], [[CST16]] : vector<3xi32>
114// CHECK-NEXT:       [[LHSSIGN:%.+]] = spirv.ShiftRightArithmetic [[ARG0]], [[CST16]] : vector<3xi32>
115// CHECK-NEXT:       [[LHSEXT:%.+]]  = spirv.ShiftRightLogical [[LHSSIGN]], [[CST16]] : vector<3xi32>
116// CHECK-NEXT:       [[RHSLOW:%.+]]  = spirv.BitwiseAnd [[ARG1]], [[CSTMASK]] : vector<3xi32>
117// CHECK-NEXT:       [[RHSHI:%.+]]   = spirv.ShiftRightLogical [[ARG1]], [[CST16]] : vector<3xi32>
118// CHECK-NEXT:       [[RHSSIGN:%.+]] = spirv.ShiftRightArithmetic [[ARG1]], [[CST16]] : vector<3xi32>
119// CHECK-NEXT:       [[RHSEXT:%.+]]  = spirv.ShiftRightLogical [[RHSSIGN]], [[CST16]] : vector<3xi32>
120// CHECK-DAG:                          spirv.IMul [[LHSLOW]], [[RHSLOW]]
121// CHECK-DAG:                          spirv.IMul [[LHSLOW]], [[RHSHI]]
122// CHECK-DAG:                          spirv.IMul [[LHSLOW]], [[RHSEXT]]
123// CHECK-DAG:                          spirv.IMul [[LHSHI]],  [[RHSLOW]]
124// CHECK-DAG:                          spirv.IMul [[LHSHI]],  [[RHSHI]]
125// CHECK-DAG:                          spirv.IMul [[LHSHI]],  [[RHSEXT]]
126// CHECK-DAG:                          spirv.IMul [[LHSEXT]], [[RHSLOW]]
127// CHECK-DAG:                          spirv.IMul [[LHSEXT]], [[RHSHI]]
128// CHECK:                              spirv.ShiftLeftLogical {{%.+}}, [[CST16]]
129// CHECK:                              spirv.BitwiseOr
130// CHECK:                              spirv.ShiftLeftLogical {{%.+}}, [[CST16]]
131// CHECK:                              spirv.BitwiseOr
132// CHECK-NEXT:       [[RES:%.+]]     = spirv.CompositeConstruct [[RESLOW:%.+]], [[RESHI:%.+]]
133// CHECK-NEXT:       spirv.ReturnValue [[RES]] : !spirv.struct<(vector<3xi32>, vector<3xi32>)>
134spirv.func @smul_extended_vector_i32(%arg0 : vector<3xi32>, %arg1 : vector<3xi32>)
135  -> !spirv.struct<(vector<3xi32>, vector<3xi32>)> "None" {
136  %0 = spirv.SMulExtended %arg0, %arg1 : !spirv.struct<(vector<3xi32>, vector<3xi32>)>
137  spirv.ReturnValue %0 : !spirv.struct<(vector<3xi32>, vector<3xi32>)>
138}
139
140// CHECK-LABEL: func @smul_extended_i16
141// CHECK-NEXT:       spirv.SMulExtended
142// CHECK-NEXT:       spirv.ReturnValue
143spirv.func @smul_extended_i16(%arg : i16) -> !spirv.struct<(i16, i16)> "None" {
144  %0 = spirv.SMulExtended %arg, %arg : !spirv.struct<(i16, i16)>
145  spirv.ReturnValue %0 : !spirv.struct<(i16, i16)>
146}
147
148// CHECK-LABEL: func @iaddcarry_i32
149// CHECK-SAME:       ([[A:%.+]]: i32, [[B:%.+]]: i32)
150// CHECK-NEXT:       [[ONE:%.+]]    = spirv.Constant 1 : i32
151// CHECK-NEXT:       [[ZERO:%.+]]   = spirv.Constant 0 : i32
152// CHECK-NEXT:       [[OUT:%.+]]    = spirv.IAdd [[A]], [[B]]
153// CHECK-NEXT:       [[CMP:%.+]]    = spirv.ULessThan [[OUT]], [[A]]
154// CHECK-NEXT:       [[CARRY:%.+]]  = spirv.Select [[CMP]], [[ONE]], [[ZERO]]
155// CHECK-NEXT:       [[RES:%.+]]     = spirv.CompositeConstruct [[OUT]], [[CARRY]] : (i32, i32) -> !spirv.struct<(i32, i32)>
156// CHECK-NEXT:       spirv.ReturnValue [[RES]] : !spirv.struct<(i32, i32)>
157spirv.func @iaddcarry_i32(%a : i32, %b : i32) -> !spirv.struct<(i32, i32)> "None" {
158  %0 = spirv.IAddCarry %a, %b : !spirv.struct<(i32, i32)>
159  spirv.ReturnValue %0 : !spirv.struct<(i32, i32)>
160}
161
162// CHECK-LABEL: func @iaddcarry_vector_i32
163// CHECK-SAME:       ([[A:%.+]]: vector<3xi32>, [[B:%.+]]: vector<3xi32>)
164// CHECK-NEXT:       [[ONE:%.+]]    = spirv.Constant dense<1> : vector<3xi32>
165// CHECK-NEXT:       [[ZERO:%.+]]   = spirv.Constant dense<0> : vector<3xi32>
166// CHECK-NEXT:       [[OUT:%.+]]    = spirv.IAdd [[A]], [[B]]
167// CHECK-NEXT:       [[CMP:%.+]]    = spirv.ULessThan [[OUT]], [[A]]
168// CHECK-NEXT:       [[CARRY:%.+]]  = spirv.Select [[CMP]], [[ONE]], [[ZERO]]
169// CHECK-NEXT:       [[RES:%.+]]    = spirv.CompositeConstruct [[OUT]], [[CARRY]] : (vector<3xi32>, vector<3xi32>) -> !spirv.struct<(vector<3xi32>, vector<3xi32>)>
170// CHECK-NEXT:       spirv.ReturnValue [[RES]] : !spirv.struct<(vector<3xi32>, vector<3xi32>)>
171spirv.func @iaddcarry_vector_i32(%a : vector<3xi32>, %b : vector<3xi32>)
172  -> !spirv.struct<(vector<3xi32>, vector<3xi32>)> "None" {
173  %0 = spirv.IAddCarry %a, %b : !spirv.struct<(vector<3xi32>, vector<3xi32>)>
174  spirv.ReturnValue %0 : !spirv.struct<(vector<3xi32>, vector<3xi32>)>
175}
176
177// CHECK-LABEL: func @iaddcarry_i16
178// CHECK-NEXT:       spirv.IAddCarry
179// CHECK-NEXT:       spirv.ReturnValue
180spirv.func @iaddcarry_i16(%a : i16, %b : i16) -> !spirv.struct<(i16, i16)> "None" {
181  %0 = spirv.IAddCarry %a, %b : !spirv.struct<(i16, i16)>
182  spirv.ReturnValue %0 : !spirv.struct<(i16, i16)>
183}
184
185// CHECK-LABEL: func @is_inf_f32
186// CHECK-NEXT:       [[FALSE:%.+]] = spirv.Constant false
187// CHECK-NEXT:       spirv.ReturnValue [[FALSE]] : i1
188spirv.func @is_inf_f32(%a : f32) -> i1 "None" {
189  %0 = spirv.IsInf %a : f32
190  spirv.ReturnValue %0 : i1
191}
192
193// CHECK-LABEL: func @is_inf_4xf32
194// CHECK-NEXT:       [[FALSE:%.+]] = spirv.Constant dense<false> : vector<4xi1>
195// CHECK-NEXT:       spirv.ReturnValue [[FALSE]] : vector<4xi1>
196spirv.func @is_inf_4xf32(%a : vector<4xf32>) -> vector<4xi1> "None" {
197  %0 = spirv.IsInf %a : vector<4xf32>
198  spirv.ReturnValue %0 : vector<4xi1>
199}
200
201// CHECK-LABEL: func @is_nan_f32
202// CHECK-NEXT:       [[FALSE:%.+]] = spirv.Constant false
203// CHECK-NEXT:       spirv.ReturnValue [[FALSE]] : i1
204spirv.func @is_nan_f32(%a : f32) -> i1 "None" {
205  %0 = spirv.IsNan %a : f32
206  spirv.ReturnValue %0 : i1
207}
208
209// CHECK-LABEL: func @is_nan_4xf32
210// CHECK-NEXT:       [[FALSE:%.+]] = spirv.Constant dense<false> : vector<4xi1>
211// CHECK-NEXT:       spirv.ReturnValue [[FALSE]] : vector<4xi1>
212spirv.func @is_nan_4xf32(%a : vector<4xf32>) -> vector<4xi1> "None" {
213  %0 = spirv.IsNan %a : vector<4xf32>
214  spirv.ReturnValue %0 : vector<4xi1>
215}
216
217} // end module
218