1// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
2
3spv.module Logical GLSL450 requires #spv.vce<v1.0, [CooperativeMatrixNV], [SPV_NV_cooperative_matrix]> {
4  // CHECK-LABEL: @cooperative_matrix_load
5  spv.func @cooperative_matrix_load(%ptr : !spv.ptr<i32, StorageBuffer>, %stride : i32, %b : i1) "None" {
6    // CHECK: {{%.*}} = spv.CooperativeMatrixLoadNV {{%.*}}, {{%.*}}, {{%.*}} : !spv.ptr<i32, StorageBuffer> as !spv.coopmatrix<16x8xi32, Workgroup>
7    %0 = spv.CooperativeMatrixLoadNV %ptr, %stride, %b : !spv.ptr<i32, StorageBuffer> as !spv.coopmatrix<16x8xi32, Workgroup>
8    spv.Return
9  }
10
11  // CHECK-LABEL: @cooperative_matrix_load_memaccess
12  spv.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_store
19  spv.func @cooperative_matrix_store(%ptr : !spv.ptr<i32, StorageBuffer>, %stride : i32, %m : !spv.coopmatrix<16x8xi32, Workgroup>, %b : i1) "None" {
20    // CHECK: spv.CooperativeMatrixStoreNV {{%.*}}, {{%.*}}, {{%.*}} : !spv.ptr<i32, StorageBuffer>, !spv.coopmatrix<16x8xi32, Workgroup>
21    spv.CooperativeMatrixStoreNV %ptr, %m, %stride, %b : !spv.ptr<i32, StorageBuffer>, !spv.coopmatrix<16x8xi32, Workgroup>
22    spv.Return
23  }
24
25  // CHECK-LABEL: @cooperative_matrix_store_memaccess
26  spv.func @cooperative_matrix_store_memaccess(%ptr : !spv.ptr<i32, StorageBuffer>, %m : !spv.coopmatrix<8x16xi32, Subgroup>, %stride : i32, %b : i1) "None" {
27    // CHECK: spv.CooperativeMatrixStoreNV {{%.*}}, {{%.*}}, {{%.*}} ["Volatile"] : !spv.ptr<i32, StorageBuffer>, !spv.coopmatrix<8x16xi32, Subgroup>
28    spv.CooperativeMatrixStoreNV %ptr, %m, %stride, %b ["Volatile"] : !spv.ptr<i32, StorageBuffer>, !spv.coopmatrix<8x16xi32, Subgroup>
29    spv.Return
30  }
31
32  // CHECK-LABEL: @cooperative_matrix_length
33  spv.func @cooperative_matrix_length() -> i32 "None" {
34    // CHECK: {{%.*}} = spv.CooperativeMatrixLengthNV : !spv.coopmatrix<8x16xi32, Subgroup>
35    %0 = spv.CooperativeMatrixLengthNV : !spv.coopmatrix<8x16xi32, Subgroup>
36    spv.ReturnValue %0 : i32
37  }
38
39  // CHECK-LABEL: @cooperative_matrix_muladd
40  spv.func @cooperative_matrix_muladd(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<16x8xi32, Subgroup>, %c : !spv.coopmatrix<8x8xi32, Subgroup>) "None" {
41    // CHECK: {{%.*}} = spv.CooperativeMatrixMulAddNV {{%.*}}, {{%.*}}, {{%.*}}  : !spv.coopmatrix<8x16xi32, Subgroup>, !spv.coopmatrix<16x8xi32, Subgroup> -> !spv.coopmatrix<8x8xi32, Subgroup>
42    %r = spv.CooperativeMatrixMulAddNV %a, %b, %c : !spv.coopmatrix<8x16xi32, Subgroup>, !spv.coopmatrix<16x8xi32, Subgroup> -> !spv.coopmatrix<8x8xi32, Subgroup>
43    spv.Return
44  }
45
46  // CHECK-LABEL: @cooperative_matrix_add
47  spv.func @cooperative_matrix_add(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
48    // CHECK: {{%.*}} = spv.IAdd {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
49    %r = spv.IAdd %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
50    spv.Return
51  }
52
53  // CHECK-LABEL: @cooperative_matrix_sub
54  spv.func @cooperative_matrix_sub(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
55    // CHECK: {{%.*}} = spv.ISub {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
56    %r = spv.ISub %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
57    spv.Return
58  }
59
60  // CHECK-LABEL: @cooperative_matrix_sdiv
61  spv.func @cooperative_matrix_sdiv(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
62    // CHECK: {{%.*}} = spv.SDiv {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
63    %r = spv.SDiv %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
64    spv.Return
65  }
66
67  // CHECK-LABEL: @cooperative_matrix_udiv
68  spv.func @cooperative_matrix_udiv(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
69    // CHECK: {{%.*}} = spv.UDiv {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
70    %r = spv.UDiv %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
71    spv.Return
72  }
73
74  // CHECK-LABEL: @cooperative_matrix_fadd
75  spv.func @cooperative_matrix_fadd(%a : !spv.coopmatrix<8x16xf32, Subgroup>, %b : !spv.coopmatrix<8x16xf32, Subgroup>) "None" {
76    // CHECK: {{%.*}} = spv.FAdd {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup>
77    %r = spv.FAdd %a, %b : !spv.coopmatrix<8x16xf32, Subgroup>
78    spv.Return
79  }
80
81  // CHECK-LABEL: @cooperative_matrix_fsub
82  spv.func @cooperative_matrix_fsub(%a : !spv.coopmatrix<8x16xf32, Subgroup>, %b : !spv.coopmatrix<8x16xf32, Subgroup>) "None" {
83    // CHECK: {{%.*}} = spv.FSub {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup>
84    %r = spv.FSub %a, %b : !spv.coopmatrix<8x16xf32, Subgroup>
85    spv.Return
86  }
87
88  // CHECK-LABEL: @cooperative_matrix_fdiv
89  spv.func @cooperative_matrix_fdiv(%a : !spv.coopmatrix<8x16xf32, Subgroup>, %b : !spv.coopmatrix<8x16xf32, Subgroup>) "None" {
90    // CHECK: {{%.*}} = spv.FDiv {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup>
91    %r = spv.FDiv %a, %b : !spv.coopmatrix<8x16xf32, Subgroup>
92    spv.Return
93  }
94
95  // CHECK-LABEL: @cooperative_matrix_access_chain
96  spv.func @cooperative_matrix_access_chain(%a : !spv.ptr<!spv.coopmatrix<8x16xf32, Subgroup>, Function>) -> !spv.ptr<f32, Function> "None" {
97    %0 = spv.constant 0: i32
98    // CHECK: {{%.*}} = spv.AccessChain {{%.*}}[{{%.*}}] : !spv.ptr<!spv.coopmatrix<8x16xf32, Subgroup>, Function>, i32
99    %1 = spv.AccessChain %a[%0] : !spv.ptr<!spv.coopmatrix<8x16xf32, Subgroup>, Function>, i32
100    spv.ReturnValue %1 : !spv.ptr<f32, Function>
101  }
102}
103