1// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -verify-diagnostics %s | FileCheck %s
2
3// CHECK-LABEL: @cooperative_matrix_load
4spv.func @cooperative_matrix_load(%ptr : !spv.ptr<i32, StorageBuffer>, %stride : i32, %b : i1) "None" {
5  // CHECK: {{%.*}} = spv.CooperativeMatrixLoadNV {{%.*}}, {{%.*}}, {{%.*}} : !spv.ptr<i32, StorageBuffer> as !spv.coopmatrix<16x8xi32, Workgroup>
6  %0 = spv.CooperativeMatrixLoadNV %ptr, %stride, %b : !spv.ptr<i32, StorageBuffer> as !spv.coopmatrix<16x8xi32, Workgroup>
7  spv.Return
8}
9
10// -----
11// CHECK-LABEL: @cooperative_matrix_load_memaccess
12spv.func @cooperative_matrix_load_memaccess(%ptr : !spv.ptr<i32, StorageBuffer>, %stride : i32, %b : i1) "None" {
13  // CHECK: {{%.*}} = spv.CooperativeMatrixLoadNV {{%.*}}, {{%.*}}, {{%.*}} ["Volatile"] : !spv.ptr<i32, StorageBuffer> as !spv.coopmatrix<8x16xi32, Subgroup>
14  %0 = spv.CooperativeMatrixLoadNV %ptr, %stride, %b ["Volatile"] : !spv.ptr<i32, StorageBuffer> as !spv.coopmatrix<8x16xi32, Subgroup>
15  spv.Return
16}
17
18// CHECK-LABEL: @cooperative_matrix_load_diff_ptr_type
19spv.func @cooperative_matrix_load_diff_ptr_type(%ptr : !spv.ptr<vector<4xi32>, StorageBuffer>, %stride : i32, %b : i1) "None" {
20  // CHECK: {{%.*}} = spv.CooperativeMatrixLoadNV {{%.*}}, {{%.*}}, {{%.*}} ["Volatile"] : !spv.ptr<vector<4xi32>, StorageBuffer> as !spv.coopmatrix<8x16xi32, Subgroup>
21  %0 = spv.CooperativeMatrixLoadNV %ptr, %stride, %b ["Volatile"] : !spv.ptr<vector<4xi32>, StorageBuffer> as !spv.coopmatrix<8x16xi32, Subgroup>
22  spv.Return
23}
24
25// CHECK-LABEL: @cooperative_matrix_store
26spv.func @cooperative_matrix_store(%ptr : !spv.ptr<i32, StorageBuffer>, %stride : i32, %m : !spv.coopmatrix<8x16xi32, Workgroup>, %b : i1) "None" {
27  // CHECK: spv.CooperativeMatrixStoreNV {{%.*}}, {{%.*}}, {{%.*}} : !spv.ptr<i32, StorageBuffer>, !spv.coopmatrix<8x16xi32, Workgroup>
28  spv.CooperativeMatrixStoreNV %ptr, %m, %stride, %b : !spv.ptr<i32, StorageBuffer>, !spv.coopmatrix<8x16xi32, Workgroup>
29  spv.Return
30}
31
32// CHECK-LABEL: @cooperative_matrix_store_memaccess
33spv.func @cooperative_matrix_store_memaccess(%ptr : !spv.ptr<i32, StorageBuffer>, %m : !spv.coopmatrix<8x16xi32, Subgroup>, %stride : i32, %b : i1) "None" {
34  // CHECK: spv.CooperativeMatrixStoreNV {{%.*}}, {{%.*}}, {{%.*}} ["Volatile"] : !spv.ptr<i32, StorageBuffer>, !spv.coopmatrix<8x16xi32, Subgroup>
35  spv.CooperativeMatrixStoreNV %ptr, %m, %stride, %b ["Volatile"] : !spv.ptr<i32, StorageBuffer>, !spv.coopmatrix<8x16xi32, Subgroup>
36  spv.Return
37}
38
39// CHECK-LABEL: @cooperative_matrix_length
40spv.func @cooperative_matrix_length() -> i32 "None" {
41  // CHECK: {{%.*}} = spv.CooperativeMatrixLengthNV : !spv.coopmatrix<8x16xi32, Subgroup>
42  %0 = spv.CooperativeMatrixLengthNV : !spv.coopmatrix<8x16xi32, Subgroup>
43  spv.ReturnValue %0 : i32
44}
45
46// CHECK-LABEL: @cooperative_matrix_muladd
47spv.func @cooperative_matrix_muladd(%a : !spv.coopmatrix<8x32xi8, Subgroup>, %b : !spv.coopmatrix<32x8xi8, Subgroup>, %c : !spv.coopmatrix<8x8xi32, Subgroup>) "None" {
48  // CHECK: {{%.*}} = spv.CooperativeMatrixMulAddNV {{%.*}}, {{%.*}}, {{%.*}}  : !spv.coopmatrix<8x32xi8, Subgroup>, !spv.coopmatrix<32x8xi8, Subgroup> -> !spv.coopmatrix<8x8xi32, Subgroup>
49  %r = spv.CooperativeMatrixMulAddNV %a, %b, %c : !spv.coopmatrix<8x32xi8, Subgroup>, !spv.coopmatrix<32x8xi8, Subgroup> -> !spv.coopmatrix<8x8xi32, Subgroup>
50  spv.Return
51}
52
53// CHECK-LABEL: @cooperative_matrix_add
54spv.func @cooperative_matrix_add(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
55  // CHECK: {{%.*}} = spv.IAdd {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
56  %r = spv.IAdd %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
57  spv.Return
58}
59
60// CHECK-LABEL: @cooperative_matrix_sub
61spv.func @cooperative_matrix_sub(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
62  // CHECK: {{%.*}} = spv.ISub {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
63  %r = spv.ISub %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
64  spv.Return
65}
66
67// CHECK-LABEL: @cooperative_matrix_sdiv
68spv.func @cooperative_matrix_sdiv(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
69  // CHECK: {{%.*}} = spv.SDiv {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
70  %r = spv.SDiv %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
71  spv.Return
72}
73
74// CHECK-LABEL: @cooperative_matrix_udiv
75spv.func @cooperative_matrix_udiv(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
76  // CHECK: {{%.*}} = spv.UDiv {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
77  %r = spv.UDiv %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
78  spv.Return
79}
80
81// CHECK-LABEL: @cooperative_matrix_fadd
82spv.func @cooperative_matrix_fadd(%a : !spv.coopmatrix<8x16xf32, Subgroup>, %b : !spv.coopmatrix<8x16xf32, Subgroup>) "None" {
83  // CHECK: {{%.*}} = spv.FAdd {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup>
84  %r = spv.FAdd %a, %b : !spv.coopmatrix<8x16xf32, Subgroup>
85  spv.Return
86}
87
88// CHECK-LABEL: @cooperative_matrix_fsub
89spv.func @cooperative_matrix_fsub(%a : !spv.coopmatrix<8x16xf32, Subgroup>, %b : !spv.coopmatrix<8x16xf32, Subgroup>) "None" {
90  // CHECK: {{%.*}} = spv.FSub {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup>
91  %r = spv.FSub %a, %b : !spv.coopmatrix<8x16xf32, Subgroup>
92  spv.Return
93}
94
95// CHECK-LABEL: @cooperative_matrix_fdiv
96spv.func @cooperative_matrix_fdiv(%a : !spv.coopmatrix<8x16xf32, Subgroup>, %b : !spv.coopmatrix<8x16xf32, Subgroup>) "None" {
97  // CHECK: {{%.*}} = spv.FDiv {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup>
98  %r = spv.FDiv %a, %b : !spv.coopmatrix<8x16xf32, Subgroup>
99  spv.Return
100}
101
102// -----
103
104// CHECK-LABEL: @cooperative_matrix_access_chain
105spv.func @cooperative_matrix_access_chain(%a : !spv.ptr<!spv.coopmatrix<8x16xf32, Subgroup>, Function>) -> !spv.ptr<f32, Function> "None" {
106  %0 = spv.Constant 0: i32
107  // CHECK: {{%.*}} = spv.AccessChain {{%.*}}[{{%.*}}] : !spv.ptr<!spv.coopmatrix<8x16xf32, Subgroup>, Function>, i32
108  %1 = spv.AccessChain %a[%0] : !spv.ptr<!spv.coopmatrix<8x16xf32, Subgroup>, Function>, i32
109  spv.ReturnValue %1 : !spv.ptr<f32, Function>
110}
111
112// -----
113
114spv.func @cooperative_matrix_muladd(%a : !spv.coopmatrix<16x16xi32, Subgroup>, %b : !spv.coopmatrix<16x8xi32, Subgroup>, %c : !spv.coopmatrix<8x8xi32, Subgroup>) "None" {
115  // expected-error @+1 {{'spv.CooperativeMatrixMulAddNV' op matrix size must match}}
116  %r = spv.CooperativeMatrixMulAddNV %a, %b, %c : !spv.coopmatrix<16x16xi32, Subgroup>, !spv.coopmatrix<16x8xi32, Subgroup> -> !spv.coopmatrix<8x8xi32, Subgroup>
117  spv.Return
118}
119
120// -----
121
122spv.func @cooperative_matrix_muladd(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x8xi32, Subgroup>, %c : !spv.coopmatrix<8x8xi32, Subgroup>) "None" {
123  // expected-error @+1 {{'spv.CooperativeMatrixMulAddNV' op matrix size must match}}
124  %r = spv.CooperativeMatrixMulAddNV %a, %b, %c : !spv.coopmatrix<8x16xi32, Subgroup>, !spv.coopmatrix<8x8xi32, Subgroup> -> !spv.coopmatrix<8x8xi32, Subgroup>
125  spv.Return
126}
127
128// -----
129
130spv.func @cooperative_matrix_muladd(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<16x8xi32, Workgroup>, %c : !spv.coopmatrix<8x8xi32, Subgroup>) "None" {
131  // expected-error @+1 {{'spv.CooperativeMatrixMulAddNV' op matrix scope must match}}
132  %r = spv.CooperativeMatrixMulAddNV %a, %b, %c : !spv.coopmatrix<8x16xi32, Subgroup>, !spv.coopmatrix<16x8xi32, Workgroup> -> !spv.coopmatrix<8x8xi32, Subgroup>
133  spv.Return
134}
135
136// -----
137
138spv.func @cooperative_matrix_muladd(%a : !spv.coopmatrix<8x16xf32, Subgroup>, %b : !spv.coopmatrix<16x8xi32, Subgroup>, %c : !spv.coopmatrix<8x8xi32, Subgroup>) "None" {
139  // expected-error @+1 {{matrix element type must match}}
140  %r = spv.CooperativeMatrixMulAddNV %a, %b, %c : !spv.coopmatrix<8x16xf32, Subgroup>, !spv.coopmatrix<16x8xi32, Subgroup> -> !spv.coopmatrix<8x8xi32, Subgroup>
141  spv.Return
142}
143
144// -----
145
146spv.func @cooperative_matrix_load_memaccess(%ptr : !spv.ptr<!spv.struct<(f32 [0])>, StorageBuffer>, %stride : i32, %b : i1) "None" {
147  // expected-error @+1 {{Pointer must point to a scalar or vector type}}
148  %0 = spv.CooperativeMatrixLoadNV %ptr, %stride, %b : !spv.ptr<!spv.struct<(f32 [0])>, StorageBuffer> as !spv.coopmatrix<8x16xi32, Subgroup>
149  spv.Return
150}
151
152// -----
153
154spv.func @cooperative_matrix_load_memaccess(%ptr : !spv.ptr<i32, Function>, %stride : i32, %b : i1) "None" {
155  // expected-error @+1 {{Pointer storage class must be Workgroup, StorageBuffer or PhysicalStorageBufferEXT}}
156  %0 = spv.CooperativeMatrixLoadNV %ptr, %stride, %b : !spv.ptr<i32, Function> as !spv.coopmatrix<8x16xi32, Subgroup>
157  spv.Return
158}
159