Metal TensorOps for Custom ML Kernels

import Metal
import MetalPerformanceShadersGraph

// MARK: - Host-side MTLTensor setup for quantized matmul (iOS 27 / macOS 27)

final class QuantizedMatMulDemo {
    let device: MTLDevice
    let commandQueue: MTLCommandQueue

    init?() {
        guard let dev = MTLCreateSystemDefaultDevice(),
              let queue = dev.makeCommandQueue() else { return nil }
        self.device = dev
        self.commandQueue = queue
    }

    // Build an MTLTensor holding 4-bit integer weights + FP8-E8M0 block scales
    func makeQuantizedWeightTensor(rows: Int, cols: Int, blockSize: Int = 32) throws -> MTLTensor {
        // 1. Describe the primary data plane: 4-bit signed integers
        let dataDescriptor = MTLTensorDescriptor()
        dataDescriptor.dataType   = .int4          // New in iOS 26; FP8 added iOS 27
        dataDescriptor.dimensions = [rows as NSNumber, cols as NSNumber]

        // 2. Describe the auxiliary scale plane: FP8 E8M0, one scale per block of 32
        let scaleDescriptor = MTLTensorAuxiliaryPlaneDescriptor()
        scaleDescriptor.dataType     = .float8E8M0  // New in iOS 27
        scaleDescriptor.blockFactors = [blockSize as NSNumber, 1]

        // 3. Wire the scale plane into the main tensor descriptor
        let auxMap = MTLTensorAuxiliaryPlaneMap()
        auxMap.scalePlaneDescriptor = scaleDescriptor
        dataDescriptor.auxiliaryPlaneMap = auxMap

        // 4. Allocate the tensor (data + scales packed together)
        guard let tensor = device.makeTensor(descriptor: dataDescriptor) else {
            throw NSError(domain: "QuantizedMatMulDemo",
                          code: -1,
                          userInfo: [NSLocalizedDescriptionKey: "Failed to allocate MTLTensor"])
        }
        return tensor
    }

    // Dispatch the custom MSL TensorOps kernel that performs dequantized matmul
    func runMatMul(weightTensor: MTLTensor,
                   activationTensor: MTLTensor,
                   outputTensor: MTLTensor,
                   library: MTLLibrary) throws {
        guard let kernelFn = library.makeFunction(name: "quantized_matmul_kernel"),
              let pipeline = try? device.makeComputePipelineState(function: kernelFn),
              let cmdBuf  = commandQueue.makeCommandBuffer(),
              let encoder = cmdBuf.makeComputeCommandEncoder() else {
            throw NSError(domain: "QuantizedMatMulDemo", code: -2, userInfo: nil)
        }

        encoder.setComputePipelineState(pipeline)

        // Bind MTLTensors to buffer indices; the MSL kernel reads planes via tensor_handle
        encoder.setTensor(weightTensor,     index: 0)
        encoder.setTensor(activationTensor, index: 1)
        encoder.setTensor(outputTensor,     index: 2)

        // Threadgroup geometry — tune to your tile size
        let threadgroupSize  = MTLSize(width: 32, height: 4, depth: 1)
        let rows = outputTensor.descriptor.dimensions[0].intValue
        let cols = outputTensor.descriptor.dimensions[1].intValue
        let gridSize = MTLSize(
            width:  (cols + threadgroupSize.width  - 1) / threadgroupSize.width,
            height: (rows + threadgroupSize.height - 1) / threadgroupSize.height,
            depth:  1
        )

        encoder.dispatchThreadgroups(gridSize, threadsPerThreadgroup: threadgroupSize)
        encoder.endEncoding()
        cmdBuf.commit()
        cmdBuf.waitUntilCompleted()
    }
}

/* ---------- MSL kernel (quantized_matmul_kernel.metal) ----------
   // Runs on-GPU with TensorOps; shown here as a Swift raw string for reference.
   // In a real project this lives in a .metal file.

   #include <metal_stdlib>
   #include <metal_tensor_ops>
   using namespace metal;

   using scale_plane_t   = tensor_plane<fp8_e8m0_, block_factors<32, 1>>;
   using quant_tensor_t  = tensor_handle<fp8_e4m3_, scale_plane_t>; // FP8 weights
   using float_tensor_t  = tensor_handle<half>;

   [[kernel]]
   void quantized_matmul_kernel(
       quant_tensor_t  weights     [[tensor(0)]],
       float_tensor_t  activations [[tensor(1)]],
       float_tensor_t  output      [[tensor(2)]],
       uint2           tgid        [[threadgroup_position_in_grid]])
   {
       // Slice tiles for this threadgroup
       auto w_tile  = weights.slice(tgid);
       auto a_tile  = activations.slice(tgid);
       auto out_tile = output.slice(tgid);

       // TensorOps handles FP8+E8M0 dequantization automatically
       matmul2d_descriptor desc;
       desc.M_tile = 32; desc.N_tile = 32; desc.K_tile = 64;

       auto op = matmul2d<4>(desc); // 4 simdgroups per threadgroup
       op.run(w_tile, a_tile, out_tile);
   }
   ---------------------------------------------------------------- */