[TIR] Add cooperative_tensor builtins and metal.cooperative_tensor storage scope

Add TIR builtins for Metal cooperative_tensor operations (MetalPerformancePrimitives):
- cooperative_tensor_fill: fill a cooperative_tensor with a value
- cooperative_tensor_load: load from device/threadgroup memory
- cooperative_tensor_store: store to device/threadgroup memory
- cooperative_tensor_multiply_accumulate: matrix multiply-accumulate via matmul2d

Add metal.cooperative_tensor storage scope (StorageRank::kMetalCooperativeTensor)
for buffers backed by MPP cooperative_tensor registers, analogous to the existing
metal.simdgroup scope but targeting the Metal 4 tensor operations API.

These primitives enable code generation for MetalPerformancePrimitives matmul2d,
which routes to NAX tensor cores on Apple M5 and falls back to simdgroup matrix
instructions on M1-M4.

Author: oraluben

include/tvm/tir/builtin.h

@@ -806,6 +806,51 @@ TVM_DLL const Op& simdgroup_store();
  */
 TVM_DLL const Op& simdgroup_multiply_accumulate();
 
+// Metal cooperative_tensor intrinsics (MetalPerformancePrimitives / Metal 4)
+
+/*!
+ * \brief Fill a cooperative_tensor with a given value.
+ *
+ * void cooperative_tensor_fill(Var d, PrimExpr index, PrimExpr value,
+ *                              int rows, int cols);
+ */
+TVM_DLL const Op& cooperative_tensor_fill();
+
+/*!
+ * \brief Load data from device or threadgroup memory into a cooperative_tensor.
+ *
+ * void cooperative_tensor_load(Var d, PrimExpr index, PrimExpr ptr,
+ *                              PrimExpr stride, int rows, int cols,
+ *                              bool transpose_matrix,
+ *                              int mma_M, int mma_N, int mma_K,
+ *                              int operand_role);
+ * operand_role: 0=left(A), 1=right(B), 2=destination(C)
+ */
+TVM_DLL const Op& cooperative_tensor_load();
+
+/*!
+ * \brief Store data from a cooperative_tensor to device or threadgroup memory.
+ *
+ * void cooperative_tensor_store(Var d, PrimExpr index, PrimExpr ptr,
+ *                               PrimExpr stride, int rows, int cols,
+ *                               bool transpose_matrix,
+ *                               int mma_M, int mma_N, int mma_K,
+ *                               int operand_role);
+ * operand_role: 0=left(A), 1=right(B), 2=destination(C)
+ */
+TVM_DLL const Op& cooperative_tensor_store();
+
+/*!
+ * \brief Multiply and accumulate two matrices using cooperative_tensor
+ *        (MetalPerformancePrimitives matmul2d).
+ *
+ * void cooperative_tensor_multiply_accumulate(
+ *     Var d, PrimExpr index_d, Var a, PrimExpr index_a,
+ *     Var b, PrimExpr index_b, Var c, PrimExpr index_c,
+ *     int M, int N, int K, bool transpose_a, bool transpose_b);
+ */
+TVM_DLL const Op& cooperative_tensor_multiply_accumulate();
+
 // TODO(tvm-team) replace the usage of the vector operations by Shuffle.
 /*!
  * \brief Get the high level half of the vector

python/tvm/script/ir_builder/tir/ir.py

@@ -1430,159 +1430,313 @@ def func(
 
     return func
 
+
 if TYPE_CHECKING:
+
     class int8: ...
+
     class int16: ...
+
     class int32: ...
+
     class int64: ...
+
     class int8x4: ...
+
     class int16x4: ...
+
     class int32x4: ...
+
     class int64x4: ...
+
     class int8x8: ...
+
     class int16x8: ...
+
     class int32x8: ...
+
     class int64x8: ...
+
     class int8x16: ...
+
     class int16x16: ...
+
     class int32x16: ...
+
     class int64x16: ...
+
     class int8x32: ...
+
     class int16x32: ...
+
     class int32x32: ...
+
     class int64x32: ...
+
     class int8x64: ...
+
     class int16x64: ...
+
     class int32x64: ...
+
     class int64x64: ...
+
     class uint8: ...
+
     class uint16: ...
+
     class uint32: ...
+
     class uint64: ...
+
     class uint8x4: ...
+
     class uint16x4: ...
+
     class uint32x4: ...
+
     class uint64x4: ...
+
     class uint8x8: ...
+
     class uint16x8: ...
+
     class uint32x8: ...
+
     class uint64x8: ...
+
     class uint8x16: ...
+
     class uint16x16: ...
+
     class uint32x16: ...
+
     class uint64x16: ...
+
     class uint8x32: ...
+
     class uint16x32: ...
+
     class uint32x32: ...
+
     class uint64x32: ...
+
     class uint8x64: ...
+
     class uint16x64: ...
+
     class uint32x64: ...
+
     class uint64x64: ...
+
     class float16: ...
+
     class float32: ...
+
     class float64: ...
+
     class float16x2: ...
+
     class float32x2: ...
+
     class float64x2: ...
+
     class float16x4: ...
+
     class float32x4: ...
+
     class float64x4: ...
+
     class float16x8: ...
+
     class float32x8: ...
+
     class float64x8: ...
+
     class float16x16: ...
+
     class float32x16: ...
+
     class float64x16: ...
+
     class float16x32: ...
+
     class float32x32: ...
+
     class float64x32: ...
+
     class float16x64: ...
+
     class float32x64: ...
+
     class float64x64: ...
+
     class float8_e3m4: ...
+
     class float8_e3m4x2: ...
+
     class float8_e3m4x4: ...
+
     class float8_e3m4x8: ...
+
     class float8_e3m4x16: ...
+
     class float8_e3m4x32: ...
+
     class float8_e3m4x64: ...
+
     class float8_e4m3: ...
+
     class float8_e4m3x2: ...
+
     class float8_e4m3x4: ...
+
     class float8_e4m3x8: ...
+
     class float8_e4m3x16: ...
+
     class float8_e4m3x32: ...
+
     class float8_e4m3x64: ...
+
     class float8_e4m3b11fnuz: ...
+
     class float8_e4m3b11fnuzx2: ...
+
     class float8_e4m3b11fnuzx4: ...
+
     class float8_e4m3b11fnuzx8: ...
+
     class float8_e4m3b11fnuzx16: ...
+
     class float8_e4m3b11fnuzx32: ...
+
     class float8_e4m3b11fnuzx64: ...
+
     class float8_e4m3fn: ...
+
     class float8_e4m3fnx2: ...
+
     class float8_e4m3fnx4: ...
+
     class float8_e4m3fnx8: ...
+
     class float8_e4m3fnx16: ...
+
     class float8_e4m3fnx32: ...
+
     class float8_e4m3fnx64: ...
+
     class float8_e4m3fnuz: ...
+
     class float8_e4m3fnuzx2: ...
+
     class float8_e4m3fnuzx4: ...
+
     class float8_e4m3fnuzx8: ...
+
     class float8_e4m3fnuzx16: ...
+
     class float8_e4m3fnuzx32: ...
+
     class float8_e4m3fnuzx64: ...
+
     class float8_e5m2: ...
+
     class float8_e5m2x2: ...
+
     class float8_e5m2x4: ...
+
     class float8_e5m2x8: ...
+
     class float8_e5m2x16: ...
+
     class float8_e5m2x32: ...
+
     class float8_e5m2x64: ...
+
     class float8_e5m2fnuz: ...
+
     class float8_e5m2fnuzx2: ...
+
     class float8_e5m2fnuzx4: ...
+
     class float8_e5m2fnuzx8: ...
+
     class float8_e5m2fnuzx16: ...
+
     class float8_e5m2fnuzx32: ...
+
     class float8_e5m2fnuzx64: ...
+
     class float8_e8m0fnu: ...
+
     class float8_e8m0fnux2: ...
+
     class float8_e8m0fnux4: ...
+
     class float8_e8m0fnux8: ...
+
     class float8_e8m0fnux16: ...
+
     class float8_e8m0fnux32: ...
+
     class float8_e8m0fnux64: ...
+
     class float6_e2m3fn: ...
+
     class float6_e2m3fnx2: ...
+
     class float6_e2m3fnx4: ...
+
     class float6_e2m3fnx8: ...
+
     class float6_e2m3fnx16: ...
+
     class float6_e2m3fnx32: ...
+
     class float6_e2m3fnx64: ...
+
     class float6_e3m2fn: ...
+
     class float6_e3m2fnx2: ...
+
     class float6_e3m2fnx4: ...
+
     class float6_e3m2fnx8: ...
+
     class float6_e3m2fnx16: ...
+
     class float6_e3m2fnx32: ...
+
     class float6_e3m2fnx64: ...
+
     class float4_e2m1fn: ...
+
     class float4_e2m1fnx2: ...
+
     class float4_e2m1fnx4: ...
+
     class float4_e2m1fnx8: ...
+
     class float4_e2m1fnx16: ...
+
     class float4_e2m1fnx32: ...
+
     class float4_e2m1fnx64: ...
+
     class bfloat16: ...
+
     class bfloat16x2: ...
+
     class bfloat16x4: ...
+
     class bfloat16x8: ...
+
     class bfloat16x16: ...
+
     class bfloat16x32: ...
+
     class bfloat16x64: ...
 else:
     # pylint: disable=invalid-name
@@ -2202,6 +2356,10 @@ def wrapped(*args, **kwargs):
 simdgroup_load = _op_wrapper(_tir_op.simdgroup_load)
 simdgroup_store = _op_wrapper(_tir_op.simdgroup_store)
 simdgroup_multiply_accumulate = _op_wrapper(_tir_op.simdgroup_multiply_accumulate)
+cooperative_tensor_fill = _op_wrapper(_tir_op.cooperative_tensor_fill)
+cooperative_tensor_load = _op_wrapper(_tir_op.cooperative_tensor_load)
+cooperative_tensor_store = _op_wrapper(_tir_op.cooperative_tensor_store)
+cooperative_tensor_multiply_accumulate = _op_wrapper(_tir_op.cooperative_tensor_multiply_accumulate)
 create_barriers = _op_wrapper(_tir_op.create_barriers)
 assume = _op_wrapper(_tir_op.assume)
 undef = _op_wrapper(_tir_op.undef)
@@ -2500,6 +2658,10 @@ def wrapped(*args, **kwargs):
     "simdgroup_load",
     "simdgroup_store",
     "simdgroup_multiply_accumulate",
+    "cooperative_tensor_fill",
+    "cooperative_tensor_load",
+    "cooperative_tensor_store",
+    "cooperative_tensor_multiply_accumulate",
     "create_barriers",
     "mma_store",
     "mma_fill",