spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#include <Foundation/Foundation.hpp>3#include <Metal/Metal.hpp>4#include "core/device.h"5#include "core/tensor.h"6#include "ops/cpu/cpu_ops.h"7#include "ops/metal/metal_ops.h"8#include <cmath>9#include <cstdio>10#include <random>11#include <string>12using namespace forge;13struct Params { uint32_t M, N, K; };14int main() {15    NS::AutoreleasePool* pool = NS::AutoreleasePool::alloc()->init();16    const int64_t M = 128, N = 64, K = 96;   // exact multiples of the 64x32 tile17    Tensor a = Tensor::empty({M, K}), b = Tensor::empty({K, N});18    Tensor c = Tensor::zeros({M, N}), ref = Tensor::empty({M, N});19    std::mt19937 rng(5);20    std::uniform_real_distribution<float> d(-1.f, 1.f);21    for (int64_t i = 0; i < a.numel(); ++i) a.data<float>()[i] = d(rng);22    for (int64_t i = 0; i < b.numel(); ++i) b.data<float>()[i] = d(rng);23    cpu::matmul(a, b, ref, false, false);2425    Device& dev = Device::get();26    MTL::ComputePipelineState* pso = dev.pipeline("matmul_mpp_f32");27    Params p{uint32_t(M), uint32_t(N), uint32_t(K)};28    MTL::ComputeCommandEncoder* enc = metal::Stream::get().encoder();29    enc->setComputePipelineState(pso);30    enc->setBuffer(a.buffer(), a.buffer_offset(), 0);31    enc->setBuffer(b.buffer(), b.buffer_offset(), 1);32    enc->setBuffer(c.buffer(), c.buffer_offset(), 2);33    enc->setBytes(&p, sizeof(p), 3);34    enc->dispatchThreadgroups(MTL::Size(N/32, M/64, 1), MTL::Size(128,1,1));35    metal::sync();3637    float worst = 0.f; int64_t nz = 0;38    for (int64_t i = 0; i < c.numel(); ++i) {39        worst = std::max(worst, std::fabs(c.data<float>()[i] - ref.data<float>()[i]));40        if (c.data<float>()[i] != 0.f) ++nz;41    }42    std::printf("MPP matmul2d f32 vs CPU: max abs err %.3e  (nonzero outputs %lld/%lld)\n",43                double(worst), (long long)nz, (long long)c.numel());44    const bool ok32 = worst <= 1e-4f && nz == c.numel();45    std::printf("%s\n", ok32 ? "CORRECT" : "WRONG");4647    // f16 operands, f32 accumulator: the headline throughput number, so check48    // it computes the same product within f16 input precision.49    Tensor ah = Tensor::empty({M, K}, DType::F16);50    Tensor bh = Tensor::empty({K, N}, DType::F16);51    Tensor ch = Tensor::zeros({M, N});52    Tensor refh = Tensor::empty({M, N});53    for (int64_t i = 0; i < a.numel(); ++i) ah.set_item(i, a.data<float>()[i]);54    for (int64_t i = 0; i < b.numel(); ++i) bh.set_item(i, b.data<float>()[i]);55    // reference from the ROUNDED f16 values, so we compare arithmetic not rounding56    Tensor ar = Tensor::empty({M, K}), br = Tensor::empty({K, N});57    for (int64_t i = 0; i < a.numel(); ++i) ar.data<float>()[i] = ah.item_at(i);58    for (int64_t i = 0; i < b.numel(); ++i) br.data<float>()[i] = bh.item_at(i);59    cpu::matmul(ar, br, refh, false, false);6061    MTL::ComputePipelineState* pso16 = dev.pipeline("matmul_mpp_f16");62    MTL::ComputeCommandEncoder* e2 = metal::Stream::get().encoder();63    e2->setComputePipelineState(pso16);64    e2->setBuffer(ah.buffer(), ah.buffer_offset(), 0);65    e2->setBuffer(bh.buffer(), bh.buffer_offset(), 1);66    e2->setBuffer(ch.buffer(), ch.buffer_offset(), 2);67    e2->setBytes(&p, sizeof(p), 3);68    e2->dispatchThreadgroups(MTL::Size(N/32, M/64, 1), MTL::Size(128,1,1));69    metal::sync();7071    float w16 = 0.f; int64_t nz16 = 0;72    for (int64_t i = 0; i < ch.numel(); ++i) {73        w16 = std::max(w16, std::fabs(ch.data<float>()[i] - refh.data<float>()[i]));74        if (ch.data<float>()[i] != 0.f) ++nz16;75    }76    const bool ok16 = w16 <= 1e-3f && nz16 == ch.numel();77    std::printf("MPP matmul2d f16 vs CPU(f16 inputs): max abs err %.3e  "78                "(nonzero %lld/%lld)\n%s\n", double(w16), (long long)nz16,79                (long long)ch.numel(), ok16 ? "CORRECT" : "WRONG");80    // Transposed variants — training needs nt (forward X.W^T) and tn (dW = dY^T.X).81    bool okT = true;82    struct TCase { const char* suffix; bool ta, tb; const char* tag; };83    const TCase tcases[] = {{"_nt", false, true, "nt (fwd X.Wt)"},84                            {"_tn", true, false, "tn (dW = dYt.X)"}};85    for (const TCase& tc : tcases) {86        Tensor at = tc.ta ? Tensor::empty({K, M}) : Tensor::empty({M, K});87        Tensor bt = tc.tb ? Tensor::empty({N, K}) : Tensor::empty({K, N});88        Tensor ct = Tensor::zeros({M, N}), rt = Tensor::empty({M, N});89        for (int64_t i = 0; i < at.numel(); ++i) at.data<float>()[i] = d(rng);90        for (int64_t i = 0; i < bt.numel(); ++i) bt.data<float>()[i] = d(rng);91        cpu::matmul(at, bt, rt, tc.ta, tc.tb);9293        std::string kn = std::string("matmul_mpp_f32") + tc.suffix;94        MTL::ComputePipelineState* ps = dev.pipeline(kn);95        MTL::ComputeCommandEncoder* e3 = metal::Stream::get().encoder();96        e3->setComputePipelineState(ps);97        e3->setBuffer(at.buffer(), at.buffer_offset(), 0);98        e3->setBuffer(bt.buffer(), bt.buffer_offset(), 1);99        e3->setBuffer(ct.buffer(), ct.buffer_offset(), 2);100        e3->setBytes(&p, sizeof(p), 3);101        e3->dispatchThreadgroups(MTL::Size(N/32, M/64, 1), MTL::Size(128,1,1));102        metal::sync();103104        float w = 0.f; int64_t nzt = 0;105        for (int64_t i = 0; i < ct.numel(); ++i) {106            w = std::max(w, std::fabs(ct.data<float>()[i] - rt.data<float>()[i]));107            if (ct.data<float>()[i] != 0.f) ++nzt;108        }109        const bool ok = w <= 1e-4f && nzt == ct.numel();110        okT = okT && ok;111        std::printf("MPP matmul2d f32 %-16s max abs err %.3e  (nonzero %lld/%lld)  %s\n",112                    tc.tag, double(w), (long long)nzt, (long long)ct.numel(),113                    ok ? "CORRECT" : "WRONG");114    }115116    pool->drain();117    return (ok32 && ok16 && okT) ? 0 : 1;118}119