mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td - llvm-project - Git at Google

 //===- SPIRVCooperativeMatrixOps.td - cooperative matmul ---*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This is the op definition spec of cooperative matrix multiply extension ops.
 //
 //===----------------------------------------------------------------------===//

 #ifndef MLIR_DIALECT_SPIRV_IR_COOPERATIVE_MATRIX_OPS
 #define MLIR_DIALECT_SPIRV_IR_COOPERATIVE_MATRIX_OPS

 // -----

 def SPV_CooperativeMatrixLengthNVOp : SPV_Op<"CooperativeMatrixLengthNV",
   [NoSideEffect]> {
   let summary = "See extension SPV_NV_cooperative_matrix";

   let description = [{
     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.

     ``` {.ebnf}
     cooperative-matrix-length-op ::= ssa-id `=` `spv.CooperativeMatrixLengthNV
                                     ` : ` cooperative-matrix-type
     ```

     For example:

     ```
     %0 = spv.CooperativeMatrixLengthNV : !spv.coopmatrix<Subgroup, i32, 8, 16>
     ```
   }];

   let assemblyFormat = "attr-dict `:` $type";

   let availability = [
     MinVersion<SPV_V_1_0>,
     MaxVersion<SPV_V_1_5>,
     Extension<[SPV_NV_cooperative_matrix]>,
     Capability<[SPV_C_CooperativeMatrixNV]>
   ];

   let arguments = (ins
     TypeAttr:$type
   );

   let results = (outs
     SPV_Int32:$result
   );
   let verifier = [{ return success(); }];
 }

 // -----

 def SPV_CooperativeMatrixLoadNVOp : SPV_Op<"CooperativeMatrixLoadNV", []> {
   let summary = "See extension SPV_NV_cooperative_matrix";

   let description = [{
     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.

     ### Custom assembly form

     ``` {.ebnf}
     cooperative-matrixload-op ::= ssa-id `=` `spv.CooperativeMatrixLoadNV`
                               ssa-use `,` ssa-use `,` ssa-use
                               (`[` memory-access `]`)? ` : `
                               pointer-type `as`
                               cooperative-matrix-type
     ```

     For example:

     ```
     %0 = spv.CooperativeMatrixLoadNV %ptr, %stride, %colMajor
          : !spv.ptr<i32, StorageBuffer> as !spv.coopmatrix<i32, Workgroup, 16, 8>
     ```
   }];

   let availability = [
     MinVersion<SPV_V_1_0>,
     MaxVersion<SPV_V_1_5>,
     Extension<[SPV_NV_cooperative_matrix]>,
     Capability<[SPV_C_CooperativeMatrixNV]>
   ];

   let arguments = (ins
     SPV_AnyPtr:$pointer,
     SPV_Integer:$stride,
     SPV_Bool:$columnmajor,
     OptionalAttr<SPV_MemoryAccessAttr>:$memory_access
   );

   let results = (outs
     SPV_AnyCooperativeMatrix:$result
   );

   let verifier = [{
     return verifyPointerAndCoopMatrixType(*this, pointer().getType(),
                                           result().getType());
   }];
 }

 // -----

 def SPV_CooperativeMatrixMulAddNVOp : SPV_Op<"CooperativeMatrixMulAddNV",
   [NoSideEffect, AllTypesMatch<["c", "result"]>]> {
   let summary = "See extension SPV_NV_cooperative_matrix";

   let description = [{
     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).

     ``` {.ebnf}
     cooperative-matrixmuladd-op ::= ssa-id `=` `spv.CooperativeMatrixMulAddNV`
                               ssa-use `,` ssa-use `,` ssa-use ` : `
                               a-cooperative-matrix-type,
                               b-cooperative-matrix-type ->
                               result-cooperative-matrix-type
     ```
     For example:

     ```
     %0 = spv.CooperativeMatrixMulAddNV %arg0, %arg1, %arg2,  :
       !spv.coopmatrix<Subgroup, i32, 8, 16>
     ```
   }];

   let assemblyFormat = [{
     operands attr-dict`:` type($a) `,` type($b) `->` type($c)
   }];

   let availability = [
     MinVersion<SPV_V_1_0>,
     MaxVersion<SPV_V_1_5>,
     Extension<[SPV_NV_cooperative_matrix]>,
     Capability<[SPV_C_CooperativeMatrixNV]>
   ];

   let arguments = (ins
     SPV_AnyCooperativeMatrix:$a,
     SPV_AnyCooperativeMatrix:$b,
     SPV_AnyCooperativeMatrix:$c
   );

   let results = (outs
     SPV_AnyCooperativeMatrix:$result
   );

   let verifier = [{ return verifyCoopMatrixMulAdd(*this); }];
 }

 // -----

 def SPV_CooperativeMatrixStoreNVOp : SPV_Op<"CooperativeMatrixStoreNV", []> {
   let summary = "See extension SPV_NV_cooperative_matrix";

   let description = [{
     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.

     ``` {.ebnf}
     coop-matrix-store-op ::= `spv.CooperativeMatrixStoreNV `
                               ssa-use `, ` ssa-use `, `
                               ssa-use `, ` ssa-use `, `
                               (`[` memory-access `]`)? `:`
                               pointer-type `,` spirv-element-type
     ```

     For example:

     ```
       spv.CooperativeMatrixStoreNV %arg0, %arg2, %arg1, %arg3 :
         !spv.ptr<i32, StorageBuffer>, !spv.coopmatrix<Workgroup, i32, 16, 8>
     ```
   }];

   let availability = [
     MinVersion<SPV_V_1_0>,
     MaxVersion<SPV_V_1_5>,
     Extension<[SPV_NV_cooperative_matrix]>,
     Capability<[SPV_C_CooperativeMatrixNV]>
   ];

   let arguments = (ins
     SPV_AnyPtr:$pointer,
     SPV_AnyCooperativeMatrix:$object,
     SPV_Integer:$stride,
     SPV_Bool:$columnmajor,
     OptionalAttr<SPV_MemoryAccessAttr>:$memory_access
   );

   let results = (outs);

   let verifier = [{
     return verifyPointerAndCoopMatrixType(*this, pointer().getType(),
                                           object().getType());
   }];
 }

 // -----

 #endif // MLIR_DIALECT_SPIRV_IR_COOPERATIVE_MATRIX_OPS
	//===- SPIRVCooperativeMatrixOps.td - cooperative matmul ---- tablegen --===//
	//
	// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
	//
	//===----------------------------------------------------------------------===//
	//
	// This is the op definition spec of cooperative matrix multiply extension ops.
	//
	//===----------------------------------------------------------------------===//

	#ifndef MLIR_DIALECT_SPIRV_IR_COOPERATIVE_MATRIX_OPS
	#define MLIR_DIALECT_SPIRV_IR_COOPERATIVE_MATRIX_OPS

	// -----

	def SPV_CooperativeMatrixLengthNVOp : SPV_Op<"CooperativeMatrixLengthNV",
	[NoSideEffect]> {
	let summary = "See extension SPV_NV_cooperative_matrix";

	let description = [{
	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.

	``` {.ebnf}
	cooperative-matrix-length-op ::= ssa-id `=` `spv.CooperativeMatrixLengthNV
	` : ` cooperative-matrix-type
	```

	For example:

	```
	%0 = spv.CooperativeMatrixLengthNV : !spv.coopmatrix<Subgroup, i32, 8, 16>
	```
	}];

	let assemblyFormat = "attr-dict `:` $type";

	let availability = [
	MinVersion<SPV_V_1_0>,
	MaxVersion<SPV_V_1_5>,
	Extension<[SPV_NV_cooperative_matrix]>,
	Capability<[SPV_C_CooperativeMatrixNV]>
	];

	let arguments = (ins
	TypeAttr:$type
	);

	let results = (outs
	SPV_Int32:$result
	);
	let verifier = [{ return success(); }];
	}

	// -----

	def SPV_CooperativeMatrixLoadNVOp : SPV_Op<"CooperativeMatrixLoadNV", []> {
	let summary = "See extension SPV_NV_cooperative_matrix";

	let description = [{
	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.

	### Custom assembly form

	``` {.ebnf}
	cooperative-matrixload-op ::= ssa-id `=` `spv.CooperativeMatrixLoadNV`
	ssa-use `,` ssa-use `,` ssa-use
	(`[` memory-access `]`)? ` : `
	pointer-type `as`
	cooperative-matrix-type
	```

	For example:

	```
	%0 = spv.CooperativeMatrixLoadNV %ptr, %stride, %colMajor
	: !spv.ptr<i32, StorageBuffer> as !spv.coopmatrix<i32, Workgroup, 16, 8>
	```
	}];

	let availability = [
	MinVersion<SPV_V_1_0>,
	MaxVersion<SPV_V_1_5>,
	Extension<[SPV_NV_cooperative_matrix]>,
	Capability<[SPV_C_CooperativeMatrixNV]>
	];

	let arguments = (ins
	SPV_AnyPtr:$pointer,
	SPV_Integer:$stride,
	SPV_Bool:$columnmajor,
	OptionalAttr<SPV_MemoryAccessAttr>:$memory_access
	);

	let results = (outs
	SPV_AnyCooperativeMatrix:$result
	);

	let verifier = [{
	return verifyPointerAndCoopMatrixType(*this, pointer().getType(),
	result().getType());
	}];
	}

	// -----

	def SPV_CooperativeMatrixMulAddNVOp : SPV_Op<"CooperativeMatrixMulAddNV",
	[NoSideEffect, AllTypesMatch<["c", "result"]>]> {
	let summary = "See extension SPV_NV_cooperative_matrix";

	let description = [{
	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).

	``` {.ebnf}
	cooperative-matrixmuladd-op ::= ssa-id `=` `spv.CooperativeMatrixMulAddNV`
	ssa-use `,` ssa-use `,` ssa-use ` : `
	a-cooperative-matrix-type,
	b-cooperative-matrix-type ->
	result-cooperative-matrix-type
	```
	For example:

	```
	%0 = spv.CooperativeMatrixMulAddNV %arg0, %arg1, %arg2, :
	!spv.coopmatrix<Subgroup, i32, 8, 16>
	```
	}];

	let assemblyFormat = [{
	operands attr-dict`:` type($a) `,` type($b) `->` type($c)
	}];

	let availability = [
	MinVersion<SPV_V_1_0>,
	MaxVersion<SPV_V_1_5>,
	Extension<[SPV_NV_cooperative_matrix]>,
	Capability<[SPV_C_CooperativeMatrixNV]>
	];

	let arguments = (ins
	SPV_AnyCooperativeMatrix:$a,
	SPV_AnyCooperativeMatrix:$b,
	SPV_AnyCooperativeMatrix:$c
	);

	let results = (outs
	SPV_AnyCooperativeMatrix:$result
	);

	let verifier = [{ return verifyCoopMatrixMulAdd(*this); }];
	}

	// -----

	def SPV_CooperativeMatrixStoreNVOp : SPV_Op<"CooperativeMatrixStoreNV", []> {
	let summary = "See extension SPV_NV_cooperative_matrix";

	let description = [{
	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.

	``` {.ebnf}
	coop-matrix-store-op ::= `spv.CooperativeMatrixStoreNV `
	ssa-use `, ` ssa-use `, `
	ssa-use `, ` ssa-use `, `
	(`[` memory-access `]`)? `:`
	pointer-type `,` spirv-element-type
	```

	For example:

	```
	spv.CooperativeMatrixStoreNV %arg0, %arg2, %arg1, %arg3 :
	!spv.ptr<i32, StorageBuffer>, !spv.coopmatrix<Workgroup, i32, 16, 8>
	```
	}];

	let availability = [
	MinVersion<SPV_V_1_0>,
	MaxVersion<SPV_V_1_5>,
	Extension<[SPV_NV_cooperative_matrix]>,
	Capability<[SPV_C_CooperativeMatrixNV]>
	];

	let arguments = (ins
	SPV_AnyPtr:$pointer,
	SPV_AnyCooperativeMatrix:$object,
	SPV_Integer:$stride,
	SPV_Bool:$columnmajor,
	OptionalAttr<SPV_MemoryAccessAttr>:$memory_access
	);

	let results = (outs);

	let verifier = [{
	return verifyPointerAndCoopMatrixType(*this, pointer().getType(),
	object().getType());
	}];
	}

	// -----

	#endif // MLIR_DIALECT_SPIRV_IR_COOPERATIVE_MATRIX_OPS