Name Strings

SPV_NV_cooperative_matrix

Contact

To report problems with this extension, please open a new issue at:

https://github.com/KhronosGroup/SPIRV-Registry

Contributors

  • Jeff Bolz, NVIDIA

  • Daniel Koch, NVIDIA

  • Markus Tavenrath, NVIDIA

Notice

Copyright (c) 2018-2019 NVIDIA Corp.

Status

  • Complete

Version

Last Modified Date

2019-07-12

Revision

2

Dependencies

This extension is written against the SPIR-V Specification, Version 1.3, Revision 5, Unified.

This extension requires SPIR-V 1.3.

This extension requires SPV_KHR_vulkan_memory_model.

This extension interacts with SPV_EXT_physical_storage_buffer.

Overview

This extension adds a new set of types known as "cooperative matrix" types, where the storage for and computations performed on the matrix are spread across a set of invocations such as a subgroup. These types give the implementation freedom in how to optimize matrix multiplies.

This extension introduces the types and instructions, but does not specify rules about what sizes/combinations are valid. This is left to the client API specs, and it is expected that different implementations may support different sizes. To help accommodate this, the dimensions of the cooperative types can be specialized via specialization constants. Since the scope parameter is also something that could potentially be specialized, this extension allows all scope ids to be specialization constants.

Extension Name

To use this extension within a SPIR-V module, the following OpExtension must be present in the module:

OpExtension "SPV_NV_cooperative_matrix"

Modifications to the SPIR-V Specification, Version 1.3

2.2 Terms

Add new terms to section 2.2.2 Types:

Cooperative Matrix: A two-dimensional ordered collection of scalars, whose storage is spread across multiple shader invocations.

Add Cooperative Matrix to the definition of Concrete Type.

Add Cooperative Matrix to the definition of Composite. As a composite, a cooperative matrix acts like a vector with an implementation-dependent number of components (which can be queried with OpCooperativeMatrixLengthNV). It can be used as a composite for all operations that act on composite types, including OpCompositeExtract, OpCompositeInsert, OpAccessChain, etc.. The mapping of components to invocations and indexes is implementation-dependent.

2.16 Validation Rules

Modify section 2.16.1. Universal Validation Rules:

Add OpCooperativeMatrixLoadNV and OpCooperativeMatrixStoreNV to the list of instructions under "A pointer can only be an operand to the following instructions:", when the Logical addressing model is selected and the VariablePointers capability is not declared.

Cooperative matrix types (or types containing them) can only be allocated in Function or Private storage classes.

Modify section 2.16.2. Shader Validation Rules, to relax the rules for scopes to allow specialization constants:

All <id> used for Scope and Memory Semantics must be the result of a constant instruction.

3.26 Memory Access

Modify Section 3.26, "Memory Access":

In the description of MakePointerAvailableKHR added by SPV_KHR_vulkan_memory_model, change "Not valid with OpLoad" to "Not valid with OpLoad or OpCooperativeMatrixLoadNV".

In the description of MakePointerVisibleKHR added by SPV_KHR_vulkan_memory_model, change "Not valid with OpStore" to "Not valid with OpStore or OpCooperativeMatrixStoreNV".

3.31 Capabilities

Modify Section 3.31, "Capability", adding these rows to the Capability table:

Capability Enabling Capabilities

5357

CooperativeMatrixNV
Enables cooperative matrix types and instructions operating on them.

Shader

3.32.6 Type-Declaration Instructions

OpTypeCooperativeMatrixNV

Declare a new cooperative matrix type, where all invocations in Scope cooperate to compute and store the matrix.

Component Type must be a scalar numerical type.

Rows must be a constant instruction with scalar integer type.

Columns must be a constant instruction with scalar integer type.

Capability:
CooperativeMatrixNV

6

5358

Result <id>

<id>
Component Type

Scope <id>
Scope

<id>
Rows

<id>
Columns

3.32.7 Constant-Creation Instructions

Modify OpConstantComposite to make an exception for cooperative matrix types: "If the Result Type is a cooperative matrix type, then there must be only one Constituent and it is used to initialize all members."

Modify OpSpecConstantOp to add: If the CooperativeMatrixNV capability was declared, the following opcode is also valid: OpCooperativeMatrixLengthNV. Relax the limitation on operands for this instruction to allow the operand to be a cooperative matrix type.

3.32.8 Memory Instructions

OpCooperativeMatrixLoadNV

Load a cooperative matrix through a pointer.

Result Type is the type of the loaded object. It must be a cooperative matrix type.

Pointer is a pointer into an array. Its type must be an OpTypePointer whose Type operand is a scalar or vector type. The storage class of Pointer must be Workgroup, StorageBuffer, or (if SPV_EXT_physical_storage_buffer is supported) PhysicalStorageBufferEXT.

Stride is the number of elements in the array in memory between the first component of consecutive rows (or columns) in the result. It must be a scalar integer type.

ColumnMajor indicates whether the values loaded from memory are arranged in column-major or row-major order. It must be a boolean constant instruction, with false indicating row major and true indicating column major.

Memory Access must be a Memory Access literal. If not present, it is the same as specifying None.

If ColumnMajor is false, then elements (row,*) of the result are taken in order from contiguous locations starting at Pointer[row*Stride]. If ColumnMajor is true, then elements (*,col) of the result are taken in order from contiguous locations starting from Pointer[col*Stride]. Any ArrayStride decoration on Pointer is ignored.

For a given dynamic instance of this instruction, all operands of this instruction must be the same for all invocations in a given scope instance (where the scope is the scope the cooperative matrix type was created with). All invocations in a given scope instance must be active or all must be inactive.

Capability:
CooperativeMatrixNV

6+variable

5359

<id>
Result Type

Result <id>

<id>
Pointer

<id>
Stride

<id>
ColumnMajor

Optional
Memory Access

OpCooperativeMatrixStoreNV

Store a cooperative matrix through a pointer.

Pointer is a pointer into an array. Its type must be an OpTypePointer whose Type operand is a scalar or vector type. The storage class of Pointer must be Workgroup, StorageBuffer, or (if SPV_EXT_physical_storage_buffer is supported) PhysicalStorageBufferEXT.

Object is the object to store. Its type must be an OpTypeCooperativeMatrixNV.

Stride is the number of elements in the array in memory between the first component of consecutive rows (or columns) in the result. It must be a scalar integer type.

ColumnMajor indicates whether the values stored to memory are arranged in column-major or row-major order. It must be a boolean constant instruction, with false indicating row major and true indicating column major.

Memory Access must be a Memory Access literal. If not present, it is the same as specifying None.

If ColumnMajor is false, then elements (row,*) of Object are stored in order to contiguous locations starting at Pointer[row*Stride]. If ColumnMajor is true, then elements (*,col) of Object are stored in order to contiguous locations starting from Pointer[col*Stride]. Any ArrayStride decoration Pointer is ignored.

For a given dynamic instance of this instruction, all operands of this instruction must be the same for all invocations in a given scope instance (where the scope is the scope the cooperative matrix type was created with). All invocations in a given scope instance must be active or all must be inactive.

Capability:
CooperativeMatrixNV

5+variable

5360

<id>
Pointer

<id>
Object

<id>
Stride

<id>
ColumnMajor

Optional
Memory Access

OpCooperativeMatrixLengthNV

Number of components of a cooperative matrix type accessible to each invocation when treated as a composite.

Result Type must be an OpTypeInt with 32-bit Width and 0 Signedness.

Type is a cooperative matrix type.

Capability:
CooperativeMatrixNV

4

5362

<id>
Result Type

Result <id>

<id>
Type

3.32.11 Conversion Instructions

Allow cooperative matrix types for the following conversion instructions (if the component types are appropriate): OpConvertFToU, OpConvertFToS, OpConvertSToF, OpConvertUToF, OpUConvert, OpSConvert, OpFConvert. The result type and value type must have the same scope, number of rows, and number of columns.

3.32.12 Composite Instructions

Modify OpCompositeConstruct to make an exception for cooperative matrix types: "If the Result Type is a cooperative matrix type, then there must be only one Constituent and it is used to initialize all members."

3.32.13 Arithmetic Instructions

OpCooperativeMatrixMulAddNV

Linear-algebraic matrix multiply of A by B and then component-wise add C. The order of the operations is implementation-dependent. The internal precision of floating-point operations is defined by the client API. Integer operations are performed at the precision of the Result Type and are exact unless there is overflow or underflow, in which case the result is undefined.

Result Type must be a cooperative matrix type with M rows and N columns.

A is a cooperative matrix with M rows and K columns.

B is a cooperative matrix with K rows and N columns.

C is a cooperative matrix with M rows and N columns.

The values of M, N, and K must be consistent across the result and operands. This is referred to as an MxNxK matrix multiply.

A, B, C, and Result Type must have the same scope, and this defines the scope of the operation. A, B, C, and Result Type need not necessarily have the same component type, this is defined by the client API.

If the Component Type of any matrix operand is an integer type, then its components are treated as signed if its Component Type has Signedness of 1 and are treated as unsigned otherwise.

For a given dynamic instance of this instruction, all invocations in a given scope instance must be active or all must be inactive (where the scope is the scope of the operation).

Capability:
CooperativeMatrixNV

6

5361

<id>
Result Type

Result <id>

<id>
A

<id>
B

<id>
C'

Allow cooperative matrix types for the following instructions (if the component type is appropriate): OpSNegate, OpFNegate, OpIAdd, OpFAdd, OpISub, OpFSub, OpFDiv, OpSDiv, and OpUDiv. Allow cooperative matrix types for OpMatrixTimesScalar.

Issues

1) Should this functionality hardwire a scope (subgroup) or be more flexible?

Discussion: Volta hardware used a smaller scope (8 threads), and it is plausible that other implemenations could also want a smaller scope.

Resolution: Allow a specialization constant scope.

2) Should we have capabilities for each MxNxK matrix multiply "size" that is supported?

Discussion: It’s nice for validation if the shader instructions can be validated solely based on the OpCapability instructions. But that already breaks down for spec-constant-defined cooperative matrix types.

Resolution: Just one capability for the overall feature.

3) Should stride be in bytes or elements?

Discussion: Using elements helps avoid the unsupportable (or more difficult to support) cases.

Resolution: Stride is in elements of the pointee type (which can be different than the matrix component type).

4) Should we allow matrices to be stored in an opaque layout in shared memory?

Discussion: Currently matrices can be stored to shared memory as an array of floats, in row- or column-major order. It might be beneficial to let shaders spill matrices to shared memory in an opaque, implementation-dependent layout. There are a few possible ways to handle this: (1) Reuse the existing OpCooperativeMatrixLoad/Store opcodes with a flag or a trick like Stride=0, (2) new instructions without the stride parameter, (3) let the cooperative matrix types be placed in shared memory and use OpLoad/OpStore.

Resolution: Not supported in the current extension.

5) Should the "column major" operand be a literal constant, or a constant instruction?

Discussion: Constant instructions are more general, and easier for code generation.

Resolution: Constant instruction.

6) Should we allow OpTranspose on cooperative matrix types?

Discussion: In NVIDIA’s initial implementation, we’ll support a pretty restricted set of sizes where the transpose of a matrix will sometimes not be a valid type. So it’s unclear if this is useful.

Resolved: Not supported in this extension.

7) What should the Pointer operand to a cooperative Load/Store be?

Discussion: The spec currently chooses to have the Pointer parameter point at the first element of the matrix in memory, and this pointer is assumed to be in the middle of an array. Another option would be to have the Pointer parameter be a pointer to the whole array, and have an additional "Element" parameter to the instructions, which indicates where the matrix starts in the array.

The alternative option’s main benefit is that you don’t end up with a pointer parameter being used to access something it does not point to. However, it effectively splits out the last element of the access chain into the load/store instruction, which is kind of weird. And in the first option, the pointer to the array is still there implicitly in the access chain.

Resolution: Pointer points to the first element of the array.

8) Should we allow the Pointer type and matrix component type to mismatch?

Resolution: Yes, this makes it easier to efficiently load matrix data into shared memory, which can be declared to use a larger type (e.g. uvec4). The Stride parameter is interpreted in units of the pointed-to type, not in units of the matrix’s component type.

Revision History

Rev Date Author Changes

1

2019-01-30

Jeff Bolz

Initial revision

2

2019-07-12

Jeff Bolz

Added details for integer operations