tvm::backend::metal Namespace Reference

Functions
TVM_TIRX_DLL const Op &	cooperative_tensor_fill_op ()
	Fill a cooperative_tensor with a given value.
TVM_TIRX_DLL const Op &	cooperative_tensor_load_op ()
	Load data from device or threadgroup memory into a cooperative_tensor.
TVM_TIRX_DLL const Op &	cooperative_tensor_store_op ()
	Store data from a cooperative_tensor to device or threadgroup memory.
TVM_TIRX_DLL const Op &	cooperative_tensor_multiply_accumulate_op ()
	Multiply and accumulate two matrices using cooperative_tensor (MetalPerformancePrimitives matmul2d).

Function Documentation

cooperative_tensor_fill_op()

TVM_TIRX_DLL const Op & tvm::backend::metal::cooperative_tensor_fill_op

(

)

Fill a cooperative_tensor with a given value.

Arguments, in order:

• args[0]: d, The D operand.
• args[1]: index, The index.
• args[2]: value, The value to use.
• args[3]: rows, The number of rows.
• args[4]: cols, The number of columns.

cooperative_tensor_load_op()

TVM_TIRX_DLL const Op & tvm::backend::metal::cooperative_tensor_load_op

(

)

Load data from device or threadgroup memory into a cooperative_tensor.

Arguments, in order:

• args[0]: d, The D operand.
• args[1]: index, The index.
• args[2]: ptr, The pointer.
• args[3]: stride, The stride.
• args[4]: rows, The number of rows.
• args[5]: cols, The number of columns.
• args[6]: transpose_matrix, Whether to transpose the matrix.
• args[7]: mma_M, The M dimension of the matrix operation.
• args[8]: mma_N, The N dimension of the matrix operation.
• args[9]: mma_K, The K dimension of the matrix operation.
• args[10]: operand_role, 0 for left (A), 1 for right (B), 2 for destination (C).

cooperative_tensor_multiply_accumulate_op()

TVM_TIRX_DLL const Op & tvm::backend::metal::cooperative_tensor_multiply_accumulate_op

(

)

Multiply and accumulate two matrices using cooperative_tensor (MetalPerformancePrimitives matmul2d).

Arguments, in order:

• args[0]: d, The D operand.
• args[1]: index_d, The D fragment index.
• args[2]: a, The A operand.
• args[3]: index_a, The A fragment index.
• args[4]: b, The B operand.
• args[5]: index_b, The B fragment index.
• args[6]: c, The C operand.
• args[7]: index_c, The C fragment index.
• args[8]: M, The M dimension.
• args[9]: N, The N dimension.
• args[10]: K, The K dimension.
• args[11]: transpose_a, Whether to transpose A.
• args[12]: transpose_b, Whether to transpose B.

cooperative_tensor_store_op()

TVM_TIRX_DLL const Op & tvm::backend::metal::cooperative_tensor_store_op

(

)

Store data from a cooperative_tensor to device or threadgroup memory.

Arguments, in order:

• args[0]: d, The D operand.
• args[1]: index, The index.
• args[2]: ptr, The pointer.
• args[3]: stride, The stride.
• args[4]: rows, The number of rows.
• args[5]: cols, The number of columns.
• args[6]: transpose_matrix, Whether to transpose the matrix.
• args[7]: mma_M, The M dimension of the matrix operation.
• args[8]: mma_N, The N dimension of the matrix operation.
• args[9]: mma_K, The K dimension of the matrix operation.
• args[10]: operand_role, 0 for left (A), 1 for right (B), 2 for destination (C).