• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1//===- SPIRVCooperativeMatrixOps.td - cooperative matmul ---*- tablegen -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This is the op definition spec of cooperative matrix multiply extension ops.
10//
11//===----------------------------------------------------------------------===//
12
13#ifndef SPIRV_COOPERATIVE_MATRIX_OPS
14#define SPIRV_COOPERATIVE_MATRIX_OPS
15
16// -----
17
18def SPV_CooperativeMatrixLengthNVOp : SPV_Op<"CooperativeMatrixLengthNV",
19  [NoSideEffect]> {
20  let summary = "See extension SPV_NV_cooperative_matrix";
21
22  let description = [{
23    Number of components of a cooperative matrix type accessible to each
24    invocation when treated as a composite.
25
26    Result Type must be an OpTypeInt with 32-bit Width and 0 Signedness.
27
28    Type is a cooperative matrix type.
29
30    ``` {.ebnf}
31    cooperative-matrix-length-op ::= ssa-id `=` `spv.CooperativeMatrixLengthNV
32                                    ` : ` cooperative-matrix-type
33    ```
34
35    For example:
36
37    ```
38    %0 = spv.CooperativeMatrixLengthNV : !spv.coopmatrix<Subgroup, i32, 8, 16>
39    ```
40  }];
41
42  let assemblyFormat = "attr-dict `:` $type";
43
44  let availability = [
45    MinVersion<SPV_V_1_0>,
46    MaxVersion<SPV_V_1_5>,
47    Extension<[SPV_NV_cooperative_matrix]>,
48    Capability<[SPV_C_CooperativeMatrixNV]>
49  ];
50
51  let arguments = (ins
52    TypeAttr:$type
53  );
54
55  let results = (outs
56    SPV_Int32:$result
57  );
58  let verifier = [{ return success(); }];
59}
60
61// -----
62
63def SPV_CooperativeMatrixLoadNVOp : SPV_Op<"CooperativeMatrixLoadNV", []> {
64  let summary = "See extension SPV_NV_cooperative_matrix";
65
66  let description = [{
67    Load a cooperative matrix through a pointer.
68
69    Result Type is the type of the loaded object. It must be a cooperative
70    matrix type.
71
72    Pointer is a pointer into an array. Its type must be an OpTypePointer whose
73    Type operand is a scalar or vector type. The storage class of Pointer must
74    be Workgroup, StorageBuffer, or (if SPV_EXT_physical_storage_buffer is
75    supported) PhysicalStorageBufferEXT.
76
77    Stride is the number of elements in the array in memory between the first
78    component of consecutive rows (or columns) in the result. It must be a
79    scalar integer type.
80
81    ColumnMajor indicates whether the values loaded from memory are arranged in
82    column-major or row-major order. It must be a boolean constant instruction,
83    with false indicating row major and true indicating column major.
84
85    Memory Access must be a Memory Access literal. If not present, it is the
86    same as specifying None.
87
88    If ColumnMajor is false, then elements (row,*) of the result are taken in
89    order from contiguous locations starting at Pointer[row*Stride]. If
90    ColumnMajor is true, then elements (*,col) of the result are taken in order
91    from contiguous locations starting from Pointer[col*Stride]. Any ArrayStride
92    decoration on Pointer is ignored.
93
94    For a given dynamic instance of this instruction, all operands of this
95    instruction must be the same for all invocations in a given scope instance
96    (where the scope is the scope the cooperative matrix type was created with).
97    All invocations in a given scope instance must be active or all must be
98    inactive.
99
100    ### Custom assembly form
101
102    ``` {.ebnf}
103    cooperative-matrixload-op ::= ssa-id `=` `spv.CooperativeMatrixLoadNV`
104                              ssa-use `,` ssa-use `,` ssa-use
105                              (`[` memory-access `]`)? ` : `
106                              pointer-type `as`
107                              cooperative-matrix-type
108    ```
109
110    For example:
111
112    ```
113    %0 = spv.CooperativeMatrixLoadNV %ptr, %stride, %colMajor
114         : !spv.ptr<i32, StorageBuffer> as !spv.coopmatrix<i32, Workgroup, 16, 8>
115    ```
116  }];
117
118  let availability = [
119    MinVersion<SPV_V_1_0>,
120    MaxVersion<SPV_V_1_5>,
121    Extension<[SPV_NV_cooperative_matrix]>,
122    Capability<[SPV_C_CooperativeMatrixNV]>
123  ];
124
125  let arguments = (ins
126    SPV_AnyPtr:$pointer,
127    SPV_Integer:$stride,
128    SPV_Bool:$columnmajor,
129    OptionalAttr<SPV_MemoryAccessAttr>:$memory_access
130  );
131
132  let results = (outs
133    SPV_AnyCooperativeMatrix:$result
134  );
135
136  let verifier = [{
137    return verifyPointerAndCoopMatrixType(*this, pointer().getType(),
138                                          result().getType());
139  }];
140}
141
142// -----
143
144def SPV_CooperativeMatrixMulAddNVOp : SPV_Op<"CooperativeMatrixMulAddNV",
145  [NoSideEffect, AllTypesMatch<["c", "result"]>]> {
146  let summary = "See extension SPV_NV_cooperative_matrix";
147
148  let description = [{
149    Linear-algebraic matrix multiply of A by B and then component-wise add C.
150    The order of the operations is implementation-dependent. The internal
151    precision of floating-point operations is defined by the client API.
152    Integer operations are performed at the precision of the Result Type and are
153    exact unless there is overflow or underflow, in which case the result is
154    undefined.
155
156    Result Type must be a cooperative matrix type with M rows and N columns.
157
158    A is a cooperative matrix with M rows and K columns.
159
160    B is a cooperative matrix with K rows and N columns.
161
162    C is a cooperative matrix with M rows and N columns.
163
164    The values of M, N, and K must be consistent across the result and operands.
165    This is referred to as an MxNxK matrix multiply.
166
167    A, B, C, and Result Type must have the same scope, and this defines the
168    scope of the operation. A, B, C, and Result Type need not necessarily have
169    the same component type, this is defined by the client API.
170
171    If the Component Type of any matrix operand is an integer type, then its
172    components are treated as signed if its Component Type has Signedness of 1
173    and are treated as unsigned otherwise.
174
175    For a given dynamic instance of this instruction, all invocations in a given
176    scope instance must be active or all must be inactive (where the scope is
177    the scope of the operation).
178
179    ``` {.ebnf}
180    cooperative-matrixmuladd-op ::= ssa-id `=` `spv.CooperativeMatrixMulAddNV`
181                              ssa-use `,` ssa-use `,` ssa-use ` : `
182                              a-cooperative-matrix-type,
183                              b-cooperative-matrix-type ->
184                              result-cooperative-matrix-type
185    ```
186    For example:
187
188    ```
189    %0 = spv.CooperativeMatrixMulAddNV %arg0, %arg1, %arg2,  :
190      !spv.coopmatrix<Subgroup, i32, 8, 16>
191    ```
192  }];
193
194  let assemblyFormat = [{
195    operands attr-dict`:` type($a) `,` type($b) `->` type($c)
196  }];
197
198  let availability = [
199    MinVersion<SPV_V_1_0>,
200    MaxVersion<SPV_V_1_5>,
201    Extension<[SPV_NV_cooperative_matrix]>,
202    Capability<[SPV_C_CooperativeMatrixNV]>
203  ];
204
205  let arguments = (ins
206    SPV_AnyCooperativeMatrix:$a,
207    SPV_AnyCooperativeMatrix:$b,
208    SPV_AnyCooperativeMatrix:$c
209  );
210
211  let results = (outs
212    SPV_AnyCooperativeMatrix:$result
213  );
214
215  let verifier = [{ return verifyCoopMatrixMulAdd(*this); }];
216}
217
218// -----
219
220def SPV_CooperativeMatrixStoreNVOp : SPV_Op<"CooperativeMatrixStoreNV", []> {
221  let summary = "See extension SPV_NV_cooperative_matrix";
222
223  let description = [{
224    Store a cooperative matrix through a pointer.
225
226    Pointer is a pointer into an array. Its type must be an OpTypePointer whose
227    Type operand is a scalar or vector type. The storage class of Pointer must
228    be Workgroup, StorageBuffer, or (if SPV_EXT_physical_storage_buffer is
229    supported) PhysicalStorageBufferEXT.
230
231    Object is the object to store. Its type must be an
232    OpTypeCooperativeMatrixNV.
233
234    Stride is the number of elements in the array in memory between the first
235    component of consecutive rows (or columns) in the result. It must be a
236    scalar integer type.
237
238    ColumnMajor indicates whether the values stored to memory are arranged in
239    column-major or row-major order. It must be a boolean constant instruction,
240    with false indicating row major and true indicating column major.
241
242    Memory Access must be a Memory Access literal. If not present, it is the
243    same as specifying None.
244
245    ``` {.ebnf}
246    coop-matrix-store-op ::= `spv.CooperativeMatrixStoreNV `
247                              ssa-use `, ` ssa-use `, `
248                              ssa-use `, ` ssa-use `, `
249                              (`[` memory-access `]`)? `:`
250                              pointer-type `,` spirv-element-type
251    ```
252
253    For example:
254
255    ```
256      spv.CooperativeMatrixStoreNV %arg0, %arg2, %arg1, %arg3 :
257        !spv.ptr<i32, StorageBuffer>, !spv.coopmatrix<Workgroup, i32, 16, 8>
258    ```
259  }];
260
261  let availability = [
262    MinVersion<SPV_V_1_0>,
263    MaxVersion<SPV_V_1_5>,
264    Extension<[SPV_NV_cooperative_matrix]>,
265    Capability<[SPV_C_CooperativeMatrixNV]>
266  ];
267
268  let arguments = (ins
269    SPV_AnyPtr:$pointer,
270    SPV_AnyCooperativeMatrix:$object,
271    SPV_Integer:$stride,
272    SPV_Bool:$columnmajor,
273    OptionalAttr<SPV_MemoryAccessAttr>:$memory_access
274  );
275
276  let results = (outs);
277
278  let verifier = [{
279    return verifyPointerAndCoopMatrixType(*this, pointer().getType(),
280                                          object().getType());
281  }];
282}
283
284// -----
285
286#endif // SPIRV_COOPERATIVE_MATRIX_OPS
287