From c08b5eff9af676b5480acb15c23982c5666f0a4a Mon Sep 17 00:00:00 2001 From: ZHOU Jie Date: Mon, 28 Sep 2026 04:46:21 -0700 Subject: [PATCH 1/8] Metal: pure-C# macOS GPU backend (DET/CLS/REC on GPU, full-graph schedule) - ObjC/Metal P/Invoke layer (no native deps), MetalBackend probe + auto-select on macOS, MetalSession/MetalRunner/MetalGraph mirroring the Vulkan seam: compiled managed MetalSchedule (shape-keyed LRU, zero-cost per run) + grow-only shared device arena. - ~40 MSL compute kernels embedded as resources; GEMM family incl. simdgroup MMA 64x64 (mm_sg_dd direct-load, mm_ic_sg implicit-GEMM conv), dot-product conv kernels, SE squeeze-excite two-phase split (paravirt lacks acq_rel atomics). - OcrBackend.Metal (additive); session-create fallback to CPU permanent. - GpuBench PipelineBench: System.Drawing -> ImageSharp (cross-platform). - docs/metal-m4.md: 3-tier baselines (tiny 1.9x / small 4.7x / medium 5.2x vs sharp, 4w), correctness gates (lines identical, --conc 0/72, --dtr bitwise clean), tuning log and paravirt caps. --- docs/metal-m4.md | 73 + .../Backends/Metal/Metal.cs | 286 ++ .../Backends/Metal/MetalBackend.cs | 56 + .../Backends/Metal/MetalGraph.cs | 2438 +++++++++++++++++ .../Backends/Metal/MetalRunner.cs | 550 ++++ .../Backends/Metal/MetalSession.cs | 324 +++ src/Sdcb.SimdPaddleOCR/Backends/Metal/ObjC.cs | 229 ++ .../Backends/Metal/Shaders/a_common.metal | 42 + .../Backends/Metal/Shaders/conv.metal | 358 +++ .../Backends/Metal/Shaders/dot.metal | 334 +++ .../Backends/Metal/Shaders/elem.metal | 175 ++ .../Backends/Metal/Shaders/gemm.metal | 602 ++++ .../Backends/Metal/Shaders/space.metal | 350 +++ .../Backends/Metal/Shaders/xform.metal | 236 ++ src/Sdcb.SimdPaddleOCR/OcrBackend.cs | 2 + .../OnnxSharp/OcrSessionFactory.cs | 9 +- .../Sdcb.SimdPaddleOCR.csproj | 2 + .../PipelineBench.cs | 13 +- test/Sdcb.SimdPaddleOCR.Tests/IBenchEngine.cs | 3 +- test/Sdcb.SimdPaddleOCR.Tests/Program.cs | 6 +- 20 files changed, 6074 insertions(+), 14 deletions(-) create mode 100644 docs/metal-m4.md create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/Metal/Metal.cs create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalBackend.cs create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalGraph.cs create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalRunner.cs create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalSession.cs create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/Metal/ObjC.cs create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/a_common.metal create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/conv.metal create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/dot.metal create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/elem.metal create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/gemm.metal create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/space.metal create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/xform.metal diff --git a/docs/metal-m4.md b/docs/metal-m4.md new file mode 100644 index 0000000..429ac45 --- /dev/null +++ b/docs/metal-m4.md @@ -0,0 +1,73 @@ +# SimdPaddleOCR — Apple Silicon Metal 后端实测（Devin VM / M4 paravirt） + +实测机：macOS VM（Apple M4 paravirt GPU，8 vCPU），SDK .NET 10，`test/Sdcb.SimdPaddleOCR.Tests` bench（`--workers 4 --benchmark-kind simd --engine sharp|metal`，warmup=1，n=99，dataset/ 100 张固定图、张张不同尺寸）。commit：working tree @ feature/2.0 + metal-backend 工作区（未提交）。 + +**结论：三档 Metal 全部净胜 CPU（tiny 1.9× / small 4.7× / medium 5.2×），检测行数与 CPU 完全一致，文本差异仅 fp16 级噪声（3 图各 1 字符）。受 paravirt 虚拟化所限绝对数偏保守——实测 MMA≈fp32≈3TFLOPS、copy ~84GB/s、dispatch ~45µs，均低于真机 M4，真机数字预期更好。** + +## 端到端（4 workers，median ms/图，越低越好） + +| 模型 | CPU(sharp) | Metal | 比值 | CPU img/s | GPU img/s | +|---|---:|---:|---:|---:|---:| +| tiny | 73.9 | **38.9** | 1.90× | 13.5 | 25.7 | +| small | 269.7 | **57.1** | 4.72× | 3.7 | 17.5 | +| medium | 935.5 | **178.4** | 5.24× | 1.07 | 5.6 | + +注：与 Vulkan/B580 表口径相同（dataset 变 shape 是最不利 GPU 的场景）。Metal 走 `OcrBackend.Metal`（Auto 在 macOS+Metal 可用时自动选）。 + +## 分阶段耗时（mean ms/图） + +| 模型 | 阶段 | CPU | Metal | 倍率 | +|---|---|---:|---:|---:| +| tiny | det_graph | 35.5 | 12.7 | 2.8× | +| tiny | cls_graph | 11.3 | 3.25 | 3.5× | +| tiny | rec_graph | 99.7 | 15.9 | 6.3× | +| small | det_graph | 104.7 | 14.9 | 7.0× | +| small | cls_graph | 12.2 | 2.45 | 5.0× | +| small | rec_graph | 585.5 | 33.3 | 17.6× | +| medium | det_graph | 315.1 | 59.6 | 5.3× | +| medium | cls_graph | 11.9 | 2.61 | 4.6× | +| medium | rec_graph | 2322.2 | 112.6 | 20.6× | + +精度：metal/sharp 检出 lines 完全一致（1016/1009/1013）；逐图 hash 3 张差（fp16 阈值临界像素，与 Vulkan 同型已知限制）；文本 3 图各差 1 字符（CER 0.67% vs CPU 0.14%，在 fp16 噪声范围内）。`--conc`（4 线程 × 72 runs/模型）0 mismatch；`--dtr`（det medium 228 节点 × 8 跑）逐节点 bitwise 一致。 + +## 内存与分配 + +| 模型 | CPU peak WS | Metal peak WS | +|---|---:|---:| +| tiny | 691 MB | 574 MB | +| small | 857 MB | 647 MB | +| medium | 1691 MB | 948 MB | + +GPU 侧 arena 是共享 grow-only device buffer + 托管 schedule（MetalSchedule 纯数据 LRU=512，逐 shape 录制代价为零），与 Vulkan 的 plan/arena 形态同构。托管堆分配 ~1.4–2.5MB/图。 + +## GPU 能力实测（paravirt VM，--mprobe） + +| 项 | 实测 | +|---|---| +| fp32 FMA 峰值 | ~3.19 TFLOPS | +| fp16 MMA（half8x8×float-acc） | ~2.5–3.0 TFLOPS（≈fp32，疑似模拟实现） | +| copy 带宽 | ~84 GB/s | +| dispatch 摊销（单 CB 内） | ~40–48µs | +| commit+wait 往返 | ~250–300µs | +| threadgroup mem | 32KB；execWidth=32；maxTg=1024 | + +**含义：paravirt 上 fp16 MMA 无 tensor-core 加速**——kernel 收益主要来自减少 DRAM 流量和 dispatch/barrier 次数，而不是算力切换。 + +## M4 实测调优记录（保留的改动） + +| 优化 | 机制 | 实测 | +|---|---|---| +| `mm_sg_dd` | 64×64 MMA tile 直接从 device 内存 simdgroup_load，省掉 threadgroup staging 往返；emit gate `K%16==0`（边界 tile 的越界读被 arena ≥1MB 尾 slack + W /128 pad + epilogue 掩码兜住，数学不变） | 5760×512×1024 GEMM 2.63→2.31ms；det tiny 8.6→6.8ms、small 20.8→17.7ms | +| `mm_ic_sg`（隐式 GEMM 卷积） | kxk conv 不再物化 `[M,Kp]` im2col 矩阵：A staging 内联 tap 寻址 gather；同时接管 nb==1 的 convk_dot 分支 | det medium 960² 86.2→**80.0ms**（-7%）；消掉 ~1.28GB 无谓读写 | +| SE 两阶段定稿 | `se_part`（S WG 写 partial，零原子）+ buffer barrier + `se_join`（1 WG/batch）| 修掉跨 WG ticket 竞态（paravirt MSL 无 acq_rel 原子，握手不可能）；`--dtr` 全绿 | +| conv_dw4a flat | 平铺 depthwise 省 threadgroup staging | **未采纳**：实测比 conv_dw4t 慢（staging 省 halo 重读仍有收益）→ 留 `SIMD_OCR_DWA` 开关 | +| `mm_ic_sg32`（BK=32） | staging barrier 减半 | **未采纳**：medium −1ms 但 tiny/small 更差 → 留 `SIMD_OCR_IC32` 开关 | +| NODOT（conv1x1 全 MMA） | 关 dot kernel | 与现状持平 → 阈值保持 B580 标定值 | +| rec bias-free SE 融合 | rec 的 fc1/fc2 conv 无 bias（Inputs[2]==MaxValue），IsPwConv 拒绝 | **跳过**（Vulkan 同缺），记为遗留机会 | + +## 已知限制 / 未验证项 + +- DET bs>1 resize 不支持（同 Vulkan）。 +- paravirt 独有：MSL `memory_order` 只有 relaxed → 跨 WG 定序走两阶段 split + `memoryBarrierWithScope:Buffers`；`char16/uchar16` 不可用。 +- 序列化 `--prof` 有 ~0.4–0.5ms/dispatch 地板，绝对值偏大；真实比例看 `--time` 墙钟与 `--mgemm` N-reps-in-one-CB。 +- 无 CI Metal 依赖；单元测试在非 macOS 上 SKIP。 diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Metal.cs b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Metal.cs new file mode 100644 index 0000000..f444c0a --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Metal.cs @@ -0,0 +1,286 @@ +using System.Runtime.InteropServices; + +namespace Sdcb.SimdPaddleOCR.Backends.Metal; + +// Thin Metal wrappers over ObjC msgSend. +1-owned objects are released on +// Dispose; autoreleased objects (commandBuffer, computeCommandEncoder, …) are +// only valid inside the current NSAutoreleasePool — all public entry points +// that create them take their own pool. + +// MTLResourceOptions: StorageModeShared = 0 << 4 = 0 (the zero value), which +// is what the whole backend uses — unified memory means the same allocation +// is CPU- and GPU-visible without didModifyRange/coherence calls. +internal static class MtlStorageMode { public const nuint Shared = 0; } + +// MTLDataType values (MTLDataType.h) +internal static class MtlDataType +{ + public const nuint Float = 3; + public const nuint Int = 29; + public const nuint UInt = 30; + public const nuint Bool = 53; +} + +internal sealed class MtlDevice : IDisposable +{ + public readonly IntPtr H; + private readonly IntPtr _queue; + + private MtlDevice(IntPtr h) + { + H = h; + // newCommandQueue is +1 — owned, released in Dispose. + _queue = ObjC.Send0(h, ObjC.S.NewCommandQueue); + if (_queue == 0) throw new InvalidOperationException("Metal: newCommandQueue returned nil"); + } + + public IntPtr Queue => _queue; + + public string Name => ObjC.NsToString(ObjC.Send0(H, ObjC.S.Name)) ?? "?"; + + // MTLGPUFamily enum values (MTLGPUFamily.h) + public bool SupportsFamily(int family) => ObjC.SendBool1L(H, ObjC.S.SupportsFamily, family) != 0; + public bool SupportsApple1 => SupportsFamily(1001); + public bool SupportsApple2 => SupportsFamily(1002); + public bool SupportsApple3 => SupportsFamily(1003); + public bool SupportsApple4 => SupportsFamily(1004); + public bool SupportsApple5 => SupportsFamily(1005); + public bool SupportsApple6 => SupportsFamily(1006); + public bool SupportsApple7 => SupportsFamily(1007); + public bool SupportsApple8 => SupportsFamily(1008); + public bool SupportsApple9 => SupportsFamily(1009); + public bool SupportsMetal3 => SupportsFamily(5001); + public bool SupportsMetal4 => SupportsFamily(5002); + + public nuint MaxThreadgroupMemoryLength => ObjC.SendNuint(H, ObjC.S.MaxThreadgroupMemoryLength); + + /// Defensive check for optional selectors — an unrecognized + /// selector on this paravirt driver raises an ObjC exception that kills + /// the process, so never call a maybe-absent method unconditionally. + public bool RespondsTo(IntPtr sel) => ObjC.SendBool1P(H, ObjC.S.RespondsToSelector, sel) != 0; + public bool HasUnifiedMemory => ObjC.SendBool(H, ObjC.S.HasUnifiedMemory) != 0; + public ulong RecommendedMaxWorkingSetSize => (ulong)ObjC.SendNuint(H, ObjC.S.RecommendedMaxWorkingSetSize); + public nuint CurrentAllocatedSize => ObjC.SendNuint(H, ObjC.S.CurrentAllocatedSize); + public ulong RegistryID => (ulong)ObjC.SendNuint(H, ObjC.S.RegistryID); + + public static MtlDevice? Probe() + { + if (!OperatingSystem.IsMacOS()) return null; + using var pool = AutoReleasePool.Create(); + IntPtr h = ObjC.CreateSystemDefaultDevice(); + if (h == 0) return null; + ObjC.Retain(h); // C-function result is autoreleased; promote to +1. + return new MtlDevice(h); + } + + public MtlBuffer NewBuffer(nuint bytes) + { + // MTLResourceStorageModeShared = 0 options. + IntPtr b = ObjC.Send1N1N(H, ObjC.S.NewBufferWithLengthOptions, bytes, 0); + if (b == 0) throw new InvalidOperationException($"Metal: newBufferWithLength({bytes}) failed"); + return new MtlBuffer(b, bytes); + } + + public unsafe MtlBuffer NewBuffer(void* src, nuint bytes) + { + IntPtr b = ObjC.Send1P2N(H, ObjC.S.NewBufferWithBytesLengthOptions, (IntPtr)src, bytes, 0); + if (b == 0) throw new InvalidOperationException($"Metal: newBufferWithBytes({bytes}) failed"); + return new MtlBuffer(b, bytes); + } + + /// Compile MSL source into a library. Throws with the NSError + /// localizedDescription on failure — MSL compile errors must never be + /// silently swallowed. + public unsafe MtlLibrary NewLibrary(string source, bool fastMath = false) + { + using var pool = AutoReleasePool.Create(); + IntPtr src = ObjC.NsStr(source); + IntPtr opts = MakeCompileOptions(fastMath); + IntPtr err = 0; + IntPtr lib; + lib = ObjC.Send3P(H, ObjC.S.NewLibraryWithSourceOptionsError, src, opts, (IntPtr)(&err)); + ObjC.Release(src); + ObjC.Release(opts); + if (lib == 0) throw new InvalidOperationException("MSL compile failed: " + DescribeError(err)); + return new MtlLibrary(lib); + } + + private static IntPtr MakeCompileOptions(bool fastMath) + { + // MTLCompileOptions: mathMode is a private-ish key; the public knob is + // fastMathEnabled (deprecated in MSL 3.x but still honored) plus + // languageVersion. MSL 3.1 = 0x30001 << 16? No — MTLLanguageVersion is + // (major<<16)+minor: 2.4=0x20004, 3.1=0x30001. We request 3.1 when + // supported for simdgroup_matrix; harmless if clamped. + IntPtr o = ObjC.Send0(ObjC.Send0(ObjC.ClassMTLCompileOptions, ObjC.S.Alloc), ObjC.S.Init); + ObjC.SendV1N(o, ObjC.S.SetLanguageVersion, (nuint)0x00030001); // MSL 3.1 + ObjC.SendV1B(o, ObjC.S.SetFastMathEnabled, (byte)(fastMath ? 1 : 0)); + return o; + } + + public unsafe MtlPipeline NewPipeline(MtlLibrary lib, string fnName, int[]? intConstants = null) + { + using var pool = AutoReleasePool.Create(); + IntPtr name = ObjC.NsStr(fnName); + IntPtr fn; + if (intConstants is { Length: > 0 }) + { + IntPtr fc = ObjC.Send0(ObjC.Send0(ObjC.ClassMTLFunctionConstantValues, ObjC.S.Alloc), ObjC.S.Init); + for (int i = 0; i < intConstants.Length; ++i) + { + int v = intConstants[i]; + // setConstantValue:type:atIndex: — (const void*, MTLDataType, NSUInteger). + ObjC.Send3P(fc, ObjC.S.SetConstantValueTypeAtIndex, (IntPtr)(&v), (IntPtr)(nint)(long)MtlDataType.Int, (IntPtr)(nint)i); + } + IntPtr err = 0; + fn = ObjC.Send3P(lib.H, ObjC.S.NewFunctionWithNameConstantValuesError, name, fc, (IntPtr)(&err)); + ObjC.Release(fc); + if (fn == 0) { ObjC.Release(name); throw new InvalidOperationException($"Metal: function '{fnName}' failed: {DescribeError(err)}"); } + } + else + { + fn = ObjC.Send1P(lib.H, ObjC.S.NewFunctionWithName, name); + } + ObjC.Release(name); + if (fn == 0) throw new InvalidOperationException($"Metal: no function '{fnName}' in library"); + + IntPtr err2 = 0; + IntPtr pso = ObjC.Send2P(H, ObjC.S.NewComputePipelineStateWithFunctionError, fn, (IntPtr)(&err2)); + ObjC.Release(fn); + if (pso == 0) throw new InvalidOperationException($"Metal: PSO '{fnName}' failed: {DescribeError(err2)}"); + return new MtlPipeline(pso, fnName); + } + + internal static unsafe string DescribeError(IntPtr err) + { + if (err == 0) return "(no NSError)"; + IntPtr desc = ObjC.Send0(err, ObjC.S.LocalizedDescription); + return ObjC.NsToString(desc) ?? "(unreadable NSError)"; + } + + /// Begin a command buffer + compute encoder pair. Both are + /// autoreleased — caller must hold an AutoReleasePool open. + public MtlEncoder BeginCommands() + { + IntPtr cmd = ObjC.Send0(_queue, ObjC.S.CommandBuffer); + IntPtr enc = ObjC.Send0(cmd, ObjC.S.ComputeCommandEncoder); + return new MtlEncoder(cmd, enc); + } + + /// Begin encoder with a specific dispatch type + /// (MTLDispatchType: 0=Serial, 1=Concurrent — concurrent allows the GPU to + /// overlap dispatches that don't depend on each other; serial inserts an + /// implicit barrier between dispatches which is what the schedule assumes). + public MtlEncoder BeginCommands(bool concurrent) + { + IntPtr cmd = ObjC.Send0(_queue, ObjC.S.CommandBuffer); + IntPtr enc = concurrent + ? ObjC.Send1N(cmd, ObjC.S.ComputeCommandEncoderWithDispatchType, 1) + : ObjC.Send0(cmd, ObjC.S.ComputeCommandEncoder); + return new MtlEncoder(cmd, enc); + } + + public void Dispose() + { + ObjC.Release(_queue); + ObjC.Release(H); + } +} + +internal sealed class MtlBuffer : IDisposable +{ + public IntPtr H; + public readonly nuint Length; + public MtlBuffer(IntPtr h, nuint length) { H = h; Length = length; } + public IntPtr Contents => ObjC.Send0(H, ObjC.S.Contents); + public void Dispose() { ObjC.Release(H); H = 0; } +} + +internal sealed class MtlLibrary : IDisposable +{ + public readonly IntPtr H; + public MtlLibrary(IntPtr h) => H = h; + public void Dispose() => ObjC.Release(H); +} + +internal sealed class MtlPipeline : IDisposable +{ + public readonly IntPtr Pso; + public readonly string Name; + public readonly nuint MaxTotalThreadsPerThreadgroup; + public readonly nuint ThreadExecutionWidth; + public MtlPipeline(IntPtr pso, string name) + { + Pso = pso; + Name = name; + MaxTotalThreadsPerThreadgroup = ObjC.SendNuint(pso, ObjC.S.MaxTotalThreadsPerThreadgroup); + ThreadExecutionWidth = ObjC.SendNuint(pso, ObjC.S.ThreadExecutionWidth); + } + public void Dispose() => ObjC.Release(Pso); +} + +/// One-shot command-buffer + compute-encoder pair. The objects are +/// autoreleased; keep the enclosing AutoReleasePool alive until Commit has +/// been issued (and waitUntilCompleted returned if you read results). +internal readonly struct MtlEncoder +{ + public readonly IntPtr Cmd; + public readonly IntPtr Enc; + + public MtlEncoder(IntPtr cmd, IntPtr enc) { Cmd = cmd; Enc = enc; } + + public void SetPipeline(MtlPipeline p) => ObjC.Send1P(Enc, ObjC.S.SetComputePipelineState, p.Pso); + + public void SetBuffer(MtlBuffer buf, nuint byteOffset, int index) => + ObjC.SendV1P2N(Enc, ObjC.S.SetBufferOffsetAtIndex, buf.H, byteOffset, (nuint)index); + + public unsafe void SetPc(in T value, int index) where T : unmanaged + { + fixed (T* p = &value) + ObjC.SendV1P2N(Enc, ObjC.S.SetBytesLengthAtIndex, (IntPtr)p, (nuint)sizeof(T), (nuint)index); + } + + public unsafe void SetPc(void* data, nuint bytes, int index) => + ObjC.SendV1P2N(Enc, ObjC.S.SetBytesLengthAtIndex, (IntPtr)data, bytes, (nuint)index); + + public void SetThreadgroupMemory(nuint bytes, int index) => + ObjC.SendV1N1N(Enc, ObjC.S.SetThreadgroupMemoryLengthAtIndex, bytes, (nuint)index); + + public void Dispatch(nuint gx, nuint gy, nuint gz, nuint tx, nuint ty, nuint tz) => + ObjC.SendV2Size(Enc, ObjC.S.DispatchThreadgroupsThreadsPerThreadgroup, + new ObjC.MTLSize(gx, gy, gz), new ObjC.MTLSize(tx, ty, tz)); + + /// MTLBarrierScope buffers(=1): visibility barrier between + /// dispatches — needed when a later dispatch consumes a buffer written by + /// an earlier one and we used a Concurrent encoder. Serial encoders imply + /// this already. + public void BarrierBuffers() => ObjC.SendV1N(Enc, ObjC.S.MemoryBarrierWithScope, 1); + + public void End() => ObjC.Send0(Enc, ObjC.S.EndEncoding); + + public void Commit() => ObjC.Send0(Cmd, ObjC.S.Commit); + + /// Wait for completion and check status; throws with NSError on + /// any non-completed terminal state (incl. GPU watchdog kills). + public void CommitAndWait() + { + Commit(); + ObjC.Send0(Cmd, ObjC.S.WaitUntilCompleted); + long status = (long)ObjC.SendLong(Cmd, ObjC.S.Status); + // MTLCommandBufferStatus: 4=Completed; 5=Error; others shouldn't occur post-wait. + if (status != 4) + { + IntPtr err = ObjC.Send0(Cmd, ObjC.S.Error); + throw new InvalidOperationException($"Metal command buffer failed (status={status}): {MtlDevice.DescribeError(err)}"); + } + } + + /// GPU-side wall time of the buffer (seconds). 0 when timestamps + /// are unavailable (paravirt drivers often return 0). + public double GpuDurationSeconds() + { + double s = ObjC.SendDouble(Cmd, ObjC.S.GpuStartTime); + double e = ObjC.SendDouble(Cmd, ObjC.S.GpuEndTime); + return e > s ? e - s : 0; + } +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalBackend.cs b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalBackend.cs new file mode 100644 index 0000000..050d75e --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalBackend.cs @@ -0,0 +1,56 @@ +using Sdcb.SimdPaddleOCR.OnnxSharp; + +namespace Sdcb.SimdPaddleOCR.Backends.Metal; + +/// +/// Process-wide Metal device for OCR graphs plus the backend-resolution rules +/// (option → SIMD_OCR_BACKEND env → probe). On Apple platforms Metal wins the +/// Auto slot ahead of Vulkan; on other platforms Metal is never selected. +/// +internal static class MetalBackend +{ + private static readonly object s_probeLock = new(); + private static MtlDevice? s_device; + private static bool s_probed; + + /// The shared device, or null when Metal is unavailable (not + /// macOS, no GPU, or the ObjC bridge can't be reached). + internal static MtlDevice? TryGetDevice() + { + if (s_probed) return s_device; + lock (s_probeLock) + { + if (s_probed) return s_device; + try { s_device = OperatingSystem.IsMacOS() ? MtlDevice.Probe() : null; } + catch { s_device = null; } + s_probed = true; + return s_device; + } + } + + /// Whether the given option resolves to the Metal path. + internal static bool IsMetalSelected(OcrBackend backend) => backend switch + { + OcrBackend.Metal => true, + OcrBackend.Cpu or OcrBackend.Vulkan => false, + _ => Environment.GetEnvironmentVariable("SIMD_OCR_BACKEND") switch + { + { } s when s.Equals("metal", StringComparison.OrdinalIgnoreCase) => true, + { } s when s.Equals("cpu", StringComparison.OrdinalIgnoreCase) + || s.Equals("vulkan", StringComparison.OrdinalIgnoreCase) => false, + // Auto: Metal is the GPU of choice on macOS only + _ => OperatingSystem.IsMacOS(), + }, + }; + + /// Creates a session on Metal; CPU on any failure. + internal static IOcrSession CreateSession(CompiledModel compiled, OcrBackend backend) + { + if (IsMetalSelected(backend) && TryGetDevice() is { } dev) + { + try { return new MetalSession(dev, compiled); } + catch { /* fall through to CPU */ } + } + return compiled.CreateRequest(); + } +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalGraph.cs b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalGraph.cs new file mode 100644 index 0000000..dae67a8 --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalGraph.cs @@ -0,0 +1,2438 @@ +using System.Buffers.Binary; +using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; +using Sdcb.SimdPaddleOCR.OnnxSharp; + +namespace Sdcb.SimdPaddleOCR.Backends.Metal; + +/// +/// +/// Metal counterpart of the Metal MetalGraphModel: shared per (device, +/// compiled model); owns the embedded-MSL library pipelines, fp16 weight +/// uploads and the graph compiler that emits a +/// — pure managed metadata (pipeline, arena-relative binding offsets, push +/// constants, grid), no Metal object per shape. Sessions +/// () own the grow-only arena/IO buffers and +/// re-record a command buffer each run. +/// +internal sealed class MetalGraphModel +{ + // Binding sentinels: schedules reference session-owned buffers by role; + // the session substitutes its live handle when recording. + internal static readonly MtlBuffer RoleArena = new(0, 0), RoleIn = new(0, 0), + RoleOut = new(0, 0), RolePart = new(0, 0); + + private static readonly ConditionalWeakTable s_models = new(); + private int _refs; + + /// Ref-counted shared model for (device, compiled); released by + /// when the last session goes away. + internal static MetalGraphModel Acquire(MtlDevice dev, CompiledModel compiled) + { + lock (s_models) + { + if (!s_models.TryGetValue(compiled, out MetalGraphModel? m) || m._refs == 0 || m._dev != dev) + { + m = new MetalGraphModel(dev, compiled); + s_models.AddOrUpdate(compiled, m); + } + m._refs++; + return m; + } + } + + internal void Release() + { + lock (s_models) + { + if (--_refs > 0) return; + if (s_models.TryGetValue(_compiled, out MetalGraphModel? cur) && cur == this) + s_models.Remove(_compiled); + lock (this) + { + foreach (MtlBuffer b in _allBufs) b.Dispose(); + _allBufs.Clear(); + _schedules.Clear(); + } + } + } + + private readonly MtlDevice _dev; + private readonly Model _model; + private readonly CompiledModel _compiled; + internal MtlDevice Device => _dev; + internal MtlPipeline OutPipe => _pOut; + + private readonly MtlPipeline _pSoftmax; + private readonly MtlPipeline _pDot, + _pDw, _pDw4, _pDw4T, _pDense, _pConvT, _pConvT4, _pElem, _pElem4, + _pReduce, _pReduce4, _pReduce4b, _pPool, _pPool4, + _pResize, _pResize4, _pResize4Add, _pConcat, _pConcat4, _pNchw, _pOut, _pIm2col, + _pAddPs, _pSePart, _pSeJoin, _pConvD, _pConvDF32, _pCatRes, + _pAvg4, _pAffine, _pLn, _pAttn, _pDotSk; + // Flat (unstaged) depthwise — threadgroup staging buys nothing on this + // unified-memory GPU; SIMD_OCR_NODWA=1 keeps the tiled variant for A/B. + private readonly MtlPipeline? _pDw4A; + + // fp16 copies of constant tensors (weights, biases, scalars), lazy + private readonly Dictionary<(int, int, int, int), MtlBuffer> _constF16 = new(); + private readonly Dictionary _constF32 = new(); + private readonly Dictionary _constTap = new(); + + // conv weight repack to tap-major [Cout, (dy*kW+dx)*Cin+ci], rows padded to Kp + private MtlBuffer ConstTapMajor(int tensorIndex, int cout, int cin, + int kH, int kW, int kp, int cinPad = 0) + { + if (cinPad == 0) cinPad = cin; + if (_constTap.TryGetValue(tensorIndex, out MtlBuffer? cached)) return cached; + ReadOnlySpan f32 = + MemoryMarshal.Cast(_compiled.GetTensor(tensorIndex).Constant); + int kSpan = kH * kW; + int rowsPad = (cout + 127) / 128 * 128; + Half[] h = new Half[rowsPad * kp]; + for (int co = 0; co < cout; co++) + for (int ci = 0; ci < cinPad; ci++) + for (int tap = 0; tap < kSpan; tap++) + if (ci < cin) + h[co * kp + tap * cinPad + ci] = + (Half)f32[co * cin * kSpan + ci * kSpan + tap]; + MtlBuffer buf = Up16(h); + _constTap[tensorIndex] = buf; + _allBufs.Add(buf); + return buf; + } + + // depthwise weight repack to tap-major [tap*C + c] fp16 (vec4-readable) + private readonly Dictionary _constDw = new(); + private MtlBuffer ConstDwTap(int tensorIndex, int c, int kSpan) + { + if (_constDw.TryGetValue(tensorIndex, out MtlBuffer? cached)) return cached; + ReadOnlySpan f32 = + MemoryMarshal.Cast(_compiled.GetTensor(tensorIndex).Constant); + Half[] h = new Half[c * kSpan]; + for (int cc = 0; cc < c; cc++) + for (int tap = 0; tap < kSpan; tap++) + h[tap * c + cc] = (Half)f32[cc * kSpan + tap]; + MtlBuffer buf = Up16(h); + _constDw[tensorIndex] = buf; + _allBufs.Add(buf); + return buf; + } + + // k-major weight repack for conv1x1_dot: [k*Cout + n], k = tap*cinPad + ci + private readonly Dictionary _constKm = new(); + internal unsafe float[]? DbgConstKm(int tensorIndex, int n) + { + if (!_constKm.TryGetValue(tensorIndex, out MtlBuffer? b)) return null; + var sp = new ReadOnlySpan((void*)b.Contents, n); + var f = new float[n]; + for (int i = 0; i < n; i++) f[i] = (float)BitConverter.UInt16BitsToHalf(sp[i]); + return f; + } + private MtlBuffer ConstKMajor(int tensorIndex, int cout, int cin, + int kSpan, int kp, int cinPad) + { + if (_constKm.TryGetValue(tensorIndex, out MtlBuffer? cached)) return cached; + ReadOnlySpan f32 = + MemoryMarshal.Cast(_compiled.GetTensor(tensorIndex).Constant); + Half[] h = new Half[kp * cout]; + for (int co = 0; co < cout; co++) + for (int ci = 0; ci < cinPad; ci++) + for (int tap = 0; tap < kSpan; tap++) + if (ci < cin) + h[(tap * cinPad + ci) * cout + co] = + (Half)f32[co * cin * kSpan + ci * kSpan + tap]; + MtlBuffer buf = Up16(h); + _constKm[tensorIndex] = buf; + _allBufs.Add(buf); + return buf; + } + + // ConvTranspose weight repack [tap*Cin + ci, co] (tap=dy*2+dx) fp16 + private readonly Dictionary _constCt = new(); + private MtlBuffer ConstConvT(int tensorIndex, int cin, int cout) + { + if (_constCt.TryGetValue(tensorIndex, out MtlBuffer? cached)) return cached; + ReadOnlySpan f32 = + MemoryMarshal.Cast(_compiled.GetTensor(tensorIndex).Constant); + Half[] h = new Half[4 * cin * cout]; + for (int tap = 0; tap < 4; tap++) + for (int ci = 0; ci < cin; ci++) + for (int co = 0; co < cout; co++) + h[(tap * cin + ci) * cout + co] = + (Half)f32[(ci * cout + co) * 4 + tap]; + MtlBuffer buf = Up16(h); + _constCt[tensorIndex] = buf; + _allBufs.Add(buf); + return buf; + } + + // reduce_hw4 / se_part phase-1 partials (fp32): a fixed + // session-owned scratch (S partitions × C channels × batch) + // partials: fp32 [0, PartCtrElem) + internal const int PartCtrElem = 1 << 18; + internal const int PartBytes = PartCtrElem * 4; + private static MtlBuffer PartBuf(int floats) => RolePart; + private readonly List _allBufs = new(); + + private MetalGraphModel(MtlDevice dev, CompiledModel compiled) + { + _dev = dev; + _compiled = compiled; + _model = compiled.Model; + _lib = LoadLibrary(); + + // one probe per capability tier: the simdgroup-matrix kernels get a + // PSO + smoke dispatch; failure downgrades the cm paths to the + // portable dot kernels (correct, just slower). + _sgm = ProbeSimdMatrix(); + _pMm = Pipe("mm_sg"); + _pMmPs = Pipe("mm_ps_sg"); + _pMmDot = Pipe("mm_dot"); // portable cm-path fallback + _pMmDd = Pipe("mm_sg_dd"); // direct fragment loads, no staging + _pMmIc = Pipe("mm_ic_sg"); // implicit-GEMM kxk conv (tap gather in A stage) + _pMmIc32 = Environment.GetEnvironmentVariable("SIMD_OCR_IC32") != null + ? Pipe("mm_ic_sg32") : _pMmIc; // measured: BK=32 loses except det-medium + if (!_sgm) { _pMm = _pMmDot; _pMmPs = _pMmDot; _pMmDd = _pMmDot; _pMmIc = _pMmDot; _pMmIc32 = _pMmDot; } + + _pDot = Pipe("conv1x1_dot"); + _pDw = Pipe("conv_dw"); + _pDw4 = Pipe("conv_dw4"); + _pDw4T = Pipe("conv_dw4t"); + _pDw4A = Environment.GetEnvironmentVariable("SIMD_OCR_DWA") != null + ? Pipe("conv_dw4a") : null; // measured: flat loses to conv_dw4t + _pDense = Pipe("conv_dense"); + _pIm2col = Pipe("im2col"); + _pConvT = Pipe("convt2s2"); + _pConvT4 = Pipe("convt4"); + _pReduce4 = Pipe("reduce_hw4"); + _pReduce4b = Pipe("reduce_hw4b"); + _pPool4 = Pipe("maxpool4"); + _pResize4 = Pipe("resize4"); + _pResize4Add = Pipe("resize4add"); + _pConcat4 = Pipe("concat4"); + _pCatRes = Pipe("catresize"); + _pElem = Pipe("elem"); + _pElem4 = Pipe("elem4"); + _pAddPs = Pipe("addps"); + _pSePart = Pipe("se_part"); + _pSeJoin = Pipe("se_join"); + _pConvD = Pipe("convk_dot"); + _pConvDF32 = Pipe("convk_f32n"); + _pReduce = Pipe("reduce_hw"); + _pPool = Pipe("maxpool2e"); + _pResize = Pipe("resize_nn"); + _pConcat = Pipe("concat_c"); + _pNchw = Pipe("nchw2nhwc"); + _pOut = Pipe("sigmoid_out"); + _pAvg4 = Pipe("avgpool4"); + _pAffine = Pipe("affine4"); + _pSoftmax = Pipe("softmax"); + _pLn = Pipe("layernorm"); + _pAttn = Pipe("attn"); + _pDotSk = Pipe("mmdot_sk"); + } + + private readonly MtlLibrary _lib; + private readonly bool _sgm; // simdgroup_matrix 8x8 fp16 MMA verified live + // sg32-family fields kept for emit parity with the Vulkan schedule + // (DwTiled uses the sg32 kernel cap; PackSlabs runs unconditionally here) + private readonly bool _sg32 = true; + private readonly bool _lite = false; + private readonly MtlPipeline _pMm, _pMmPs, _pMmDot, _pMmDd, _pMmIc, _pMmIc32; + + // Concatenated .metal sources embedded as manifest resources, sorted so + // a_common.metal's helpers come first. + private MtlLibrary LoadLibrary() + { + var asm = typeof(MetalGraphModel).Assembly; + var names = asm.GetManifestResourceNames() + .Where(n => n.StartsWith("Sdcb.SimdPaddleOCR.Backends.Metal.Shaders.") + && n.EndsWith(".metal")) + .OrderBy(n => n).ToArray(); + if (names.Length == 0) + throw new FileNotFoundException("Metal: no embedded .metal shaders"); + var sb = new System.Text.StringBuilder(); + foreach (var n in names) + { + using Stream s = asm.GetManifestResourceStream(n)!; + using var r = new StreamReader(s); + sb.AppendLine(r.ReadToEnd()); + } + return _dev.NewLibrary(sb.ToString()); + } + + /// PSO + live correctness smoke for simdgroup_matrix kernels: + /// a PSO that links is not proof the driver executes it — the probe runs + /// a 64x64 tile GEMM against a canned input and compares one element. + /// + private unsafe bool ProbeSimdMatrix() + { + try + { + MtlPipeline p = _dev.NewPipeline(_lib, "mm_sg"); + const int M = 64, N = 64, K = 16; + var a = new Half[M * K]; var w = new Half[N * K]; + var rnd = new Random(1234); + for (int i = 0; i < a.Length; i++) a[i] = (Half)((rnd.NextDouble() - 0.5) * 0.25); + for (int i = 0; i < w.Length; i++) w[i] = (Half)((rnd.NextDouble() - 0.5) * 0.25); + using MtlBuffer dum = _dev.NewBuffer(256); + using MtlBuffer ob = _dev.NewBuffer((nuint)(M * N * 2)); + uint[] pc = [(uint)M, (uint)N, (uint)K, (uint)(K / 4), 0u, 0u, 0u]; + fixed (Half* pa = a) + fixed (Half* pw = w) + { + using MtlBuffer xa = _dev.NewBuffer(pa, (nuint)(a.Length * 2)); + using MtlBuffer wb = _dev.NewBuffer(pw, (nuint)(w.Length * 2)); + using (var pool = AutoReleasePool.Create()) + { + MtlEncoder e = _dev.BeginCommands(); + e.SetPipeline(p); + e.SetBuffer(xa, 0, 0); e.SetBuffer(wb, 0, 1); e.SetBuffer(dum, 0, 2); + e.SetBuffer(dum, 0, 3); e.SetBuffer(ob, 0, 4); e.SetBuffer(dum, 0, 5); + e.SetBuffer(dum, 0, 6); + fixed (uint* pp = pc) e.SetPc(pp, (nuint)(pc.Length * 4), 7); + e.Dispatch(1, 1, 1, 256, 1, 1); + e.End(); e.CommitAndWait(); + } + } + // verify a corner element: o[0,0] = sum_k a[0,k]*w[0,k] + float expect = 0; + for (int k = 0; k < K; k++) expect += (float)a[k] * (float)w[k]; + float got = (float)*(Half*)ob.Contents; + return MathF.Abs(got - expect) <= 0.05f * MathF.Max(1, MathF.Abs(expect)); + } + catch { return false; } + } + + private readonly Dictionary _pipes = new(); + private MtlPipeline Pipe(string name) + { + if (!_pipes.TryGetValue(name, out MtlPipeline? p)) + { + p = _dev.NewPipeline(_lib, name); + _pipes[name] = p; + } + return p; + } + + // constant upload helpers (Shared-mode buffers: contents is a plain + // pointer on unified memory — the repack writes straight in) + private unsafe MtlBuffer Up16(Half[] h) + { + fixed (Half* p = h) return _dev.NewBuffer(p, (nuint)(h.Length * 2)); + } + private unsafe MtlBuffer Up32(ReadOnlySpan f) + { + fixed (float* p = f) return _dev.NewBuffer(p, (nuint)(f.Length * 4)); + } + + + // fp16 device copy of a constant tensor; padded to padRows*rowElemsPad + // elements. rowPad>rowElems repacks [rows, rowElems] into [rows, rowPad]. + private MtlBuffer ConstF16(int tensorIndex, int padRows = 0, int rowElems = 0, + int rowPad = 0) + { + var key = (tensorIndex, padRows, rowElems, rowPad); + if (_constF16.TryGetValue(key, out MtlBuffer? cached)) return cached; + ReadOnlySpan f32 = + MemoryMarshal.Cast(_compiled.GetTensor(tensorIndex).Constant); + int n = f32.Length; + int alloc; + Half[] h; + if (rowPad > rowElems && rowElems > 0) + { + int rows = Math.Max(n / rowElems, padRows); + alloc = rows * rowPad; + h = new Half[alloc]; + for (int r = 0; r * rowElems < n; r++) + for (int c = 0; c < rowElems && r * rowElems + c < n; c++) + h[r * rowPad + c] = (Half)f32[r * rowElems + c]; + } + else + { + alloc = Math.Max(n, padRows * rowElems); + h = new Half[alloc]; + for (int i = 0; i < n; i++) h[i] = (Half)f32[i]; + } + MtlBuffer buf = Up16(h); + _constF16[key] = buf; + _allBufs.Add(buf); + return buf; + } + + // fp32 device copy of a constant tensor (small tensors needing exact math). + private MtlBuffer ConstF32(int tensorIndex) + { + if (_constF32.TryGetValue(tensorIndex, out MtlBuffer? cached)) return cached; + ReadOnlySpan f32 = + MemoryMarshal.Cast(_compiled.GetTensor(tensorIndex).Constant); + MtlBuffer buf = Up32(f32); + _constF32[tensorIndex] = buf; + _allBufs.Add(buf); + return buf; + } + + private ReadOnlySpan CstF32(int tensorIndex) => + MemoryMarshal.Cast(_compiled.GetTensor(tensorIndex).Constant); + + // fp16 device copy of a computed vector (e.g. folded BN affine params) + // keyed by (node, slot) so re-compiling another shape reuses the upload + private readonly Dictionary<(int, int), MtlBuffer> _vecF16 = new(); + private MtlBuffer VecF16((int, int) key, float[] v) + { + if (_vecF16.TryGetValue(key, out MtlBuffer? cached)) return cached; + Half[] h = new Half[v.Length]; + for (int i = 0; i < v.Length; i++) h[i] = (Half)v[i]; + MtlBuffer buf = Up16(h); + _vecF16[key] = buf; + _allBufs.Add(buf); + return buf; + } + + // MatMul weight repack [K,N] row-major -> [N,K] fp16 (coopmat B layout), + // N padded to a 128-row tile. + private readonly Dictionary _constGw = new(); + private MtlBuffer ConstGemmW(int tensorIndex, int K, int N) + { + if (_constGw.TryGetValue(tensorIndex, out MtlBuffer? cached)) return cached; + ReadOnlySpan f32 = + MemoryMarshal.Cast(_compiled.GetTensor(tensorIndex).Constant); + int nPad = (N + 127) / 128 * 128; + Half[] h = new Half[nPad * K]; + for (int k = 0; k < K; k++) + for (int n = 0; n < N; n++) + h[n * K + k] = (Half)f32[k * N + n]; + MtlBuffer buf = Up16(h); + _constGw[tensorIndex] = buf; + _allBufs.Add(buf); + return buf; + } + + // Managed schedule cache: pure metadata (no Metal objects), shared by + // every session of this model. The LRU bound only caps managed memory — + // eviction never touches a GPU resource and nothing is ever invalidated. + private readonly Dictionary _schedules = new(); + private const int MaxSchedules = 512; + private long _tick; + + internal static MetalSchedule.Key KeyOf(int[] inputShape, int nodeLimit, int outTensor) + => new(inputShape[0], inputShape[2], inputShape[3], nodeLimit, outTensor); + + // SIMD_OCR_NOCATABS=1 disables sole-consumer concat-slice writes — + // producers then write their own slot so DebugStats stays accurate. + private static readonly bool _noCatAbs = + Environment.GetEnvironmentVariable("SIMD_OCR_NOCATABS") != null; + private static readonly bool s_dbgTime = + Environment.GetEnvironmentVariable("SIMD_OCR_GPU_TIME") == "1"; + // SIMD_OCR_NOSK=1: narrow-M MatMuls use the plain dot kernel (bisecting) + private static readonly bool s_noSplitK = + Environment.GetEnvironmentVariable("SIMD_OCR_NOSK") != null; + // SIMD_OCR_NODD=1: keep the staged mm_sg tile even when K%16==0 (A/B) + private static readonly bool _noDd = + Environment.GetEnvironmentVariable("SIMD_OCR_NODD") != null; + + // SIMD_OCR_NOIC=1: materialize im2col + plain mm instead of implicit-GEMM (A/B) + private static readonly bool _noIc = + Environment.GetEnvironmentVariable("SIMD_OCR_NOIC") != null; + /// Schedule for (shape, nodeLimit, outTensor); compiled on first + /// use. nodeLimit truncates the emit loop and outTensor overrides the + /// readback tensor (REC stops before the vocab projection). + internal MetalSchedule GetSchedule(int[] inputShape, int nodeLimit = int.MaxValue, int outTensor = -1) + { + MetalSchedule.Key key = KeyOf(inputShape, nodeLimit, outTensor); + lock (this) + { + if (!_schedules.TryGetValue(key, out MetalSchedule? s)) + { + if (_schedules.Count >= MaxSchedules) + { + MetalSchedule.Key oldest = default; + long min = long.MaxValue; + foreach ((MetalSchedule.Key k, MetalSchedule v) in _schedules) + if (v.LastTick < min) { min = v.LastTick; oldest = k; } + _schedules.Remove(oldest); + } + long t0 = System.Diagnostics.Stopwatch.GetTimestamp(); + s = BuildPlan(inputShape, nodeLimit, outTensor); + _schedules[key] = s; + if (s_dbgTime) + Console.WriteLine($"[t] compile {key} {System.Diagnostics.Stopwatch.GetElapsedTime(t0).TotalMilliseconds:F2} ms"); + } + s.LastTick = ++_tick; + return s; + } + } + + internal MetalSchedule? TryGetSchedule(MetalSchedule.Key key) + { + lock (this) return _schedules.TryGetValue(key, out MetalSchedule? s) ? s : null; + } + + private unsafe MetalSchedule BuildPlan(int[] inputShape, int nodeLimit, int outTensor) + { + int[][] shapes = _compiled.ResolveShapesFor(inputShape); + NodeRecord[] nodes = _model.Nodes; + int nT = shapes.Length; + int outIdx = outTensor >= 0 ? outTensor + : checked((int)_model.GraphOutputs[0]); + int inIdx = checked((int)_model.GraphInputs[0]); + int emitLimit = Math.Min(nodeLimit, nodes.Length); + int nb = inputShape[0]; // batch — spatial kernels take it via gy + + var alias = new int[nT]; // metadata-only ops alias producer slot + var isConst = new bool[nT]; + var numel = new long[nT]; + var used = new bool[nT]; + for (int i = 0; i < nT; i++) + { + alias[i] = i; + numel[i] = 1; + foreach (int d in shapes[i]) numel[i] *= Math.Max(d, 1); + isConst[i] = _compiled.GetTensor(i).IsConstant; + } + int Phys(int t) { while (alias[t] != t) t = alias[t]; return t; } + + // refcount / producer / consumers for emit-time fusion decisions + var refCount = new int[nT]; + var prodNode = new int[nT]; + Array.Fill(prodNode, -1); + var consumers = new List[nT]; + // members of conv-headed fused groups are folded into the conv dispatch; + // groups headed by other ops emit every member standalone, so a Mul + // inside e.g. a HardSigmoid+Mul (hardswish) group is still foldable. + var inConvGroup = new bool[nodes.Length]; + for (int ni = 0; ni < nodes.Length; ni++) + { + int skip = _compiled.FusedSkip(ni); + if (nodes[ni].Operator is OperatorId.Conv or OperatorId.ConvTranspose) + for (int j = ni + 1; j <= ni + skip && j < nodes.Length; j++) + inConvGroup[j] = true; + foreach (uint inp in nodes[ni].Inputs) + { + if (inp == uint.MaxValue) continue; + refCount[(int)inp]++; + (consumers[(int)inp] ??= new List()).Add(ni); + } + prodNode[checked((int)nodes[ni].Outputs[0])] = ni; + } + + // ---- fused MHSA (SVTR mixer): the Transpose5D -> Slice/Squeeze -> + // q*scale, k^T -> MatMul -> Softmax -> MatMul -> Transpose chain + // collapses into one `attn` dispatch reading the packed qkv tensor. + // Members get no arena slots; the chain's out-Transpose slot is the + // kernel's destination (the trailing Reshape aliases onto it). + var attnAt = new Dictionary(); + var attnSkip = new HashSet(); + if (Environment.GetEnvironmentVariable("SIMD_OCR_NOATTN") == null) + { + var mem = new List(); + for (int ni = 0; ni < nodes.Length; ni++) + if (nodes[ni].Operator == OperatorId.Transpose + && MatchAttention(ni, mem) is { } spec) + { + attnAt[ni] = spec; + foreach (int m in mem) attnSkip.Add(m); + } + } + + (int qkv, int outT, int T, int H, int D, float scale)? MatchAttention(int tr, List mem) + { + mem.Clear(); + int Sole(int t) => refCount[t] == 1 && consumers[t] is { Count: 1 } c ? c[0] : -1; + bool Perm(int ni, params int[] perm) + { + ReadOnlySpan tp = _model.GetParameters(nodes[ni]); + if (U16(tp, 2) != perm.Length) return false; + for (int i = 0; i < perm.Length; i++) + if (I32(tp, 4 + 4 * i) != perm[i]) return false; + return true; + } + NodeRecord t5 = nodes[tr]; + int in5 = checked((int)t5.Inputs[0]); + int[] s5 = shapes[in5]; + if (s5.Length != 5 || s5[0] != nb || s5[2] != 3 || s5[4] > 32 + || !Perm(tr, 2, 0, 3, 1, 4)) return null; + int r5 = prodNode[in5]; + if (r5 < 0 || nodes[r5].Operator != OperatorId.Reshape) return null; + int qkvT = checked((int)nodes[r5].Inputs[0]); + int T = s5[1], H = s5[3], D = s5[4]; + if (shapes[qkvT].Length != 3 || numel[qkvT] != numel[in5]) return null; + int tOut = checked((int)t5.Outputs[0]); + List? cs = consumers[tOut]; + if (cs is not { Count: 3 }) return null; + int[] role = [-1, -1, -1]; + mem.Add(tr); + foreach (int sl in cs) + { + if (nodes[sl].Operator != OperatorId.Slice) return null; + (int[] st, int[] stp) = _compiled.ResolveSliceBounds(shapes[tOut], nodes[sl], out _); + int[] so = shapes[checked((int)nodes[sl].Outputs[0])]; + if (st.Length != 5 || so.Length != 5 || so[0] != 1 || (uint)st[0] > 2 + || role[st[0]] != -1) return null; + for (int a = 0; a < 5; a++) + if (stp[a] != 1 || (a > 0 && (st[a] != 0 || so[a] != shapes[tOut][a]))) + return null; + int sq = Sole(checked((int)nodes[sl].Outputs[0])); + if (sq < 0 || nodes[sq].Operator != OperatorId.Squeeze) return null; + role[st[0]] = checked((int)nodes[sq].Outputs[0]); + mem.Add(sl); mem.Add(sq); + } + // q: Mul(scalar const) [-> Reshape] -> MatMul1.A + int mul = Sole(role[0]); + if (mul < 0 || nodes[mul].Operator != OperatorId.Mul) return null; + int sc = nodes[mul].Inputs[0] == (uint)role[0] + ? checked((int)nodes[mul].Inputs[1]) : checked((int)nodes[mul].Inputs[0]); + if (!isConst[sc] || numel[sc] != 1) return null; + float scale = CstF32(sc)[0]; + mem.Add(mul); + int qt = checked((int)nodes[mul].Outputs[0]); + int mm1 = Sole(qt); + if (mm1 >= 0 && nodes[mm1].Operator == OperatorId.Reshape + && numel[checked((int)nodes[mm1].Outputs[0])] == numel[qt]) + { + mem.Add(mm1); + qt = checked((int)nodes[mm1].Outputs[0]); + mm1 = Sole(qt); + } + if (mm1 < 0 || nodes[mm1].Operator != OperatorId.MatMul + || nodes[mm1].Inputs[0] != (uint)qt) return null; + // k: Transpose[0,1,3,2] -> MatMul1.B + int kt = Sole(role[1]); + if (kt < 0 || nodes[kt].Operator != OperatorId.Transpose || !Perm(kt, 0, 1, 3, 2) + || nodes[mm1].Inputs[1] != nodes[kt].Outputs[0]) return null; + int smx = Sole(checked((int)nodes[mm1].Outputs[0])); + if (smx < 0 || nodes[smx].Operator != OperatorId.Softmax) return null; + int ax = I32(_model.GetParameters(nodes[smx]), 4); + if (ax != -1 && ax != 3) return null; + int mm2 = Sole(checked((int)nodes[smx].Outputs[0])); + if (mm2 < 0 || nodes[mm2].Operator != OperatorId.MatMul + || nodes[mm2].Inputs[0] != nodes[smx].Outputs[0] + || nodes[mm2].Inputs[1] != (uint)role[2] || Sole(role[2]) != mm2) return null; + int ot = Sole(checked((int)nodes[mm2].Outputs[0])); + if (ot < 0 || nodes[ot].Operator != OperatorId.Transpose || !Perm(ot, 0, 2, 1, 3)) + return null; + mem.Add(kt); mem.Add(mm1); mem.Add(smx); mem.Add(mm2); mem.Add(ot); + return (qkvT, checked((int)nodes[ot].Outputs[0]), T, H, D, scale); + } + + // ---- pass 1: aliases + which tensors need arena slots ---- + for (int ni = 0; ni < nodes.Length; ni++) + { + NodeRecord node = nodes[ni]; + int skip = _compiled.FusedSkip(ni); + if (attnSkip.Contains(ni)) continue; + // [n,1,a,b] -> [n,b,1,a] (perm 0,3,1,2): with a unit dim-1 the + // row-major input flat order already equals the NHWC output's. + if (node.Operator == OperatorId.Transpose + && shapes[checked((int)node.Inputs[0])] is { Length: 4 } ts4 && ts4[1] == 1 + && _model.GetParameters(node) is ReadOnlySpan p4 + && U16(p4, 2) == 4 && I32(p4, 4) == 0 && I32(p4, 8) == 3 + && I32(p4, 12) == 1 && I32(p4, 16) == 2) + { + alias[checked((int)node.Outputs[0])] = checked((int)node.Inputs[0]); + continue; + } + if (node.Operator is OperatorId.LayoutConvert or OperatorId.Squeeze + or OperatorId.Unsqueeze or OperatorId.Reshape) + { + alias[checked((int)node.Outputs[0])] = checked((int)node.Inputs[0]); + continue; + } + // rank-3 [n,c,w] -> [n,w,c] transpose is a no-op under the + // channel-last physical layout: out flat (n*w'+c') index equals the + // producer's (n*w+c) index because w'=w rows and c'=c channels swap + // positions but keep the same flat offset. + if (node.Operator == OperatorId.Transpose + && shapes[checked((int)node.Inputs[0])].Length == 3 + && _model.GetParameters(node) is ReadOnlySpan p2 + && U16(p2, 2) == 3 && I32(p2, 4) == 0 && I32(p2, 8) == 2 + && I32(p2, 12) == 1) + { + alias[checked((int)node.Outputs[0])] = checked((int)node.Inputs[0]); + continue; + } + foreach (uint inp in node.Inputs) + if (inp != uint.MaxValue && !isConst[checked((int)inp)]) + used[checked((int)inp)] = true; + // the slot that gets written = the fused group's last node output + int sink = checked((int)nodes[ni + skip].Outputs[0]); + used[sink] = true; + // fused members reading tensors produced BEFORE the group head + // (residuals, dual-write pre-gelu taps) need real slots too — + // in-group reads never touch the arena. + for (int j = ni + 1; j <= ni + skip; j++) + foreach (uint inp2 in nodes[j].Inputs) + { + if (inp2 == uint.MaxValue) continue; + int ip = checked((int)inp2); + if (!isConst[ip] && prodNode[ip] >= 0 && prodNode[ip] < ni) + used[Phys(ip)] = true; + } + // conv+bias+gelu groups dual-write the pre-gelu tensor when it also + // feeds a residual outside the group — reserve its slot + if (skip == 6 && node.Operator == OperatorId.Conv + && nodes[ni + 1].Operator == OperatorId.Add + && nodes[ni + 5].Operator == OperatorId.Mul) + { + // preG has exactly 2 in-group consumers (Div ni+2, Mul ni+5); + // anything beyond that is an external reader needing a slot + int preG = checked((int)nodes[ni + 1].Outputs[0]); + if (refCount[preG] > 2) used[preG] = true; + } + ni += skip; + } + foreach (uint g in _model.GraphOutputs) used[checked((int)g)] = true; + foreach (var at in attnAt.Values) { used[at.outT] = true; used[Phys(at.qkv)] = true; } + + // graph input gets physically padded to 4 channels when Cin%4!=0 so the + // input conv's im2col can use the vector path; weights pack zeros there. + int inCinPad = shapes[inIdx][1] % 4 != 0 ? 4 : 0; + + var off = new long[nT]; + var slabElems = new long[nT]; + long cursor = 0; + long maxIm2col = 0; // shared scratch: max M*K_pad over emitted im2col convs + for (int i = 0; i < nT; i++) + { + off[i] = -1; + if (!used[i] || isConst[i] || alias[i] != i) continue; + off[i] = cursor; + // pad every slab to a 128-row tile multiple: coopmat reads pad rows; + // slabs are 8-element (16 B) aligned: the sg32 coopmat stages + // 16 B uvec4 loads straight from slab starts + long sz = numel[i]; + if (i == inIdx && inCinPad != 0) + sz = numel[i] / shapes[i][1] * inCinPad; + slabElems[i] = (sz + 7) / 8 * 8 + 128 * 64; + cursor += slabElems[i]; + } + // the im2col scratch sits past every slab; its size (and the arena + // total) is settled after emit, so convs that take a direct or + // implicit-GEMM path reserve nothing + long im2colOff = (cursor + 63) / 64 * 64; + // Role sentinels: the session binds its own grow-only buffers at + // record time, sized from the schedule's requirements. + MtlBuffer arena = RoleArena, outF32 = RoleOut, inF32 = RoleIn; + bool inF32Consumed = false; // stem conv reads fp32 NCHW directly + bool convTOut = false; // last convT wrote fp32 out directly + + + // ---- pre-emit fusion detection ---- + // Mul(x, seChannelVec) (the SE gate) gets skipped when EVERY consumer can + // fold the gating: a pointwise Conv applies it as input-prescale, an Add + // applies it as scaled-residual-add (addps). Any unhandled consumer keeps + // the Mul emitted. + var prescale = new Dictionary(); // convNi -> tensors + var addScale = new Dictionary(); // addNi -> tensors + var skipEmit = new HashSet(); + for (int ni = 0; ni < nodes.Length; ni++) + { + NodeRecord nd = nodes[ni]; + if (nd.Operator != OperatorId.Mul || _compiled.FusedSkip(ni) != 0 || inConvGroup[ni]) + continue; + uint a = nd.Inputs[0], b = nd.Inputs[1]; + if (a == uint.MaxValue || b == uint.MaxValue + || isConst[(int)a] || isConst[(int)b]) continue; + // se = input whose numel equals the other's channel count + // (per-image: nb*C for a batch — se tensors are [n,C,1,1]) + long na = numel[(int)a], nelb = numel[(int)b]; + int[] sa = shapes[(int)a], sb = shapes[(int)b]; + uint xT, seT; + if (na == (sa.Length >= 2 ? (long)sb[1] * nb : -1) && nelb > na) { xT = b; seT = a; } + else if (nelb == (sb.Length >= 2 ? (long)sa[1] * nb : -1) && na > nelb) { xT = a; seT = b; } + else continue; + if (numel[(int)seT] % 4 != 0) continue; + + var pendPrescale = new List(); + var pendAdd = new List<(int c, uint oth)>(); + bool ok = true; + var stk = new Stack(); + var seen = new HashSet(); + stk.Push(nd.Outputs[0]); + while (stk.Count > 0 && ok) + { + uint t = stk.Pop(); + if (!seen.Add(t)) continue; + List? cs = consumers[(int)t]; + if (cs == null) { ok = false; break; } + foreach (int c in cs) + { + NodeRecord cn = nodes[c]; + if (cn.Operator is OperatorId.LayoutConvert or OperatorId.Squeeze + or OperatorId.Unsqueeze or OperatorId.Reshape) + { stk.Push(cn.Outputs[0]); continue; } + if (cn.Operator == OperatorId.Conv && !inConvGroup[c] + && Phys((int)cn.Inputs[0]) == Phys((int)nd.Outputs[0])) + { + ReadOnlySpan cp = _model.GetParameters(cn); + int ccin = shapes[(int)xT][1]; + int ccout = shapes[(int)cn.Outputs[0]][1]; + if (U32(cp, 4) == 1 && I32(cp, 8) == 1 && I32(cp, 12) == 1 + && I32(cp, 16) == 1 && I32(cp, 20) == 1 + && ccin % 4 == 0 && ccout % 4 == 0 && ccin <= 256 + && Environment.GetEnvironmentVariable("SIMD_OCR_NOPRESCALE") == null) + pendPrescale.Add(c); + else ok = false; + continue; + } + if (cn.Operator == OperatorId.Add && !inConvGroup[c] + && _compiled.FusedSkip(c) == 0 + && Environment.GetEnvironmentVariable("SIMD_OCR_NOADDPS") == null) + { + uint oth = cn.Inputs[0] == t ? cn.Inputs[1] + : cn.Inputs[1] == t ? cn.Inputs[0] : uint.MaxValue; + if (oth != uint.MaxValue && !isConst[(int)oth] + && numel[(int)oth] == numel[(int)xT] + && numel[(int)xT] % 4 == 0) + { pendAdd.Add((c, oth)); continue; } + ok = false; + continue; + } + ok = false; + } + } + if (!ok) + { + foreach (int c in pendPrescale) prescale.Remove(c); + continue; + } + foreach (int c in pendPrescale) prescale[c] = (xT, seT); + foreach ((int c, uint oth) in pendAdd) addScale[c] = (xT, seT, oth); + if (Environment.GetEnvironmentVariable("SIMD_OCR_PS_KEEP") != null) + continue; // debug: folds registered but mul still emitted + skipEmit.Add(ni); + } + + // SE chain fusion: ReduceMean -> Conv1x1(fc1,+bias) -> Conv1x1(fc2,+bias) + // -> HardSigmoid emits as a single `se` dispatch writing the gate vector + // directly to the HardSigmoid output slot. + var seEmit = new Dictionary(); + var concatAbsorbed = new HashSet(); // phys tensors written directly into a concat slice + + if (Environment.GetEnvironmentVariable("SIMD_OCR_NOSE") == null) + { + bool IsPwConv(int ni) + { + if (nodes[ni].Operator != OperatorId.Conv || _compiled.FusedSkip(ni) != 0 + || inConvGroup[ni]) return false; + ReadOnlySpan cp = _model.GetParameters(nodes[ni]); + return U32(cp, 4) == 1 && I32(cp, 8) == 1 && I32(cp, 12) == 1 + && I32(cp, 16) == 1 && I32(cp, 20) == 1 + && nodes[ni].Inputs.Length > 2 && nodes[ni].Inputs[2] != uint.MaxValue; + } + bool dbg = Environment.GetEnvironmentVariable("SIMD_OCR_GPU_DUMP") == "1"; + for (int hs = 0; hs < nodes.Length; hs++) + { + if (nodes[hs].Operator != OperatorId.HardSigmoid + || _compiled.FusedSkip(hs) != 0 || inConvGroup[hs]) continue; + int fc2 = prodNode[Phys(checked((int)nodes[hs].Inputs[0]))]; + if (dbg) + Console.Error.WriteLine($" se? hs=n{hs} fc2=n{fc2} fc2pw={(fc2 >= 0 && IsPwConv(fc2))} rc2={(fc2 >= 0 ? refCount[Phys(checked((int)nodes[fc2].Outputs[0]))] : -1)}"); + if (fc2 < 0 || !IsPwConv(fc2) + || refCount[Phys(checked((int)nodes[fc2].Outputs[0]))] != 1) continue; + int fc1 = prodNode[Phys(checked((int)nodes[fc2].Inputs[0]))]; + int reluNi = -1; + if (fc1 >= 0 && nodes[fc1].Operator == OperatorId.Relu + && refCount[Phys(checked((int)nodes[fc1].Outputs[0]))] == 1) + { reluNi = fc1; fc1 = prodNode[Phys(checked((int)nodes[fc1].Inputs[0]))]; } + // fc1 may be a Conv+Relu compiler group (skip=1 covering reluNi) + bool fc1ok = fc1 >= 0 && reluNi >= 0 + && nodes[fc1].Operator == OperatorId.Conv && !inConvGroup[fc1] + && _compiled.FusedSkip(fc1) == 1 && fc1 + 1 == reluNi + && nodes[fc1].Inputs.Length > 2 && nodes[fc1].Inputs[2] != uint.MaxValue; + if (fc1ok) + { + ReadOnlySpan cp = _model.GetParameters(nodes[fc1]); + fc1ok = U32(cp, 4) == 1 && I32(cp, 8) == 1 && I32(cp, 12) == 1 + && I32(cp, 16) == 1 && I32(cp, 20) == 1; + } + if (dbg) + Console.Error.WriteLine($" fc1=n{fc1} relu=n{reluNi} fc1ok={fc1ok}"); + if (!fc1ok + || refCount[Phys(checked((int)nodes[fc1].Outputs[0]))] != 1) continue; + int rm = prodNode[Phys(checked((int)nodes[fc1].Inputs[0]))]; + if (dbg) + Console.Error.WriteLine($" rm=n{rm} op={(rm >= 0 ? nodes[rm].Operator.ToString() : "?")} rcIn={(rm >= 0 ? refCount[Phys(checked((int)nodes[fc1].Inputs[0]))] : -1)}"); + if (rm < 0 || nodes[rm].Operator != OperatorId.ReduceMean + || _compiled.FusedSkip(rm) != 0 || inConvGroup[rm] + || refCount[Phys(checked((int)nodes[fc1].Inputs[0]))] != 1) continue; + int[] xs = shapes[Phys(checked((int)nodes[rm].Inputs[0]))]; + int[] rs = shapes[Phys(checked((int)nodes[fc1].Outputs[0]))]; + int rDim = rs.Length >= 2 ? rs[1] : 0; + if (dbg) + Console.Error.WriteLine($" xs=[{string.Join('x', xs)}] rDim={rDim}"); + if (xs.Length != 4 || xs[1] % 4 != 0 || xs[1] > 256 + || rDim <= 0 || rDim > 64) continue; + seEmit[rm] = (fc1, fc2, hs); + for (int m = fc1; m <= fc1 + _compiled.FusedSkip(fc1); m++) + skipEmit.Add(m); // conv16 + its fused relu + skipEmit.Add(fc2); skipEmit.Add(hs); + } + } + if (Environment.GetEnvironmentVariable("SIMD_OCR_GPU_DUMP") == "1") + { + Console.Error.WriteLine($"prescale convs={prescale.Count} addps={addScale.Count} seEmit={seEmit.Count} skipEmit={skipEmit.Count}"); + foreach (int rm in seEmit.Keys.OrderBy(k => k)) + Console.Error.WriteLine($" se chain rm=n{rm} fc1=n{seEmit[rm].fc1} fc2=n{seEmit[rm].fc2} hs=n{seEmit[rm].hs}"); + foreach (var kv in addScale) + { + int px = prodNode[(int)kv.Value.x], ps = prodNode[(int)kv.Value.se], po = prodNode[(int)kv.Value.oth]; + Console.Error.WriteLine($" addps n{kv.Key}: x=t{kv.Value.x}[{string.Join('x', shapes[(int)kv.Value.x])}] prod={(px < 0 ? "?" : nodes[px].Operator.ToString() + px)} " + + $"se=t{kv.Value.se}[{string.Join('x', shapes[(int)kv.Value.se])}] prod={(ps < 0 ? "?" : nodes[ps].Operator.ToString() + ps)} " + + $"oth=t{kv.Value.oth} prod={(po < 0 ? "?" : nodes[po].Operator.ToString() + po)}"); + } + foreach (var kv in prescale) + Console.Error.WriteLine($" prescale conv n{kv.Key}: x=t{kv.Value.x} se=t{kv.Value.se}"); + } + + // ---- pass 2: emit ---- + int LastEmitNi = -1; + List recs = new(); + MetalRec Emit(MtlPipeline pipe, string tag, (MtlBuffer buf, long elemOff, int elemSize)[] binds, + uint[] pc, uint gx, uint gy = 1) + { + var bb = new MetalBind[binds.Length]; + for (int b = 0; b < binds.Length; b++) + bb[b] = new MetalBind(binds[b].buf, (ulong)(binds[b].elemOff * binds[b].elemSize)); + MetalRec r = new() { Pipe = pipe, Binds = bb, Pc = Pcu(pc), Gx = gx, Gy = gy, + Tag = tag, Tx = ReferenceEquals(pipe, _pLn) || ReferenceEquals(pipe, _pAttn) ? 64u : 256u }; + recs.Add(r); + return r; + } + uint Div256(long n) => (uint)((n + 255) / 256); + // depthwise kernel choice (emit and addps absorb must agree): lite + // parts run k <= 3 on the flat kernel with the absorb read — its + // channel-quad-consecutive loads beat the tile's per-quad halo + // there; the rest goes to the shared-tile kernel, which sg32 also + // uses for the 9x9 taps + bool DwFlatA(int kh, int kw, int sh, int sw) => + _pDw4A is not null && kh <= 3 && kw <= 3 && sh <= 2 && sw <= 2; + bool DwTiled(int kh, int kw, int sh, int sw) => !DwFlatA(kh, kw, sh, sw) + && kh <= (_sg32 ? 9 : 5) && kw <= (_sg32 ? 9 : 5) && sh <= 2 && sw <= 2; + // spatial-reduce partitions: ~512 px per workgroup (a 16-WG cap left + // most of the GPU idle on large DET maps), bounded by partial storage + int PartSplits(int hw, int c) => + Math.Max(1, Math.Min(Math.Min((hw + 511) / 512, 256), PartCtrElem / (c * nb))); + // Metal GEMM emit: "cm" path → mm_sg (simdgroup 64x64 tile, K-step 16) + // when the live probe passed, else mm_dot (portable m4n4). Both read + // n-major fp16 weights [N,K] and take pc {M,N,K,kv4,flags,mImg,scMod}. + MetalRec EmitMm(bool ps, string tag, + (MtlBuffer buf, long elemOff, int elemSize)[] binds, + uint M, uint N, uint K, uint flags, uint mImg, uint scMod) + { + uint[] pc = [M, N, K, (uint)(K / 4), flags, mImg, scMod]; + if (_sgm) + // mm_sg_dd reads fragment rows past M/N on edge tiles — safe: + // A overread stays inside the arena's ≥1MB tail slack, W rows + // are padded to /128, and the epilogue masks mm(); + if (Environment.GetEnvironmentVariable("SIMD_OCR_NOAPS") == null) + foreach ((int addNi, (uint x, uint se, uint oth) asc) in addScale) + { + int outP = Phys(checked((int)nodes[addNi].Outputs[0])); + if (refCount[outP] != 1) continue; + List? cs = consumers[outP]; + if (cs == null || cs.Count != 1 || Absorbable(cs[0], outP) == false) + continue; + addpsSrc[outP] = (asc.oth, asc.x, asc.se); + } + + bool Absorbable(int cn, int srcPhys) + { + NodeRecord c = nodes[cn]; + switch (c.Operator) + { + case OperatorId.Conv: + { + if (Phys(checked((int)c.Inputs[0])) != srcPhys) return false; + ReadOnlySpan cp = _model.GetParameters(c); + int[] ish2 = shapes[checked((int)c.Inputs[0])]; + int[] osh2 = shapes[checked((int)c.Outputs[0])]; + int cin2 = ish2[1], cout2 = osh2[1]; + int kH2 = I32(cp, 8), kW2 = I32(cp, 12), + sH2 = I32(cp, 16), sW2 = I32(cp, 20); + int grp2 = Math.Max(1, checked((int)U32(cp, 4))); + if (grp2 == cin2 && cout2 == cin2) + return cin2 % 4 == 0 && (DwTiled(kH2, kW2, sH2, sW2) // conv_dw4t + || DwFlatA(kH2, kW2, sH2, sW2)); // conv_dw4a + if (grp2 == 1 && kH2 == 1 && kW2 == 1 && sH2 == 1 && sW2 == 1) + { + int cinI2 = (srcPhys == inIdx && cin2 % 4 != 0) ? 4 : cin2; + if (cinI2 % 4 != 0 || cout2 % 4 != 0 || cinI2 > 128) + return false; + int cskip = _compiled.FusedSkip(cn); + if (cskip == 1 && nodes[cn + 1].Operator == OperatorId.Add) + { + NodeRecord ad2 = nodes[cn + 1]; + int bt = ad2.Inputs[0] == c.Outputs[0] + ? checked((int)ad2.Inputs[1]) + : checked((int)ad2.Inputs[0]); + if (numel[bt] == 1) return false; // scalarBias + } + return !prescale.ContainsKey(cn) + && Environment.GetEnvironmentVariable( + "SIMD_OCR_NODOT") == null; + } + return false; + } + case OperatorId.Resize: + return Phys(checked((int)c.Inputs[0])) == srcPhys + && shapes[checked((int)c.Outputs[0])][1] % 4 == 0; + case OperatorId.Concat: + { + int[] ospc = shapes[checked((int)c.Outputs[0])]; + if (ospc[1] % 4 != 0) return false; + int cOff2 = 0; + foreach (uint inp2 in c.Inputs) + { + int ci2 = shapes[checked((int)inp2)][1]; + if (Phys(checked((int)inp2)) == srcPhys) + return ci2 % 4 == 0 && cOff2 % 4 == 0; + cOff2 += ci2; + } + return false; + } + case OperatorId.Add: + { + // consumer Add is the tail of a resize4add fusion: the + // other input must come from a Resize that will itself + // hit the fused path (sole-consumer, %4, same numel). + if (addScale.ContainsKey(cn)) return false; + int i0 = Phys(checked((int)c.Inputs[0])); + int i1 = Phys(checked((int)c.Inputs[1])); + int othP2 = i0 == srcPhys ? i1 : i1 == srcPhys ? i0 : -1; + if (othP2 < 0) return false; + int prod = prodNode[othP2]; + if (prod < 0 || nodes[prod].Operator != OperatorId.Resize) + return false; + int[] osp3 = shapes[checked((int)nodes[prod].Outputs[0])]; + if (osp3[1] % 4 != 0) return false; + if (refCount[othP2] != 1 || consumers[othP2]?.Count != 1 + || consumers[othP2][0] != cn) return false; + return numel[srcPhys] == numel[othP2]; + } + default: return false; + } + } + + // ---- catresize fusion: a Concat whose inputs are sole-consumed + // Resize nodes (or plain tensors) collapses to ONE dispatch that + // resamples each source straight into its output slice. The resizes' + // own recs are skipped (they were the writers, now the kernel writes). + var catRes = new Dictionary(); + // catresize kernel is batch-free: batched graphs keep the plain path + for (int ci2 = 0; nb == 1 && ci2 < nodes.Length; ci2++) + { + NodeRecord c = nodes[ci2]; + if (c.Operator != OperatorId.Concat || c.Inputs.Length > 4) + continue; + int[] osp = shapes[checked((int)c.Outputs[0])]; + if (osp[1] % 4 != 0) continue; + var spec = new (long, long, long, uint, uint, uint, uint, uint, uint) + [c.Inputs.Length]; + int cOff3 = 0; bool okc = true; var skipTmp = new List(); + for (int ii = 0; ii < c.Inputs.Length; ii++) + { + uint inp2 = c.Inputs[ii]; + int ip = Phys(checked((int)inp2)); + int ci3 = shapes[checked((int)inp2)][1]; + if (ci3 % 4 != 0 || cOff3 % 4 != 0) { okc = false; break; } + uint fh3 = 1, fw3 = 1, inW3 = (uint)osp[3]; + int srcP = ip; + int prod = prodNode[ip]; + if (prod >= 0 && nodes[prod].Operator == OperatorId.Resize + && refCount[ip] == 1 && consumers[ip]?.Count == 1 + && consumers[ip][0] == ci2) + { + int[] rish = shapes[checked((int)nodes[prod].Inputs[0])]; + if (rish[1] != ci3) { okc = false; break; } + fh3 = (uint)(osp[2] / rish[2]); + fw3 = (uint)(osp[3] / rish[3]); + if (rish[2] * (int)fh3 != osp[2] || rish[3] * (int)fw3 != osp[3]) + { okc = false; break; } + inW3 = (uint)rish[3]; + srcP = Phys(checked((int)nodes[prod].Inputs[0])); + skipTmp.Add(prod); + } + else if (refCount[ip] == 1 && consumers[ip]?.Count == 1 + && consumers[ip][0] == ci2 && prod >= 0 + && !addpsSrc.ContainsKey(ip)) + { + // sole-consumed producer emits a slice-write into this + // concat instead of materializing its own tensor — + // catRes cannot read it back; keep the old path. + okc = false; break; + } + uint fl = 0; long xP = off[srcP], raP = 0, seP = 0; + if (addpsSrc.TryGetValue(srcP, out var ads)) + { + xP = off[Phys(checked((int)ads.x))]; + raP = off[Phys(checked((int)ads.a))]; + seP = off[Phys(checked((int)ads.se))]; + fl = 1; + } + spec[ii] = (xP, raP, seP, fl, inW3, fh3, fw3, + (uint)(cOff3 / 4), (uint)(ci3 / 4)); + cOff3 += ci3; + } + if (!okc) continue; + catRes[ci2] = spec; + foreach (int p2 in skipTmp) skipEmit.Add(p2); + } + + for (int ni = 0; ni < emitLimit; ni++) + { + NodeRecord node = nodes[ni]; + if (skipEmit.Contains(ni)) continue; + if (attnAt.TryGetValue(ni, out var att)) + { + LastEmitNi = ni; + Emit(_pAttn, $"attn n{ni} T{att.T} H{att.H} D{att.D}", + [(arena, off[Phys(att.qkv)], 2), (arena, off[att.outT], 2)], + [(uint)att.T, (uint)att.H, (uint)att.D, + BitConverter.SingleToUInt32Bits(att.scale)], + (uint)((att.T + 63) / 64), (uint)(nb * att.H)); + continue; + } + if (attnSkip.Contains(ni)) continue; + try + { + int skip = _compiled.FusedSkip(ni); + ReadOnlySpan p = _model.GetParameters(node); + int outPhys = Phys(checked((int)nodes[ni + skip].Outputs[0])); + LastEmitNi = ni; + long SlotOf(uint t) => off[Phys(checked((int)t))]; + + switch (node.Operator) + { + case OperatorId.LayoutConvert: + case OperatorId.Squeeze: + case OperatorId.Unsqueeze: + case OperatorId.Reshape: + break; + + case OperatorId.Conv: + { + int[] ishp = shapes[node.Inputs[0]]; + int[] oshp = shapes[node.Outputs[0]]; + int cin = ishp[1], cout = oshp[1]; + int inH = ishp[2], inW = ishp[3], outH = oshp[2], outW = oshp[3]; + int kH = I32(p, 8), kW = I32(p, 12), sH = I32(p, 16), sW = I32(p, 20); + int pt = I32(p, 32), pl = I32(p, 36); + int group = Math.Max(1, checked((int)U32(p, 4))); + + int act = 0; + MtlBuffer biasBuf = arena; + uint resT = uint.MaxValue, dualT = uint.MaxValue; + uint hasBias = 0, scalarBias = 0, flags16 = 0; + bool dualRaw = false; + if (node.Inputs.Length > 2 && node.Inputs[2] != uint.MaxValue) + { biasBuf = ConstF16(checked((int)node.Inputs[2])); hasBias = 1; } + if (skip == 1 && nodes[ni + 1].Operator == OperatorId.Relu) + act = 1; + else if (skip == 1 && nodes[ni + 1].Operator == OperatorId.Add) + { + // conv (no bias) + channel-bias add + NodeRecord add = nodes[ni + 1]; + int bt = add.Inputs[0] == node.Outputs[0] + ? checked((int)add.Inputs[1]) : checked((int)add.Inputs[0]); + biasBuf = ConstF16(bt); hasBias = 1; + if (numel[bt] == 1) scalarBias = 8u; + } + else if (skip == 2 && nodes[ni + 1].Operator == OperatorId.HardSigmoid + && nodes[ni + 2].Operator == OperatorId.Mul) + { + act = 3; + ReadOnlySpan hp = _model.GetParameters(nodes[ni + 1]); + float hsA = F32(hp, 4), hsB = F32(hp, 8); + if (hsB != 0.5f) + throw new NotSupportedException($"hardsigmoid beta={hsB} at node {ni + 1}"); + flags16 |= (uint)Half2Bits(hsA) << 16; + } + else if (skip == 6 + && nodes[ni + 1].Operator == OperatorId.Add + && nodes[ni + 2].Operator == OperatorId.Div + && nodes[ni + 3].Operator == OperatorId.Erf + && nodes[ni + 4].Operator == OperatorId.Add + && nodes[ni + 5].Operator == OperatorId.Mul + && nodes[ni + 6].Operator == OperatorId.Mul + && Phys(checked((int)nodes[ni + 2].Inputs[0])) + == Phys(checked((int)nodes[ni + 1].Outputs[0])) + && (Phys(checked((int)nodes[ni + 5].Inputs[0])) + == Phys(checked((int)nodes[ni + 1].Outputs[0])) + || Phys(checked((int)nodes[ni + 5].Inputs[1])) + == Phys(checked((int)nodes[ni + 1].Outputs[0])))) + { + // compiler-fused conv + bias-Add + gelu chain: + // out = 0.5*x*(1+erf(x/√2)) where x = conv + bias. + NodeRecord add = nodes[ni + 1]; + int bt = add.Inputs[0] == node.Outputs[0] + ? checked((int)add.Inputs[1]) + : checked((int)add.Inputs[0]); + biasBuf = ConstF16(bt); hasBias = 1; + if (numel[bt] == 1) scalarBias = 8u; + act = 2; + int preG = Phys(checked((int)nodes[ni + 1].Outputs[0])); + if (refCount[preG] > 2) + { + // conv+bias also feeds a residual outside the + // group: sink slot gets gelu(acc) via o; the preG + // slot gets raw acc via o2 (flag bit10). + dualT = (uint)preG; dualRaw = true; + } + } + else if (skip == 2 + && nodes[ni + 1].Operator == OperatorId.Add + && nodes[ni + 2].Operator == OperatorId.Add) + { + // conv + channel-bias Add + residual Add + NodeRecord add = nodes[ni + 1]; + int bt = add.Inputs[0] == node.Outputs[0] + ? checked((int)add.Inputs[1]) + : checked((int)add.Inputs[0]); + if (!isConst[bt]) + throw new NotSupportedException( + $"conv+add+add mid-add non-const at node {ni + 1}"); + biasBuf = ConstF16(bt); hasBias = 1; + if (numel[bt] == 1) scalarBias = 8u; + NodeRecord res = nodes[ni + 2]; + int mid = Phys(checked((int)add.Outputs[0])); + uint oth = Phys(checked((int)res.Inputs[0])) == mid + ? res.Inputs[1] + : Phys(checked((int)res.Inputs[1])) == mid + ? res.Inputs[0] : uint.MaxValue; + if (oth == uint.MaxValue || isConst[(int)oth] + || numel[(int)oth] != numel[mid]) + throw new NotSupportedException( + $"conv+add+add residual form at node {ni + 2}"); + resT = oth; + } + else if (skip != 0) + throw new NotSupportedException( + $"conv fused skip={skip} at node {ni} (next={nodes[ni + 1].Operator})"); + + // ---- emit-time fusion extensions (compiler didn't group these) ---- + // pointwise conv + residual Add (+ optional gelu pattern) — + // cm/dot kernels implement res binding (flags bit1) and act (bits4-6). + if (group == 1 && kH == 1 && kW == 1 && sH == 1 && sW == 1 + && Environment.GetEnvironmentVariable("SIMD_OCR_NOEXTEND") == null) + { + int ge = ni + skip; + int j = ge + 1; + // residual Add may only be absorbed when the conv group + // carries no activation: kernels evaluate act(conv+res), + // but act(conv)+res is the ONNX semantics (e.g. relu then + // add must NOT become relu of the sum). + if (j < nodes.Length && nodes[j].Operator == OperatorId.Add + && !inConvGroup[j] && _compiled.FusedSkip(j) == 0 + && !addScale.ContainsKey(j) && act == 0 + && Environment.GetEnvironmentVariable("SIMD_OCR_NOEXTR") == null) + { + NodeRecord ad = nodes[j]; + int prev = Phys(checked((int)nodes[ge].Outputs[0])); + int ia = Phys(checked((int)ad.Inputs[0])); + int ib = Phys(checked((int)ad.Inputs[1])); + uint other = ia == prev ? ad.Inputs[1] + : ib == prev ? ad.Inputs[0] : uint.MaxValue; + if (other != uint.MaxValue && !isConst[(int)other] + && numel[(int)other] == numel[prev] + && refCount[prev] == 1) + { resT = other; ge = j; j++; } + } + if (j < nodes.Length + && nodes[j].Operator == OperatorId.Div + && Environment.GetEnvironmentVariable("SIMD_OCR_NOEXTG") == null + && _compiled.FusedSkip(j) == 4 + && Phys(checked((int)nodes[j].Inputs[0])) + == Phys(checked((int)nodes[ge].Outputs[0]))) + { + int preG = Phys(checked((int)nodes[ge].Outputs[0])); + if (refCount[preG] == 1 && act == 0) + { act = 2; ge = j + 4; } + else if (refCount[preG] == 2) + { + // out also feeds a residual elsewhere — conv dual- + // writes plain out (preG) and gelu(out) (gelu sink) + dualT = (uint)Phys(checked((int)nodes[j + 4].Outputs[0])); + outPhys = preG; + ge = j + 4; + } + } + if (ge != ni + skip) + { + if (dualT == uint.MaxValue) + outPhys = Phys(checked((int)nodes[ge].Outputs[0])); + skip = ge - ni; + } + } + bool hasPs = prescale.TryGetValue(ni, out (uint x, uint se) psv); + uint flags = hasBias | scalarBias | (uint)(act << 4) | flags16 + | (resT != uint.MaxValue ? 2u : 0u) + | (hasPs ? 4u : 0u) + | (dualT != uint.MaxValue ? 8u : 0u) + | (dualRaw ? 1024u : 0u); + + if (kH == 1 && kW == 1 && sH == 1 && sW == 1 && group == 1) + { + int inT = Phys(checked((int)node.Inputs[0])); + int cinIn = (inT == inIdx && cin % 4 != 0) ? 4 : cin; + if (cinIn % 4 != 0) + throw new NotSupportedException($"conv1x1 cin={cin} not %4 at node {ni}"); + MtlBuffer wbuf = ConstF16(checked((int)node.Inputs[1]), + padRows: ((cout + 127) / 128) * 128, rowElems: cin, + rowPad: cinIn); + uint mImg = (uint)(outH * outW); + uint M = mImg * (uint)nb; // flat rows span the batch + bool dotOk = cout % 4 == 0 && scalarBias == 0 + && (cinIn <= 128 || (hasPs && cinIn <= 256)) + && Environment.GetEnvironmentVariable("SIMD_OCR_NODOT") == null; + // sg32 coopmat (SE prescale included) outruns the dot + // kernel once the output leaves split-K range (lite: + // a quarter of that already); + // addps-absorbed inputs need the dot kernel's fused read + if (dotOk && _sg32 && (long)M * cout > 1L << (_lite ? 16 : 18) + && (hasPs || !addpsSrc.ContainsKey(Phys(checked((int)node.Inputs[0]))))) + dotOk = false; + if (dotOk) + { + // small-K pointwise conv: direct dot kernel beats coopmat. + // prescale: read x (not the mul output) + per-channel se scale. + long aOff = hasPs ? SlotOf(psv.x) : SlotOf(node.Inputs[0]); + bool aps = addpsSrc.TryGetValue( + Phys(checked((int)node.Inputs[0])), out var ads) + && !hasPs; + if (aps) aOff = SlotOf(ads.x); + MtlBuffer wk = ConstKMajor(checked((int)node.Inputs[1]), + cout, cin, 1, cinIn, cinIn); + // se tensors are per-batch [n, C] — mImg lets the + // kernel pick the batch's scale vector; 0 = batch-free. + uint mIm = hasPs || aps ? mImg : 0u; + uint dFlags = flags | (aps ? 512u : 0u); + long dTiles = ((M + 3) / 4) * (long)(cout / 4); + bool dSplit = cinIn >= 64 && dTiles <= 16384 && !s_noSplitK + && (dFlags & (8u | 512u | 1024u)) == 0; + Emit(dSplit ? _pDotSk : _pDot, $"conv1x1d{(dSplit ? "sk" : "")} n{ni} {M}x{cinIn}x{cout}", + [(arena, aOff, 2), (wk, 0, 2), + (biasBuf, 0, 2), + (arena, resT != uint.MaxValue ? SlotOf(resT) : 0, 2), + (arena, off[outPhys], 2), + (arena, hasPs ? SlotOf(psv.se) + : aps ? SlotOf(ads.se) : 0, 2), + (arena, dualT != uint.MaxValue ? SlotOf(dualT) : 0, 2), + (arena, aps ? SlotOf(ads.a) : 0, 2)], + [M, (uint)cout, (uint)cinIn, (uint)(cinIn / 4), + dFlags, (uint)(cinIn / 4), mIm], + dSplit ? (uint)((dTiles + 15) / 16) : Div256(dTiles)); + } + else + { + EmitMm(hasPs, $"conv1x1 n{ni} {M}x{cin}x{cout}", + [(arena, hasPs ? SlotOf(psv.x) : SlotOf(node.Inputs[0]), 2), (wbuf, 0, 2), + (biasBuf, 0, 2), + (arena, resT != uint.MaxValue ? SlotOf(resT) : 0, 2), + (arena, off[outPhys], 2), + (arena, dualT != uint.MaxValue ? SlotOf(dualT) : 0, 2), + (arena, hasPs ? SlotOf(psv.se) : 0, 2)], + M, (uint)cout, (uint)cinIn, flags, + hasPs ? mImg : 0u, (uint)(cinIn / 4)); + } + } + else if (group == cin && cout == cin) + { + if (cin % 4 == 0) + { + if (DwTiled(kH, kW, sH, sW)) + { + // shared-tile variant: ~kH*kW less global traffic + uint tx = (uint)((outW + 15) / 16), + ty = (uint)((outH + 15) / 16); + bool aps = addpsSrc.TryGetValue( + Phys(checked((int)node.Inputs[0])), out var ads); + Emit(_pDw4T, $"conv_dw4t n{ni} {cout}ch k{kH}s{sH}", + [(arena, aps ? SlotOf(ads.x) + : SlotOf(node.Inputs[0]), 2), + (ConstDwTap(checked((int)node.Inputs[1]), cout, kH * kW), 0, 2), + (biasBuf, 0, 2), (arena, off[outPhys], 2), + (arena, aps ? SlotOf(ads.a) : 0, 2), + (arena, aps ? SlotOf(ads.se) : 0, 2)], + [(uint)outW, (uint)outH, (uint)cout, (uint)kH, (uint)kW, + (uint)sH, (uint)sW, (uint)pt, (uint)pl, + (uint)inW, (uint)inH, + flags | (aps ? 256u : 0u)], + tx * ty * (uint)(cout / 4), (uint)nb); + } + else if (DwFlatA(kH, kW, sH, sW)) + { + bool aps = addpsSrc.TryGetValue( + Phys(checked((int)node.Inputs[0])), out var ads); + Emit(_pDw4A!, $"conv_dw4a n{ni} {cout}ch k{kH}s{sH}", + [(arena, aps ? SlotOf(ads.x) + : SlotOf(node.Inputs[0]), 2), + (ConstDwTap(checked((int)node.Inputs[1]), cout, kH * kW), 0, 2), + (biasBuf, 0, 2), (arena, off[outPhys], 2), + (arena, aps ? SlotOf(ads.a) : 0, 2), + (arena, aps ? SlotOf(ads.se) : 0, 2)], + [(uint)outW, (uint)outH, (uint)cout, (uint)kH, (uint)kW, + (uint)sH, (uint)sW, (uint)pt, (uint)pl, + (uint)inW, (uint)inH, + flags | (aps ? 256u : 0u)], + Div256((long)outW * outH * (cout / 4)), (uint)nb); + } + else + Emit(_pDw4, $"conv_dw4 n{ni} {cout}ch k{kH}s{sH}", + [(arena, SlotOf(node.Inputs[0]), 2), + (ConstDwTap(checked((int)node.Inputs[1]), cout, kH * kW), 0, 2), + (biasBuf, 0, 2), (arena, off[outPhys], 2)], + [(uint)outW, (uint)outH, (uint)cout, (uint)kH, (uint)kW, + (uint)sH, (uint)sW, (uint)pt, (uint)pl, + (uint)inW, (uint)inH, flags], + Div256((long)outW * outH * (cout / 4)), + (uint)nb); + } + else + { + if (nb > 1) + throw new NotSupportedException( + $"scalar dw conv lacks batch support (n{ni})"); + Emit(_pDw, $"conv_dw n{ni} {cout}ch k{kH}s{sH}", + [(arena, SlotOf(node.Inputs[0]), 2), + (ConstF16(checked((int)node.Inputs[1])), 0, 2), + (biasBuf, 0, 2), (arena, off[outPhys], 2)], + [(uint)outW, (uint)outH, (uint)cout, (uint)kH, (uint)kW, + (uint)sH, (uint)sW, (uint)pt, (uint)pl, + (uint)inW, (uint)inH, flags], + Div256((long)outW * outH * cout)); + } + } + else if (group == 1) + { + // dense kxk conv → explicit im2col (tap-major) + coopmat GEMM. + // graph-input channels are physically padded to 4 (cinIn). + int inT = Phys(checked((int)node.Inputs[0])); + int cinIn = (inT == inIdx && cin % 4 != 0) ? 4 : cin; + int K = cinIn * kH * kW, Kp = (K + 15) / 16 * 16; + uint M = (uint)(outH * outW); + // im2col-free direct conv: tap addressing inline in the + // dot kernel; skips materializing the K-expanded matrix. + // batched graphs always take it: the im2col fallback is batch-free + // sg32: implicit-GEMM coopmat conv instead (same tap + // addressing, tensor cores; any K, batched or not) + // Metal has no implicit-GEMM kxk kernel: the dot + // path covers every case (batched or K<=1024); the + // rest falls to im2col + mm below. + bool convkCm = false; + if (convkCm || (cout % 4 == 0 && scalarBias == 0 && cinIn % 4 == 0 + && (nb > 1 || Kp <= (Environment.GetEnvironmentVariable( + "SIMD_OCR_CONVD_KMAX") is string km + ? int.Parse(km) : 1024)) + && Environment.GetEnvironmentVariable("SIMD_OCR_NODCONV") == null)) + { + // stem conv on the fp32 NCHW input: skip nchw2nhwc + // when this is the sole consumer of the graph input + bool f32In = !convkCm && inT == inIdx && cinIn == 4 && cin <= 4 + && Environment.GetEnvironmentVariable( + "SIMD_OCR_NOF32IN") == null; + if (f32In && (consumers[inIdx]?.Count ?? 0) == 1) + inF32Consumed = true; + // sole-consumer concat: write straight into its slice + long wO = off[outPhys]; uint dcv = 0, dco = 0; + { + List? cs = consumers[outPhys]; + if (!_noCatAbs && refCount[outPhys] == 1 && cs != null + && cs.Count == 1 + && nodes[cs[0]].Operator == OperatorId.Concat) + { + int cn = cs[0]; + int[] ospc = shapes[nodes[cn].Outputs[0]]; + int co = 0; bool hit = false; + foreach (uint cin_ in nodes[cn].Inputs) + { + if (Phys(checked((int)cin_)) == outPhys) + { hit = true; break; } + co += shapes[checked((int)cin_)][1]; + } + if (hit && co % 4 == 0 && ospc[1] % 4 == 0) + { + wO = off[Phys(checked((int)nodes[cn].Outputs[0]))]; + dcv = (uint)(ospc[1] / 4); + dco = (uint)(co / 4); + concatAbsorbed.Add(outPhys); + } + } + } + if (_sgm && !_noIc && nb == 1 && !f32In + && !concatAbsorbed.Contains(outPhys)) + { + // unbatched conv: implicit-GEMM (tap gather in + // the MMA staging loop) beats the scalar-dot + // convd kernel — same weights, tap-major pack. + MtlBuffer wic = ConstTapMajor(checked((int)node.Inputs[1]), + cout, cin, kH, kW, Kp, cinPad: cinIn); + Emit(_pMmIc32, $"icconv n{ni} {M}x{Kp}x{cout}", + [(arena, SlotOf(node.Inputs[0]), 2), + (wic, 0, 2), + (biasBuf, 0, 2), (arena, 0, 2), + (arena, off[outPhys], 2), + (arena, 0, 2), (arena, 0, 2)], + [M, (uint)cout, (uint)Kp, (uint)(Kp / 4), + flags, 0u, 0u, + (uint)outW, (uint)inW, (uint)inH, + (uint)sH, (uint)sW, (uint)pt, (uint)pl, + (uint)kW, (uint)cinIn], + ((uint)cout + 63) / 64, (M + 63) / 64); + } + else + Emit(f32In ? _pConvDF32 : _pConvD, + $"convd n{ni} {M}x{K}x{cout}", + [(f32In ? inF32 : arena, + f32In ? 0 : SlotOf(node.Inputs[0]), + f32In ? 4 : 2), + (ConstKMajor(checked((int)node.Inputs[1]), + cout, cin, kH * kW, Kp, cinIn), 0, 2), + (biasBuf, 0, 2), + (arena, resT != uint.MaxValue ? SlotOf(resT) : 0, 2), + (arena, wO, 2)], + f32In + ? [M, (uint)cout, (uint)outW, (uint)inW, (uint)inH, + (uint)(cinIn / 4), (uint)(kH * kW), (uint)kW, + (uint)sH, (uint)sW, (uint)pt, (uint)pl, flags, + dcv, dco, (uint)cin, (uint)(inH * inW)] + : [M, (uint)cout, (uint)outW, (uint)inW, (uint)inH, + (uint)(cinIn / 4), (uint)(kH * kW), (uint)kW, + (uint)sH, (uint)sW, (uint)pt, (uint)pl, flags, + dcv, dco], + Div256(((M + 3) / 4) * (long)(cout / 4)), + (uint)nb); + } + else if (_sgm && !_noIc && cinIn % 4 == 0 && nb == 1) + { + // implicit-GEMM conv: mm_ic_sg gathers A through + // tap-major im2col addressing inside the staging + // loop — no [M,Kp] materialization round-trip. + MtlBuffer wic = ConstTapMajor(checked((int)node.Inputs[1]), + cout, cin, kH, kW, Kp, cinPad: cinIn); + Emit(_pMmIc32, $"icconv n{ni} {M}x{Kp}x{cout}", + [(arena, SlotOf(node.Inputs[0]), 2), + (wic, 0, 2), + (biasBuf, 0, 2), (arena, 0, 2), + (arena, off[outPhys], 2), + (arena, 0, 2), (arena, 0, 2)], + [M, (uint)cout, (uint)Kp, (uint)(Kp / 4), + flags, 0u, 0u, + (uint)outW, (uint)inW, (uint)inH, + (uint)sH, (uint)sW, (uint)pt, (uint)pl, + (uint)kW, (uint)cinIn], + ((uint)cout + 63) / 64, (M + 63) / 64); + } + else + { + if (nb > 1) + throw new NotSupportedException( + $"im2col conv path lacks batch support (n{ni})"); + maxIm2col = Math.Max(maxIm2col, (long)M * Kp); + Emit(_pIm2col, $"im2col n{ni} M{M} K{K}", + [(arena, SlotOf(node.Inputs[0]), 2), + (arena, im2colOff, 2)], + [(uint)outW, (uint)outH, (uint)cinIn, (uint)kH, (uint)kW, + (uint)sH, (uint)sW, (uint)pt, (uint)pl, + (uint)inW, (uint)inH, (uint)K, (uint)Kp, + cinIn % 4 != 0 ? 1u : 0u], + Div256(M * (long)(Kp / 4))); + MtlBuffer wbuf = ConstTapMajor(checked((int)node.Inputs[1]), + cout, cin, kH, kW, Kp, cinPad: cinIn); + if (cout % 4 == 0 && scalarBias == 0 && Kp <= 128) + { + MtlBuffer wk = ConstKMajor(checked((int)node.Inputs[1]), + cout, cin, kH * kW, Kp, cinIn); + Emit(_pDot, $"im2c_dot n{ni} {M}x{Kp}x{cout}", + [(arena, im2colOff, 2), (wk, 0, 2), + (biasBuf, 0, 2), (arena, 0, 2), + (arena, off[outPhys], 2), (arena, 0, 2), + (arena, 0, 2)], + [M, (uint)cout, (uint)Kp, (uint)(Kp / 4), + flags, 0u, 0u], + Div256(((M + 3) / 4) * (long)(cout / 4))); + } + else + { + // ConstTapMajor W is already n-major [Cout][Kp] + // (row = tap*cinPad+ci) — mm_* read it directly. + EmitMm(false, $"im2colconv n{ni} {M}x{Kp}x{cout}", + [(arena, im2colOff, 2), (wbuf, 0, 2), + (biasBuf, 0, 2), (arena, 0, 2), (arena, off[outPhys], 2), + (arena, 0, 2), (arena, 0, 2)], + M, (uint)cout, (uint)Kp, flags, 0u, 0u); + } + } + } + else + { + if (nb > 1) + throw new NotSupportedException( + $"dense conv path lacks batch support (n{ni})"); + Emit(_pDense, $"conv3x3 n{ni} {cin}->{cout} k{kH}s{sH}", + [(arena, SlotOf(node.Inputs[0]), 2), + (ConstF16(checked((int)node.Inputs[1])), 0, 2), + (biasBuf, 0, 2), (arena, off[outPhys], 2)], + [(uint)outW, (uint)outH, (uint)cout, (uint)cin, + (uint)kH, (uint)kW, (uint)sH, (uint)sW, + (uint)pt, (uint)pl, (uint)inW, (uint)inH, flags], + Div256((long)outW * outH * cout)); + } + ni += skip; + break; + } + + case OperatorId.Div when skip == 4: + { + // fused GELU: write to the group sink (last Mul output) + long n = numel[Phys(checked((int)nodes[ni + 4].Outputs[0]))]; + Emit(n % 4 == 0 ? _pElem4 : _pElem, $"gelu n{ni}", + [(arena, SlotOf(node.Inputs[0]), 2), + (arena, 0, 2), (arena, off[outPhys], 2)], + [(uint)n, 1u, 2u, 0u, 0u], Div256(n % 4 == 0 ? n / 4 : n)); + ni += 4; + break; + } + + case OperatorId.Add or OperatorId.Mul or OperatorId.Sub or OperatorId.Div + or OperatorId.Pow: + { + long n = numel[outPhys]; + if (addScale.TryGetValue(ni, out (uint x, uint se, uint oth) asc)) + { + if (addpsSrc.ContainsKey(outPhys)) break; // consumer absorbs + // a + x*se — the SE Mul got skipped, fold it here + // per-batch se vector: imgV = vec4 count of one image + long seN = numel[(int)asc.se]; + Emit(_pAddPs, $"addps n{ni}", + [(arena, SlotOf(asc.oth), 2), (arena, SlotOf(asc.x), 2), + (arena, SlotOf(asc.se), 2), (arena, off[outPhys], 2)], + [(uint)n, (uint)(seN / 4 / nb), + (uint)(seN == n ? 0 : n / 4 / nb)], Div256(n / 4)); + break; + } + int[] osp = shapes[checked((int)node.Outputs[0])]; + int chan = osp.Length >= 2 ? osp[1] : (int)n; + (MtlBuffer bA, long oA, uint mA) = Operand(checked((int)node.Inputs[0]), chan, n); + (MtlBuffer bB, long oB, uint mB) = Operand(checked((int)node.Inputs[1]), chan, n); + // vector broadcast along the contiguous (last physical) dim: + // rank-3 tails like [n,T,C] + [C] need C as the modulus — + // the vector operand's own numel, not osp[1]. + { + long neA = numel[checked((int)node.Inputs[0])]; + long neB = numel[checked((int)node.Inputs[1])]; + // per-image channel vector [n,C,1,1] (nel == chan*nb): + // mode 3 = kernel picks b[(i/perImg)*C + i%C] using aux; + // flat tail vector [C] vs [n,T,C]: mode 1, chan = neB. + if (mB == 2 && neB > 1 && neB < n) + { if (neB == (long)chan * nb) mB = 3; else { mB = 1; chan = (int)neB; } } + else if (mA == 2 && neA > 1 && neA < n) + { if (neA == (long)chan * nb) mA = 3; else { mA = 1; chan = (int)neA; } } + } + uint op2 = node.Operator switch + { + OperatorId.Add => 8u, OperatorId.Mul => 9u, OperatorId.Sub => 10u, + OperatorId.Div => 11u, OperatorId.Pow => 13u, _ => 8u, + }; + bool v4 = n % 4 == 0 && (mA is not (1 or 3) || chan % 4 == 0) + && (mB is not (1 or 3) || chan % 4 == 0); + Emit(v4 ? _pElem4 : _pElem, $"bin{node.Operator} n{ni}", + [(bA, oA, 2), (bB, oB, 2), (arena, off[outPhys], 2)], + [(uint)n, (uint)chan, op2, mA, mB, (uint)(n / nb)], Div256(v4 ? n / 4 : n)); + break; + } + + case OperatorId.Relu or OperatorId.Sigmoid or OperatorId.HardSigmoid + or OperatorId.Erf or OperatorId.Sqrt: + { + // compiler-fused HS+Mul (hardswish, x*hs(x)): emit op 14 + // against the group's sink instead of a bare hardsigmoid. + bool hswish = node.Operator == OperatorId.HardSigmoid + && skip == 1 + && nodes[ni + 1].Operator == OperatorId.Mul + && Phys(checked((int)nodes[ni + 1].Inputs[0])) + == Phys(checked((int)node.Outputs[0])) + && Phys(checked((int)nodes[ni + 1].Inputs[1])) + == Phys(checked((int)node.Inputs[0])); + if (hswish) ni += 1; + long n = numel[outPhys]; + uint op2 = node.Operator switch + { + OperatorId.Relu => 1u, OperatorId.HardSigmoid => 3u, + OperatorId.Sigmoid => 4u, OperatorId.Erf => 5u, + OperatorId.Sqrt => 6u, _ => 0u, + }; + if (hswish) op2 = 14u; + uint aux = 0u; + if (node.Operator == OperatorId.HardSigmoid) + aux = HsAux(_model.GetParameters(node)); + Emit(n % 4 == 0 ? _pElem4 : _pElem, $"un{node.Operator} n{ni}", + [(arena, SlotOf(node.Inputs[0]), 2), + (arena, 0, 2), (arena, off[outPhys], 2)], + [(uint)n, 1u, op2, 0u, 0u, aux], Div256(n % 4 == 0 ? n / 4 : n)); + break; + } + + case OperatorId.ReduceMean when skip == 8: + { + // compiler-fused LayerNorm over the last axis + int xT = checked((int)node.Inputs[0]); + int[] xs = shapes[xT]; + int cN = xs[^1]; + int xP = Phys(xT); + // last axis must be physically contiguous: rank-3 row-major, + // or an NHWC slab whose channel dim is that axis + bool contiguous = xs.Length == 3 + && (shapes[xP].Length != 4 || shapes[xP][1] == cN); + int ConstIn(NodeRecord nd, int other) + { + int a0 = checked((int)nd.Inputs[0]), a1 = checked((int)nd.Inputs[1]); + int c = Phys(a0) == Phys(other) ? a1 : Phys(a1) == Phys(other) ? a0 : -1; + return c >= 0 && isConst[c] ? c : -1; + } + NodeRecord nSub = nodes[ni + 1], nPow = nodes[ni + 2], nRm = nodes[ni + 3], + nEps = nodes[ni + 4], nSq = nodes[ni + 5], nDiv = nodes[ni + 6], + nG = nodes[ni + 7], nB = nodes[ni + 8]; + int powC = nPow.Operator == OperatorId.Pow + ? ConstIn(nPow, checked((int)nSub.Outputs[0])) : -1; + int epsC = nEps.Operator == OperatorId.Add + ? ConstIn(nEps, checked((int)nRm.Outputs[0])) : -1; + int gC = nG.Operator == OperatorId.Mul + ? ConstIn(nG, checked((int)nDiv.Outputs[0])) : -1; + int bC = nB.Operator == OperatorId.Add + ? ConstIn(nB, checked((int)nG.Outputs[0])) : -1; + bool lastAxis = U16(p, 2) == 1 + && (I32(p, 12) == -1 || I32(p, 12) == xs.Length - 1); + if (!contiguous || !lastAxis || cN > 1024 || numel[xT] / cN > 65535 + || nSub.Operator != OperatorId.Sub + || nRm.Operator != OperatorId.ReduceMean || nSq.Operator != OperatorId.Sqrt + || nDiv.Operator != OperatorId.Div || powC < 0 || epsC < 0 || gC < 0 + || bC < 0 || CstF32(powC)[0] != 2f || numel[epsC] != 1 + || numel[gC] != cN || numel[bC] != cN) + throw new NotSupportedException($"layernorm pattern at node {ni}"); + long rowsLn = numel[xT] / cN; + Emit(_pLn, $"layernorm n{ni} {rowsLn}x{cN}", + [(arena, SlotOf(node.Inputs[0]), 2), (ConstF32(gC), 0, 4), + (ConstF32(bC), 0, 4), (arena, off[outPhys], 2)], + [(uint)rowsLn, (uint)cN, BitConverter.SingleToUInt32Bits(CstF32(epsC)[0])], + (uint)rowsLn); + ni += 8; + break; + } + + case OperatorId.ReduceMean: + { + int[] ishp = shapes[node.Inputs[0]]; + int hw = ishp[2] * ishp[3], c = ishp[1]; + int cv4 = c / 4; + if (seEmit.TryGetValue(ni, out (int fc1, int fc2, int hs) sev)) + { + NodeRecord f1 = nodes[sev.fc1], f2 = nodes[sev.fc2]; + ReadOnlySpan hpp = _model.GetParameters(nodes[sev.hs]); + int rDim = shapes[Phys(checked((int)f1.Outputs[0]))][1]; + uint cvP = 1; while (cvP < (uint)(c / 4)) cvP <<= 1; + int hsOut = Phys(checked((int)nodes[sev.hs].Outputs[0])); + int s = PartSplits(hw, c); + int pp = (hw + s - 1) / s; + MtlBuffer pb2 = PartBuf(s * c * nb); + uint[] pcSe = [(uint)hw, (uint)c, (uint)s, (uint)pp, cvP, + (uint)rDim, + BitConverter.SingleToUInt32Bits(F32(hpp, 4)), + BitConverter.SingleToUInt32Bits(F32(hpp, 8))]; + // paravirt atomics are relaxed-only: two dispatches + // (partials → barrier → single-WG join) replace the + // Vulkan ticket handshake. + Emit(_pSePart, $"se_p n{ni} c{c} s{s} r{rDim}", + [(arena, SlotOf(node.Inputs[0]), 2), (pb2, 0, 2)], + pcSe, (uint)s, (uint)nb); + Emit(_pSeJoin, $"se_j n{ni} c{c} s{s} r{rDim}", + [(pb2, 0, 2), + (ConstF32(checked((int)f1.Inputs[1])), 0, 4), + (ConstF32(checked((int)f1.Inputs[2])), 0, 4), + (ConstF32(checked((int)f2.Inputs[1])), 0, 4), + (ConstF32(checked((int)f2.Inputs[2])), 0, 4), + (arena, off[hsOut], 2)], + pcSe, (uint)nb, 1); + break; + } + // lite parts take it at any size: reduce_hw's one-channel + // workgroups read 2 bytes at a C*2 stride + if (c % 4 == 0 && 256 % cv4 == 0 && (hw >= 4096 || _lite)) + { + // two-phase: S pixel partitions → fp32 partials → mean + int s = PartSplits(hw, c); + int pp = (hw + s - 1) / s; + MtlBuffer pb = PartBuf(s * c * nb); + Emit(_pReduce4, $"reduce4a n{ni} c{c}", + [(arena, SlotOf(node.Inputs[0]), 2), (pb, 0, 2)], + [(uint)hw, (uint)c, (uint)s, (uint)pp], + (uint)s, (uint)nb); + Emit(_pReduce4b, $"reduce4b n{ni} c{c}", + [(pb, 0, 2), (arena, off[outPhys], 2)], + [(uint)s, (uint)c, (uint)hw], (uint)c, (uint)nb); + } + else + Emit(_pReduce, $"reduceHW n{ni} c{c}", + [(arena, SlotOf(node.Inputs[0]), 2), (arena, off[outPhys], 2)], + [(uint)hw, (uint)c], (uint)c, (uint)nb); + break; + } + + case OperatorId.MaxPool: + { + int[] ishp = shapes[node.Inputs[0]]; + // both kernels hard-code 2x2 / stride 1 / pad-end + if (!(I32(p, 8) == 2 && I32(p, 12) == 2 && I32(p, 16) == 1 + && I32(p, 20) == 1 && I32(p, 24) == 0 && I32(p, 28) == 0 + && shapes[node.Outputs[0]][2] == ishp[2] + && shapes[node.Outputs[0]][3] == ishp[3])) + throw new NotSupportedException($"maxpool shape at node {ni}"); + if (ishp[1] % 4 == 0) + { + long wO = off[outPhys]; uint dcv = 0, dco = 0; + { + int rp = outPhys; + List? cs = consumers[rp]; + if (!_noCatAbs && refCount[rp] == 1 && cs != null && cs.Count == 1 + && nodes[cs[0]].Operator == OperatorId.Concat) + { + int cn = cs[0]; + int[] ospc = shapes[nodes[cn].Outputs[0]]; + int co = 0; bool hit = false; + foreach (uint cin_ in nodes[cn].Inputs) + { + if (Phys(checked((int)cin_)) == rp) { hit = true; break; } + co += shapes[checked((int)cin_)][1]; + } + if (hit && co % 4 == 0 && ospc[1] % 4 == 0) + { + wO = off[Phys(checked((int)nodes[cn].Outputs[0]))]; + dcv = (uint)(ospc[1] / 4); + dco = (uint)(co / 4); + concatAbsorbed.Add(rp); + } + } + } + Emit(_pPool4, $"maxpool4 n{ni}", + [(arena, SlotOf(node.Inputs[0]), 2), (arena, wO, 2)], + [(uint)ishp[3], (uint)ishp[2], (uint)ishp[1], dcv, dco], + Div256(numel[outPhys] / 4 / nb), (uint)nb); + } + else if (nb > 1) + throw new NotSupportedException($"scalar maxpool lacks batch support (n{ni})"); + else + Emit(_pPool, $"maxpool n{ni}", + [(arena, SlotOf(node.Inputs[0]), 2), (arena, off[outPhys], 2)], + [(uint)ishp[3], (uint)ishp[2], (uint)ishp[1]], + Div256(numel[outPhys])); + break; + } + + case OperatorId.Resize: + { + if (nb > 1) + throw new NotSupportedException( + $"resize lacks batch support (n{ni})"); + int[] ishp = shapes[node.Inputs[0]]; + int[] osp = shapes[node.Outputs[0]]; + if (osp[1] % 4 == 0) + { + // sole-consumer fusion: resize -> Add (FPN merge) or + // resize -> Concat (write into the channel slice) + int rp = outPhys; + List? cs = consumers[rp]; + int cn = refCount[rp] == 1 && cs != null && cs.Count == 1 + ? cs[0] : -1; + if (cn >= 0 && nodes[cn].Operator == OperatorId.Add + && _compiled.FusedSkip(cn) == 0 && !inConvGroup[cn] + && !addScale.ContainsKey(cn)) + { + NodeRecord ad = nodes[cn]; + uint oth = ad.Inputs[0] == node.Outputs[0] + ? ad.Inputs[1] : ad.Inputs[0]; + int othP = Phys(checked((int)oth)); + if (!isConst[othP] && numel[othP] == numel[rp] + && Phys(checked((int)ad.Inputs[0])) + != Phys(checked((int)ad.Inputs[1]))) + { + int ao = Phys(checked((int)ad.Outputs[0])); + // addps-absorb on either operand + long xOff4 = SlotOf(node.Inputs[0]); + long aOff4 = SlotOf(oth), raOff4 = 0, seOff4 = 0; + uint raFl = 0; + if (addpsSrc.TryGetValue( + Phys(checked((int)node.Inputs[0])), out var adx)) + { xOff4 = SlotOf(adx.x); raOff4 = SlotOf(adx.a); + seOff4 = SlotOf(adx.se); raFl |= 256u; } + if (addpsSrc.TryGetValue( + Phys(checked((int)oth)), out var ada)) + { aOff4 = SlotOf(ada.x); raOff4 = SlotOf(ada.a); + seOff4 = SlotOf(ada.se); raFl |= 512u; } + Emit(_pResize4Add, $"resizeadd n{ni}+{cn}", + [(arena, xOff4, 2), + (arena, aOff4, 2), + (arena, off[ao], 2), + (arena, raOff4, 2), (arena, seOff4, 2)], + [(uint)osp[3], (uint)osp[2], (uint)osp[1], + (uint)(osp[2] / ishp[2]), (uint)(osp[3] / ishp[3]), + raFl], + Div256(numel[rp] / 4)); + skipEmit.Add(cn); + break; + } + } + long wO = off[outPhys]; uint dcv = 0, dco = 0; + if (!_noCatAbs && cn >= 0 && nodes[cn].Operator == OperatorId.Concat) + { + int[] ospc = shapes[nodes[cn].Outputs[0]]; + int co = 0; bool hit = false; + foreach (uint cin_ in nodes[cn].Inputs) + { + if (Phys(checked((int)cin_)) == rp) { hit = true; break; } + co += shapes[checked((int)cin_)][1]; + } + if (hit && co % 4 == 0 && ospc[1] % 4 == 0) + { + wO = off[Phys(checked((int)nodes[cn].Outputs[0]))]; + dcv = (uint)(ospc[1] / 4); + dco = (uint)(co / 4); + concatAbsorbed.Add(rp); + } + } + bool aps = addpsSrc.TryGetValue( + Phys(checked((int)node.Inputs[0])), out var ads); + Emit(_pResize4, $"resize4 n{ni} x{osp[2] / ishp[2]}", + [(arena, aps ? SlotOf(ads.x) + : SlotOf(node.Inputs[0]), 2), + (arena, wO, 2), + (arena, aps ? SlotOf(ads.a) : 0, 2), + (arena, aps ? SlotOf(ads.se) : 0, 2)], + [(uint)osp[3], (uint)osp[2], (uint)osp[1], + (uint)(osp[2] / ishp[2]), (uint)(osp[3] / ishp[3]), + dcv, dco, aps ? 256u : 0u], + Div256(numel[outPhys] / 4)); + } + else + Emit(_pResize, $"resize n{ni} x{osp[2] / ishp[2]}", + [(arena, SlotOf(node.Inputs[0]), 2), (arena, off[outPhys], 2)], + [(uint)osp[3], (uint)osp[2], (uint)osp[1], + (uint)(osp[2] / ishp[2]), (uint)(osp[3] / ishp[3])], + Div256(numel[outPhys])); + break; + } + + case OperatorId.Concat: + { + int[] osp = shapes[node.Outputs[0]]; + if (osp.Length != 4 || (I32(p, 4) is int cax && cax != 1 && cax != -3)) + throw new NotSupportedException($"concat axis/rank at node {ni}"); + // NHWC channel concat: pixels are flat over the batch too + long catPx = (long)osp[2] * osp[3] * nb; + if (catRes.TryGetValue(ni, out var spec)) + { + var bb = new (MtlBuffer, long, int)[13]; + var pc2 = new uint[19]; + pc2[0] = (uint)osp[3]; pc2[1] = (uint)osp[2]; + pc2[2] = (uint)osp[1]; + for (int s = 0; s < 4; s++) + { + long xP = 0, raP = 0, seP = 0; + uint fl = 0, inW3 = 0, fhfw = 0, offc = 0; + if (s < spec.Length) + { + var t = spec[s]; + xP = t.Item1; raP = t.Item2; seP = t.Item3; + fl = t.Item4; inW3 = t.Item5; + fhfw = t.Item6 | (t.Item7 << 16); + offc = t.Item8 | (t.Item9 << 16); + } + bb[s * 3] = (arena, xP, 2); + bb[s * 3 + 1] = (arena, raP, 2); + bb[s * 3 + 2] = (arena, seP, 2); + pc2[3 + s * 4] = offc; pc2[4 + s * 4] = inW3; + pc2[5 + s * 4] = fhfw; pc2[6 + s * 4] = fl; + } + bb[12] = (arena, off[outPhys], 2); + Emit(_pCatRes, $"catres n{ni}", bb, pc2, + Div256(numel[outPhys] / 4)); + break; + } + int cOff = 0; + foreach (uint inp in node.Inputs) + { + int it = checked((int)inp); + int ci = shapes[it][1]; + if (concatAbsorbed.Contains(Phys(it))) + { cOff += ci; continue; } + if (ci % 4 == 0 && cOff % 4 == 0 && osp[1] % 4 == 0) + { + bool aps = addpsSrc.TryGetValue(Phys(it), out var ads); + if (aps && nb > 1) // concat4's se index is batch-free + throw new NotSupportedException($"batched addps concat at node {ni}"); + Emit(_pConcat4, $"concat4 n{ni} c{ci}@{cOff}", + [(arena, aps ? SlotOf(ads.x) : SlotOf(inp), 2), + (arena, off[outPhys], 2), + (arena, aps ? SlotOf(ads.a) : 0, 2), + (arena, aps ? SlotOf(ads.se) : 0, 2)], + [(uint)catPx, (uint)ci, (uint)osp[1], + (uint)cOff, aps ? 256u : 0u], + Div256(catPx * (ci / 4))); + } + else + Emit(_pConcat, $"concat n{ni} c{ci}@{cOff}", + [(arena, SlotOf(inp), 2), (arena, off[outPhys], 2)], + [(uint)catPx, (uint)ci, (uint)osp[1], (uint)cOff], + Div256(catPx * ci)); + cOff += ci; + } + break; + } + + case OperatorId.ConvTranspose: + { + if (nb > 1) + throw new NotSupportedException( + $"convtranspose lacks batch support (n{ni})"); + int[] ishp = shapes[node.Inputs[0]]; + int[] oshp = shapes[node.Outputs[0]]; + if (!(I32(p, 8) == 2 && I32(p, 12) == 2 && I32(p, 16) == 2 && I32(p, 20) == 2 + && I32(p, 32) == 0 && I32(p, 36) == 0)) + throw new NotSupportedException($"ConvTranspose shape at node {ni}"); + int act = 0; + MtlBuffer biasBuf = arena; + uint flags2 = 0; + if (skip == 2 && nodes[ni + 1].Operator == OperatorId.Add) + { + NodeRecord add = nodes[ni + 1]; + int bt = add.Inputs[0] == node.Outputs[0] + ? checked((int)add.Inputs[1]) : checked((int)add.Inputs[0]); + biasBuf = ConstF16(bt); + flags2 = 1u | (numel[bt] == 1 ? 8u : 0u); + OperatorId actOp = nodes[ni + 2].Operator; + act = actOp == OperatorId.Relu ? 1 : actOp == OperatorId.Sigmoid ? 4 : 0; + } + if (ishp[1] % 4 == 0 && oshp[1] % 4 == 0) + Emit(_pConvT4, $"convT4 n{ni} {ishp[1]}->{oshp[1]}", + [(arena, SlotOf(node.Inputs[0]), 2), + (ConstConvT(checked((int)node.Inputs[1]), ishp[1], oshp[1]), 0, 2), + (biasBuf, 0, 2), (arena, off[outPhys], 2)], + [(uint)oshp[3], (uint)oshp[2], (uint)oshp[1], (uint)ishp[1], + (uint)ishp[3], flags2 | (uint)(act << 4)], + Div256(numel[outPhys] / 4)); + else + { + // scalar-tail convT writing the graph output: emit fp32 + // straight into the readback buffer (kills the `out` rec) + if (outPhys == Phys(outIdx)) { flags2 |= 128u; convTOut = true; } + Emit(_pConvT, $"convT n{ni} {ishp[1]}->{oshp[1]}", + [(arena, SlotOf(node.Inputs[0]), 2), + (ConstF16(checked((int)node.Inputs[1])), 0, 2), + (biasBuf, 0, 2), (arena, off[outPhys], 2), + (outF32, 0, 4)], + [(uint)oshp[3], (uint)oshp[2], (uint)oshp[1], (uint)ishp[1], + (uint)ishp[3], flags2 | (uint)(act << 4)], + Div256(numel[outPhys])); + } + ni += skip; + break; + } + + case OperatorId.Transpose: + // rank-3 [0,2,1] transposes were folded into aliases in + // pass 1 (channel-last physical layout makes them free). + if (Phys(checked((int)node.Outputs[0])) + != Phys(checked((int)node.Inputs[0]))) + throw new NotSupportedException( + $"non-trivial Transpose at node {ni}"); + break; + + case OperatorId.AveragePool: + { + int[] ishp = shapes[node.Inputs[0]]; + int[] osp2 = shapes[node.Outputs[0]]; + if (ishp[1] % 4 != 0 || ishp.Length != 4) + throw new NotSupportedException( + $"avgpool C%4 at node {ni}"); + Emit(_pAvg4, $"avgpool4 n{ni}", + [(arena, SlotOf(node.Inputs[0]), 2), + (arena, off[outPhys], 2)], + [(uint)osp2[3], (uint)osp2[2], (uint)ishp[1], + (uint)I32(p, 8), (uint)I32(p, 12), + (uint)I32(p, 16), (uint)I32(p, 20), + (uint)I32(p, 24), (uint)I32(p, 28), + (uint)ishp[3], (uint)ishp[2]], + Div256(numel[outPhys] / 4), (uint)nb); + break; + } + + case OperatorId.BatchNormalization: + { + // fold to per-channel affine: o = x*s4 + t4 where + // s = scale/sqrt(var+eps), t = bias - mean*s + float eps = F32(p, 4); + int bt = checked((int)node.Inputs[1]); + ReadOnlySpan sc = CstF32(bt), + bi = CstF32(checked((int)node.Inputs[2])), + mn = CstF32(checked((int)node.Inputs[3])), + va = CstF32(checked((int)node.Inputs[4])); + int cc = sc.Length; + float[] s = new float[cc], t = new float[cc]; + for (int i = 0; i < cc; i++) + { + s[i] = sc[i] / MathF.Sqrt(va[i] + eps); + t[i] = bi[i] - mn[i] * s[i]; + } + long n = numel[outPhys]; + if (n % 4 != 0 || cc % 4 != 0) + throw new NotSupportedException($"bn vec4 at node {ni}"); + Emit(_pAffine, $"bn4 n{ni} c{cc}", + [(arena, SlotOf(node.Inputs[0]), 2), + (VecF16((ni, 0), s), 0, 2), (VecF16((ni, 1), t), 0, 2), + (arena, off[outPhys], 2)], + [(uint)(n / 4), (uint)(cc / 4)], Div256(n / 4)); + break; + } + + case OperatorId.Softmax: + { + // last-axis softmax (cls head [n,2]); one invocation per row. + int[] ssp = shapes[node.Inputs[0]]; + long cols = ssp[^1]; + long rows = numel[outPhys] / cols; + if (cols <= 0 || rows <= 0 || rows * cols != numel[outPhys]) + throw new NotSupportedException($"softmax at node {ni}"); + Emit(_pSoftmax, $"softmax n{ni} {rows}x{cols}", + [(arena, SlotOf(node.Inputs[0]), 2), + (arena, off[outPhys], 2)], + [(uint)rows, (uint)cols], Div256(rows)); + break; + } + + case OperatorId.MatMul: + { + // [n,M,K] x [K,N]: flat rows over batch feed the coopmat + // 1x1-conv GEMM as-is; weights repack to [N,K] fp16. + int[] asp = shapes[node.Inputs[0]]; + int[] bsp = shapes[node.Inputs[1]]; + int mmK = asp[^1], mmN = bsp[^1]; + if (bsp.Length != 2 || !isConst[(int)node.Inputs[1]] + || mmK % 4 != 0 || mmN % 4 != 0) + throw new NotSupportedException($"matmul at node {ni}: asp=[{string.Join(',', asp)}] bsp=[{string.Join(',', bsp)}] const={isConst[(int)node.Inputs[1]]}"); + uint mmM = (uint)(numel[checked((int)node.Inputs[0])] / mmK); + // fold a following bias Add [N] into the GEMM epilogue + uint mmFlags = 0; MtlBuffer mmBias = arena; + if (ni + 1 < emitLimit + && nodes[ni + 1].Operator == OperatorId.Add + && _compiled.FusedSkip(ni + 1) == 0) + { + NodeRecord ad = nodes[ni + 1]; + uint oth = ad.Inputs[0] == node.Outputs[0] ? ad.Inputs[1] + : ad.Inputs[1] == node.Outputs[0] ? ad.Inputs[0] + : uint.MaxValue; + if (oth != uint.MaxValue && isConst[(int)oth] + && numel[(int)oth] == mmN + && refCount[Phys(checked((int)node.Outputs[0]))] == 1) + { + mmBias = ConstF16(checked((int)oth)); + mmFlags |= 1u; + outPhys = Phys(checked((int)ad.Outputs[0])); + ni += 1; + } + } + // narrow-M token GEMMs (SVTR mixer, n*T rows): the direct + // dot kernel keeps occupancy where a 128x128 coopmat tile + // would leave most of the M dimension idle. Its epilogue + // also folds swish (Sigmoid+Mul) or a residual Add. + bool mmDot = mmM <= 16384 + && Environment.GetEnvironmentVariable("SIMD_OCR_MMCM") == null; + if (mmDot) + { + int cur = outPhys; + long resOff = 0; + int j = ni + 1; + bool Free(int t) => t != Phys(outIdx) && refCount[t] >= 1; + if (j + 1 < emitLimit && (mmFlags & 1u) != 0 + && nodes[j].Operator == OperatorId.Sigmoid + && nodes[j + 1].Operator == OperatorId.Mul + && _compiled.FusedSkip(j) == 0 && _compiled.FusedSkip(j + 1) == 0 + && Phys(checked((int)nodes[j].Inputs[0])) == cur + && refCount[cur] == 2 && Free(cur) + && refCount[Phys(checked((int)nodes[j].Outputs[0]))] == 1 + && ((Phys(checked((int)nodes[j + 1].Inputs[0])) == cur + && Phys(checked((int)nodes[j + 1].Inputs[1])) + == Phys(checked((int)nodes[j].Outputs[0]))) + || (Phys(checked((int)nodes[j + 1].Inputs[1])) == cur + && Phys(checked((int)nodes[j + 1].Inputs[0])) + == Phys(checked((int)nodes[j].Outputs[0]))))) + { + mmFlags |= 5u << 4; + outPhys = Phys(checked((int)nodes[j + 1].Outputs[0])); + ni += 2; + } + else if (j < emitLimit && nodes[j].Operator == OperatorId.Add + && _compiled.FusedSkip(j) == 0 && refCount[cur] == 1 && Free(cur) + && !addScale.ContainsKey(j)) + { + NodeRecord ad = nodes[j]; + int ia = Phys(checked((int)ad.Inputs[0])); + int ib = Phys(checked((int)ad.Inputs[1])); + int oth = ia == cur ? checked((int)ad.Inputs[1]) + : ib == cur ? checked((int)ad.Inputs[0]) : -1; + if (oth >= 0 && !isConst[oth] && numel[oth] == numel[cur] + && Phys(oth) != cur) + { + resOff = SlotOf((uint)oth); + mmFlags |= 2u; + outPhys = Phys(checked((int)ad.Outputs[0])); + ni += 1; + } + } + // fewer than ~64 workgroups of 4x4 tiles: the serial K + // loop is the latency floor, so split K over 16 lanes + long tiles = ((mmM + 3) / 4) * (long)(mmN / 4); + bool splitK = mmK >= 64 && tiles <= 16384 && !s_noSplitK; + Emit(splitK ? _pDotSk : _pDot, $"mmdot{(splitK ? "sk" : "")} n{ni} {mmM}x{mmK}x{mmN}", + [(arena, SlotOf(node.Inputs[0]), 2), + (ConstF16(checked((int)node.Inputs[1])), 0, 2), + (mmBias, 0, 2), (arena, resOff, 2), + (arena, off[outPhys], 2), (arena, 0, 2), + (arena, 0, 2), (arena, 0, 2)], + [mmM, (uint)mmN, (uint)mmK, (uint)(mmK / 4), mmFlags, 0u, 0u], + splitK ? (uint)((tiles + 15) / 16) : Div256(tiles)); + break; + } + EmitMm(false, $"matmul n{ni} {mmM}x{mmK}x{mmN}", + [(arena, SlotOf(node.Inputs[0]), 2), + (ConstGemmW(checked((int)node.Inputs[1]), mmK, mmN), 0, 2), + (mmBias, 0, 2), (arena, 0, 2), + (arena, off[outPhys], 2), + (arena, 0, 2), (arena, 0, 2)], + mmM, (uint)mmN, (uint)mmK, mmFlags, 0u, 0u); + break; + } + + default: + throw new NotSupportedException($"op {node.Operator} at node {ni}"); + } + } + catch (Exception ex) + { + throw new NotSupportedException( + $"BuildPlan failed at node {ni} op={node.Operator} " + + $"in=[{string.Join(',', node.Inputs)}] out=[{string.Join(',', node.Outputs)}]", ex); + } + } + + // finalize: graph output → fp32 readback. + // If the last node is a STANDALONE Sigmoid (its own elem dispatch was + // the last rec), replace it with a sigmoid-ing f32 writeback. If the + // sigmoid was consumed into a fused group (e.g. convT+bias+sigmoid), + // the sink slot already holds sigmoid'd values → plain copy. + int lastNi = recs.Count > 0 ? LastEmitNi : -1; + NodeRecord last = nodes[nodes.Length - 1]; + bool standaloneSigmoid = lastNi == nodes.Length - 1 + && last.Operator == OperatorId.Sigmoid + && checked((int)last.Outputs[0]) == outIdx; + int outSrc; + uint outAct; + if (standaloneSigmoid) + { + recs.RemoveAt(recs.Count - 1); + outSrc = Phys(checked((int)last.Inputs[0])); + outAct = 4u; + } + else + { + outSrc = Phys(outIdx); + outAct = 0u; + } + if (!convTOut) + Emit(_pOut, "out", [(arena, off[outSrc], 2), (outF32, 0, 4)], + [(uint)numel[outIdx], outAct], Div256(numel[outIdx])); + + // head: fp32 NCHW input → fp16 NHWC (channel-padded to inCinPad if set). + // Skipped entirely when the stem conv read the fp32 input itself. + if (!inF32Consumed) + { + int[] isp = shapes[inIdx]; + int hw = isp[2] * isp[3], c = isp[1]; + uint cout = inCinPad != 0 ? (uint)inCinPad : (uint)c; + recs.Insert(0, new MetalRec + { + Pipe = _pNchw, + Binds = [new MetalBind(inF32, 0), new MetalBind(arena, (ulong)off[inIdx] * 2)], + Gx = Div256((long)hw * cout), Gy = (uint)nb, + Pc = Pcu((uint)hw, (uint)c, cout), Tag = "nchw2nhwc", + }); + } + + if (Environment.GetEnvironmentVariable("SIMD_OCR_GPU_DUMP") == "1") + for (int ri = 0; ri < recs.Count; ri++) + { + var r = recs[ri]; + Console.Error.WriteLine($"rec[{ri}] {r.Tag} gx={r.Gx} gy={r.Gy} " + + $"binds=[{string.Join(", ", r.Binds.Select(b => $"{(ReferenceEquals(b.Buf, RoleArena) ? "arena" : ReferenceEquals(b.Buf, RoleIn) ? "in" : ReferenceEquals(b.Buf, RoleOut) ? "out" : ReferenceEquals(b.Buf, RolePart) ? "part" : "buf")}+{(long)b.ByteOffset}"))}] " + + $"pc={Convert.ToHexString(r.Pc)}"); + } + + // device-independent, but the sg16 (Arc) kernel set has not been + // verified against packed slabs yet + if (_sg32 && Environment.GetEnvironmentVariable("SIMD_OCR_NOPACK") == null) + PackSlabs(recs, arena, off, slabElems, ref im2colOff); + // +128*512: coopmat A-tile reads pad rows past M + long arenaElems = im2colOff + maxIm2col + 128 * 512 + (1L << 20); + return new MetalSchedule + { + Recs = recs.ToArray(), + OutElems = (int)numel[outIdx], InElems = numel[inIdx], ArenaElems = arenaElems, + Off = off, Alias = alias, Numel = numel, Shapes = shapes, + Im2colOff = im2colOff, + }; + } + + /// + /// Re-packs the one-slab-per-tensor arena by lifetime: a slab lives from + /// the first to the last rec that binds it, and slabs with disjoint + /// lifetimes share memory (every rec is followed by a barrier, so a later + /// writer never races an earlier reader). Rewrites the arena bindings, + /// and the im2col scratch offset in place. Leaves + /// everything untouched when a rec binds the arena anywhere but a slab + /// start or the scratch. Afterwards only slabs still live at the end of + /// the schedule keep their values for the debug readers. + /// + private static void PackSlabs(List recs, MtlBuffer arena, long[] off, + long[] slabElems, ref long im2colOff) + { + var slabAt = new Dictionary(); + for (int i = 0; i < off.Length; i++) + if (off[i] >= 0) slabAt[off[i]] = i; + var first = new int[off.Length]; + var last = new int[off.Length]; + Array.Fill(first, -1); + for (int r = 0; r < recs.Count; r++) + foreach (MetalBind b in recs[r].Binds) + { + if (!ReferenceEquals(b.Buf, arena)) continue; + long e = (long)(b.ByteOffset / 2); + if (e == im2colOff) continue; + if ((b.ByteOffset & 1) != 0 || !slabAt.TryGetValue(e, out int i)) return; + if (first[i] < 0) first[i] = r; + last[i] = r; + } + + var order = new List(); + for (int i = 0; i < off.Length; i++) + if (first[i] >= 0) order.Add(i); + order.Sort((a, b) => first[a] != first[b] ? first[a].CompareTo(first[b]) : a.CompareTo(b)); + var newOff = new long[off.Length]; + Array.Fill(newOff, -1L); + var free = new List<(long Off, long Len)>(); // sorted by offset, coalesced + var live = new List(); + long top = 0; + foreach (int i in order) + { + for (int k = live.Count - 1; k >= 0; k--) + { + int j = live[k]; + if (last[j] >= first[i]) continue; + live.RemoveAt(k); + Release(free, newOff[j], slabElems[j]); + } + long need = slabElems[i]; + int best = -1; + for (int k = 0; k < free.Count; k++) + if (free[k].Len >= need && (best < 0 || free[k].Len < free[best].Len)) best = k; + if (best >= 0) + { + newOff[i] = free[best].Off; + if (free[best].Len == need) free.RemoveAt(best); + else free[best] = (free[best].Off + need, free[best].Len - need); + } + else if (free.Count > 0 && free[^1].Off + free[^1].Len == top) + { + newOff[i] = free[^1].Off; + top = newOff[i] + need; + free.RemoveAt(free.Count - 1); + } + else + { + newOff[i] = top; + top += need; + } + live.Add(i); + } + + long newIm2col = (top + 63) / 64 * 64; + foreach (MetalRec rec in recs) + { + MetalBind[] bb = (MetalBind[])rec.Binds.Clone(); + for (int k = 0; k < bb.Length; k++) + { + if (!ReferenceEquals(bb[k].Buf, arena)) continue; + long e = (long)(bb[k].ByteOffset / 2); + bb[k] = new MetalBind(arena, (ulong)(e == im2colOff ? newIm2col : newOff[slabAt[e]]) * 2); + } + rec.Binds = bb; + } + for (int i = 0; i < off.Length; i++) + if (off[i] >= 0) off[i] = newOff[i]; + im2colOff = newIm2col; + + static void Release(List<(long Off, long Len)> free, long o, long len) + { + int k = 0; + while (k < free.Count && free[k].Off < o) k++; + free.Insert(k, (o, len)); + if (k + 1 < free.Count && free[k].Off + free[k].Len == free[k + 1].Off) + { + free[k] = (free[k].Off, free[k].Len + free[k + 1].Len); + free.RemoveAt(k + 1); + } + if (k > 0 && free[k - 1].Off + free[k - 1].Len == free[k].Off) + { + free[k - 1] = (free[k - 1].Off, free[k - 1].Len + free[k].Len); + free.RemoveAt(k); + } + } + } + + private static byte[] Pcu(params uint[] v) + { + byte[] b = new byte[v.Length * 4]; + for (int i = 0; i < v.Length; i++) + BinaryPrimitives.WriteUInt32LittleEndian(b.AsSpan(i * 4), v[i]); + return b; + } + + private static ushort U16(ReadOnlySpan p, int o) => BinaryPrimitives.ReadUInt16LittleEndian(p[o..]); + private static int I32(ReadOnlySpan p, int o) => BinaryPrimitives.ReadInt32LittleEndian(p[o..]); + private static uint U32(ReadOnlySpan p, int o) => BinaryPrimitives.ReadUInt32LittleEndian(p[o..]); + private static float F32(ReadOnlySpan p, int o) => BinaryPrimitives.ReadSingleLittleEndian(p[o..]); + // fp16 bits of alpha (lo16) packed for shaders that read aux/flags via unpackHalf2x16 + private static ushort Half2Bits(float v) => BitConverter.HalfToUInt16Bits((Half)v); + private static uint HsAux(ReadOnlySpan p) => + (uint)Half2Bits(F32(p, 4)) | ((uint)Half2Bits(F32(p, 8)) << 16); + +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalRunner.cs b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalRunner.cs new file mode 100644 index 0000000..6b3e330 --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalRunner.cs @@ -0,0 +1,550 @@ +using System.Diagnostics; +using Sdcb.SimdPaddleOCR.OnnxSharp; + +namespace Sdcb.SimdPaddleOCR.Backends.Metal; + +/// One storage-buffer binding: a fixed (shared, immutable) buffer or +/// a role sentinel resolved per session. +internal readonly record struct MetalBind(MtlBuffer Buf, ulong ByteOffset); + +internal sealed class MetalRec +{ + public required MtlPipeline Pipe; + public required MetalBind[] Binds; + public required byte[] Pc; + public uint Gx, Gy; + public uint Tx = 256; // threads per threadgroup (attn/layernorm: 64) + public string Tag = ""; +} + +/// +/// Compiled dispatch list for one (N,H,W,nodeLimit,outTensor): managed data +/// only. Arena/IO requirements are sizes, not handles, so a session growing +/// its arena never invalidates anything. +/// +internal sealed class MetalSchedule +{ + public readonly record struct Key(int N, int H, int W, int NodeLimit, int OutTensor); + + public required MetalRec[] Recs; + public int OutElems; + public long InElems, ArenaElems; + public long[] Off = []; + public int[] Alias = []; + public long[] Numel = []; + public int[][] Shapes = []; + public long Im2colOff; + public long LastTick; +} + +/// Streamed-run callback: units [first, first+count) of result +/// (unit i at offsets[i]) are complete. Return false to stop receiving +/// further batches. +internal delegate bool MetalUnitsReady(float[] result, int[] offsets, int first, int count); + +/// +/// Per-session Metal runtime over a shared : +/// owns the fp16 arena, fp32 input/output buffers and the partials scratch — +/// all Shared-mode (unified memory: the CPU writes/reads the same mapping, +/// no staging copies). Each run re-records one command buffer per wave from +/// the pure-managed schedule; dispatches go through a serial compute encoder +/// with a memory barrier after each rec (the schedule's correctness relies +/// on strict ordering just like the Vulkan per-rec barrier). +/// +internal sealed unsafe class MetalDetGraph : IDisposable +{ + private readonly MtlDevice _dev; + private readonly MetalGraphModel _model; + private MtlBuffer? _arena, _in, _out, _part; + private long _arenaElems, _inElems, _outElems; + private float* _inMap, _outMap; + private MetalSchedule? _last; + private bool _disposed; + + private static readonly bool s_prof = Environment.GetEnvironmentVariable("SIMD_OCR_GPU_PROF") == "1"; + private static readonly bool s_noBar = Environment.GetEnvironmentVariable("SIMD_OCR_GPU_NOBAR") == "1"; + private static readonly bool s_dbgTime = Environment.GetEnvironmentVariable("SIMD_OCR_GPU_TIME") == "1"; + private static readonly bool s_onlyHead = Environment.GetEnvironmentVariable("SIMD_OCR_GPU_ONLYHEAD") == "1"; + private static readonly int s_only = + int.TryParse(Environment.GetEnvironmentVariable("SIMD_OCR_GPU_ONLY"), out int o) ? o : -1; + private static readonly int s_trunc = + int.TryParse(Environment.GetEnvironmentVariable("SIMD_OCR_GPU_TRUNCATE"), out int t) ? t : -1; + + public MetalDetGraph(MtlDevice dev, CompiledModel compiled) + { + _dev = dev; + _model = MetalGraphModel.Acquire(dev, compiled); + } + + /// Compile (or fetch) the schedule for a shape without running it — + /// lets callers probe GPU support before committing to the GPU path. + public MetalSchedule Prepare(int[] inputShape, int nodeLimit = int.MaxValue, int outTensor = -1) + => _model.GetSchedule(inputShape, nodeLimit, outTensor); + + /// Run the graph: input fp32 NCHW [n,C,H,W] → fp32 [graph output]. + /// nodeLimit truncates the emit loop (nodes beyond it are not dispatched) and + /// outTensor overrides the readback tensor — used to stop the REC graph before + /// the vocab projection (activations readback) instead of the graph output. + /// + public ReadOnlySpan Run(int[] inputShape, ReadOnlySpan input, + int nodeLimit = int.MaxValue, int outTensor = -1) + { + ObjectDisposedException.ThrowIf(_disposed, this); + MetalSchedule s = _model.GetSchedule(inputShape, nodeLimit, outTensor); + return RunCore([s], input, out _).AsSpan(0, s.OutElems); + } + + /// + /// Runs several shapes of the same graph (inputs back to back in + /// , unit i = numel(shapes[i]) floats) and hands + /// the activations to in unit order as they + /// land. The units go out in a few submissions: the GPU starts on the + /// first while the rest are still being recorded, and the caller's CPU + /// tail (CTC ArgMax) for early units overlaps the GPU work of later ones. + /// Returns false when declined a batch; every + /// submission has completed by the time this returns or throws. + /// + public bool RunMany(IReadOnlyList shapes, ReadOnlySpan input, + int nodeLimit, int outTensor, MetalUnitsReady onReady) + { + ObjectDisposedException.ThrowIf(_disposed, this); + var sched = new MetalSchedule[shapes.Count]; + for (int i = 0; i < sched.Length; i++) + sched[i] = _model.GetSchedule(shapes[i], nodeLimit, outTensor); + bool accepted = true; + RunCore(sched, input, out _, (result, offsets, first, count) => + accepted = accepted && onReady(result, offsets, first, count)); + return accepted; + } + + // in/out regions of consecutive units start on 256-byte boundaries + private static long Align64(long floats) => (floats + 63) & ~63L; + + private float[] RunCore(MetalSchedule[] sched, ReadOnlySpan input, + out int[] outOffsets, MetalUnitsReady? sink = null) + { + long t0 = Stopwatch.GetTimestamp(); + int n = sched.Length; + var inBase = new long[n]; + outOffsets = new int[n]; + var outBase = new long[n]; + long inTotal = 0, outTotal = 0, arena = 0; + int outElems = 0; + for (int i = 0; i < n; i++) + { + inBase[i] = inTotal; + inTotal = Align64(inTotal + sched[i].InElems); + outBase[i] = outTotal; + outTotal = Align64(outTotal + sched[i].OutElems); + outOffsets[i] = outElems; + outElems += sched[i].OutElems; + arena = Math.Max(arena, sched[i].ArenaElems); + } + if (n == 1) inTotal = Math.Max(inTotal, input.Length); + // Level-interleaved multi-unit recording (one barrier per level) plus + // wave cuts for streaming — same shape as the Vulkan runner. + bool interleave = n > 1 && !s_prof && !s_onlyHead && s_only < 0 && s_trunc < 0 && !s_noBar; + var arenaBase = new long[n]; + var waveStart = new List { 0 }; + int maxWave = 1; + if (interleave) + { + int[] cuts = sink is null ? [] : StreamCuts(sched); + long a = 0, peak = 0; + for (int i = 0, ci = 0; i < n; i++) + { + long need = (sched[i].ArenaElems + 127) & ~127L; + bool cut = ci < cuts.Length && cuts[ci] == i; + if (cut) ci++; + if (i > 0 && (cut || a + need > InterleaveArenaBudget)) { waveStart.Add(i); a = 0; } + arenaBase[i] = a; + a += need; + peak = Math.Max(peak, a); + maxWave = Math.Max(maxWave, i - waveStart[^1] + 1); + } + arena = peak; + } + waveStart.Add(n); + int waves = waveStart.Count - 1; + EnsureBuffers(arena, inTotal, outTotal); + long consumed = 0; + fixed (float* src = input) + for (int i = 0; i < n; i++) + { + long len = n == 1 ? input.Length : sched[i].InElems; + Buffer.MemoryCopy(src + consumed, _inMap + inBase[i], len * 4, len * 4); + consumed += len; + } + long t1 = Stopwatch.GetTimestamp(); + + EnsurePart(interleave ? maxWave : 1); + if (_result.Length < outElems) _result = new float[RoundUpPow2(outElems)]; + float[] result = _result; + long t2 = 0, t3; + var done = new List(waves); + try + { + for (int w = 0; w < waves; w++) + { + int s0 = waveStart[w], s1 = waveStart[w + 1]; + IntPtr cmd; + using (var pool = AutoReleasePool.Create()) + { + // Serial encoder: dispatches run in order; one device-scope + // barrier per level/rec mirrors the Vulkan pipeline barrier. + MtlEncoder enc = _dev.BeginCommands(); + cmd = enc.Cmd; + // +1 on the command buffer so it survives the pool drain + // while we wait below. + ObjC.Retain(cmd); + MtlPipeline? bound = null; + if (interleave) + { + int levels = 0; + for (int i = s0; i < s1; i++) levels = Math.Max(levels, sched[i].Recs.Length); + for (int k = 0; k < levels; k++) + { + for (int i = s0; i < s1; i++) + if (k < sched[i].Recs.Length) + RecordOne(enc, sched[i].Recs[k], ref bound, + (ulong)inBase[i] * 4, (ulong)outBase[i] * 4, + (ulong)arenaBase[i] * 2, + (ulong)(i - s0) * MetalGraphModel.PartBytes); + if (!s_noBar) enc.BarrierBuffers(); + } + } + else + for (int i = s0; i < s1; i++) + RecordRecs(enc, Selected(sched[i]), ref bound, + (ulong)inBase[i] * 4, (ulong)outBase[i] * 4); + enc.End(); + enc.Commit(); + } + done.Add(cmd); + } + t2 = Stopwatch.GetTimestamp(); + bool more = true; + for (int w = 0; w < waves; w++) + { + IntPtr cmd = done[w]; + ObjC.Send0(cmd, ObjC.S.WaitUntilCompleted); + long status = (long)ObjC.SendLong(cmd, ObjC.S.Status); + if (status != 4) + { + IntPtr err = ObjC.Send0(cmd, ObjC.S.Error); + throw new InvalidOperationException( + $"Metal command buffer failed (status={status}): {MtlDevice.DescribeError(err)}"); + } + ObjC.Release(cmd); + done[w] = IntPtr.Zero; + int s0 = waveStart[w], s1 = waveStart[w + 1]; + for (int i = s0; i < s1; i++) + new ReadOnlySpan(_outMap + outBase[i], sched[i].OutElems) + .CopyTo(result.AsSpan(outOffsets[i])); + long tsk = s_dbgTime ? Stopwatch.GetTimestamp() : 0; + if (sink is not null && more) more = sink(result, outOffsets, s0, s1 - s0); + if (s_dbgTime) _dbgSink += Stopwatch.GetTimestamp() - tsk; + } + } + finally + { + // never leave a submission in flight over buffers the next run reuses + foreach (IntPtr cmd in done) + if (cmd != IntPtr.Zero) + { + ObjC.Send0(cmd, ObjC.S.WaitUntilCompleted); + ObjC.Release(cmd); + } + } + t3 = Stopwatch.GetTimestamp(); + _last = sched[n - 1]; + if (s_dbgTime) + { + double f = Stopwatch.Frequency / 1e3; + Console.WriteLine($"[t] units={n}{(interleave ? " il" : "")} waves={waves} write={(t1 - t0) / f:F2} record+submit={(t2 - t1) / f:F2} " + + $"wait+read+sink={(t3 - t2) / f:F2} sink={_dbgSink / f:F2} ms recs={sched.Sum(x => x.Recs.Length)}"); + _dbgSink = 0; + } + return result; + } + + // Unit indices where a streamed run opens a new submission: ~55% of the + // input volume in the first wave (GPU starts while the rest records; the + // caller's tail for it overlaps the remainder), a third wave for large + // images so the last, un-overlapped tail stays short. + private static int[] StreamCuts(MetalSchedule[] sched) + { + int n = sched.Length; + if (n < 4) return []; + double[] fractions = n >= 9 ? [0.45, 0.78] : [0.55]; + long total = 0; + foreach (MetalSchedule s in sched) total += s.InElems; + var cuts = new List(); + long acc = 0; + int f = 0; + for (int i = 0; i < n - 1 && f < fractions.Length; i++) + { + acc += sched[i].InElems; + if (acc >= fractions[f] * total) { cuts.Add(i + 1); f++; } + } + return [.. cuts]; + } + + // fp16 elements of concurrently live unit arenas in one interleaved wave (512 MB) + private const long InterleaveArenaBudget = 256L << 20; + + private float[] _result = []; + private long _dbgSink; // SIMD_OCR_GPU_TIME: CPU time spent in the streamed sink + + private static long RoundUpPow2(long v) + { + long p = 1024; + while (p < v) p <<= 1; + return p; + } + + private int _partUnits; + + private void EnsurePart(int units) + { + if (_part is not null && units <= _partUnits) return; + int alloc = Math.Max(units, _partUnits * 3 / 2); + _part?.Dispose(); + _part = null; + // newBuffer is zero-initialised on Metal: se_fused ticket counters + // start at 0 and self-reset after each run. + _part = _dev.NewBuffer((nuint)alloc * MetalGraphModel.PartBytes); + _partUnits = alloc; + } + + private void EnsureBuffers(long arenaElems, long needIn, long needOut) + { + if (_arena is null || arenaElems > _arenaElems) + { + long alloc = Math.Max(arenaElems, _arenaElems * 3 / 2); + _arena?.Dispose(); + _arena = null; + _arena = _dev.NewBuffer((nuint)alloc * 2); + _arenaElems = alloc; + } + // grow in power-of-two steps so a stream of varying DET sizes settles + if (_in is null || needIn > _inElems) + { + long alloc = RoundUpPow2(needIn); + _in?.Dispose(); + _in = _dev.NewBuffer((nuint)alloc * 4); + _inMap = (float*)_in.Contents; + _inElems = alloc; + } + if (_out is null || needOut > _outElems) + { + long alloc = RoundUpPow2(needOut); + _out?.Dispose(); + _out = _dev.NewBuffer((nuint)alloc * 4); + _outMap = (float*)_out.Contents; + _outElems = alloc; + } + } + + private MtlBuffer Resolve(MtlBuffer b) => + ReferenceEquals(b, MetalGraphModel.RoleArena) ? _arena! + : ReferenceEquals(b, MetalGraphModel.RoleIn) ? _in! + : ReferenceEquals(b, MetalGraphModel.RoleOut) ? _out! + : ReferenceEquals(b, MetalGraphModel.RolePart) ? _part! + : b; + + private IReadOnlyList Selected(MetalSchedule s) + { + MetalRec[] r = s.Recs; + if (s_onlyHead) return [r[0]]; + if (s_only >= 0 && s_only < r.Length) return Enumerable.Repeat(r[s_only], 20).ToArray(); + if (s_trunc >= 0 && s_trunc < r.Length) return r.Take(s_trunc).ToArray(); + return r; + } + + private void RecordRecs(MtlEncoder enc, IReadOnlyList recs, + ref MtlPipeline? bound, ulong inBase, ulong outBase) + { + foreach (MetalRec r in recs) + { + RecordOne(enc, r, ref bound, inBase, outBase, 0, 0); + if (!s_noBar) enc.BarrierBuffers(); + } + } + + private void RecordOne(MtlEncoder enc, MetalRec r, ref MtlPipeline? bound, + ulong inBase, ulong outBase, ulong arenaBase, ulong partBase) + { + if (!ReferenceEquals(bound, r.Pipe)) + { + enc.SetPipeline(r.Pipe); + bound = r.Pipe; + } + for (int b = 0; b < r.Binds.Length; b++) + { + MtlBuffer role = r.Binds[b].Buf; + MtlBuffer buf = Resolve(role); + ulong off = r.Binds[b].ByteOffset + + (ReferenceEquals(role, MetalGraphModel.RoleIn) ? inBase + : ReferenceEquals(role, MetalGraphModel.RoleOut) ? outBase + : ReferenceEquals(role, MetalGraphModel.RoleArena) ? arenaBase + : ReferenceEquals(role, MetalGraphModel.RolePart) ? partBase : 0); + enc.SetBuffer(buf, (nuint)off, b); + } + fixed (byte* pp = r.Pc) + enc.SetPc(pp, (nuint)r.Pc.Length, r.Binds.Length); + enc.Dispatch(r.Gx, r.Gy, 1, r.Tx, 1, 1); + } + + // ---------------- debug readers ---------------- + + private static int PhysOf(MetalSchedule s, int tensorIndex) + { + int t = tensorIndex; + while (s.Alias[t] != t) t = s.Alias[t]; + return t; + } + + /// SIMD_OCR_GPU_DUMP: read a tensor's fp16 arena content as fp32. + public float[]? ReadTensor(int tensorIndex) + { + MetalSchedule? s = _last; + if (s is null || _arena is null) return null; + int t = PhysOf(s, tensorIndex); + if (t < 0 || t >= s.Numel.Length) return null; + long n = s.Numel[tensorIndex]; + long e = s.Off[t]; + if (e < 0) return new float[n]; + var v = new float[n]; + Half* p = (Half*)_arena.Contents + e; + for (long i = 0; i < n; i++) v[i] = (float)p[i]; + return v; + } + + private MetalSchedule Sched(int[] inputShape, int nodeLimit, int outTensor) + => _model.GetSchedule(inputShape, nodeLimit, outTensor); + + /// Debug: raw values of a tensor's arena slot (NHWC fp16 → fp32), + /// valid after a Run of the matching schedule. + public float[] DebugValues(int tensorIndex, int[] inputShape, int count = 16, + int nodeLimit = int.MaxValue, int outTensor = -1) + { + MetalSchedule s = Sched(inputShape, nodeLimit, outTensor); + int t = PhysOf(s, tensorIndex); + int n = Math.Min(count, (int)s.Numel[tensorIndex]); + if (s.Off[t] < 0 || _arena is null) return new float[n]; + var v = new float[n]; + Half* p = (Half*)_arena.Contents + s.Off[t]; + for (int i = 0; i < n; i++) v[i] = (float)p[i]; + return v; + } + + public float[] DebugValuesRaw(int[] inputShape, long elemOff, int count, + int nodeLimit = int.MaxValue, int outTensor = -1) + { + Sched(inputShape, nodeLimit, outTensor); + var v = new float[count]; + if (_arena is not null) + { + Half* p = (Half*)_arena.Contents + elemOff; + for (int i = 0; i < count; i++) v[i] = (float)p[i]; + } + return v; + } + + /// Debug: min/max/mean of a tensor's arena slot after a Run. + public (float Min, float Max, double Mean) DebugStats(int tensorIndex, int[] inputShape, + int nodeLimit = int.MaxValue, int outTensor = -1) + { + MetalSchedule s = Sched(inputShape, nodeLimit, outTensor); + int t = PhysOf(s, tensorIndex); + int n = (int)s.Numel[tensorIndex]; + var v = new float[n]; + if (s.Off[t] >= 0 && _arena is not null) + { + Half* p = (Half*)_arena.Contents + s.Off[t]; + for (int i = 0; i < n; i++) v[i] = (float)p[i]; + } + float mn = float.MaxValue, mx = float.MinValue; double sum = 0; + foreach (float x in v) { mn = MathF.Min(mn, x); mx = MathF.Max(mx, x); sum += x; } + return (mn, mx, sum / n); + } + + internal float[]? DbgConstKm(int tensorIndex, int n) => _model.DbgConstKm(tensorIndex, n); + internal int PhysProbe(int tensorIndex) + { + var s = _last; if (s is null) return -2; + return PhysOf(s, tensorIndex); + } + internal long SlotProbe(int tensorIndex) + { + var s = _last; if (s is null) return -2; + int t = PhysOf(s, tensorIndex); + return t < 0 ? -3 : s.Off[t]; + } + internal unsafe float[]? DbgPart(long elemOff, int n) // fp32 partials/scale region + { + if (_part is null) return null; + var sp = new ReadOnlySpan((float*)_part.Contents + elemOff, n); + return sp.ToArray(); + } + internal float[]? DbgIn(int n) + { + if (_inMap == null) return null; + return new ReadOnlySpan(_inMap, n).ToArray(); + } + + public long Im2colOffset(int[] inputShape) + => Sched(inputShape, int.MaxValue, -1).Im2colOff; + + public int DispatchCount(int[] inputShape, int nodeLimit = int.MaxValue, int outTensor = -1) + => Sched(inputShape, nodeLimit, outTensor).Recs.Length; + + /// Per-dispatch wall timing: serializes+times each rec when + /// SIMD_OCR_GPU_PROF=1; otherwise just lists the dispatch roster. + public void DumpProfile(int[] inputShape, int nodeLimit = int.MaxValue, int outTensor = -1) + { + MetalSchedule s = Sched(inputShape, nodeLimit, outTensor); + IReadOnlyList recs = Selected(s); + if (s_prof) + { + // Serialized per-rec commits — crude but correct wall clock. + var agg = new Dictionary(); + float[] inp = new float[s.InElems]; + EnsureBuffers(s.ArenaElems, Align64(s.InElems), Align64(s.OutElems)); + EnsurePart(1); + for (int i = 0; i < recs.Count; i++) + { + long t0 = Stopwatch.GetTimestamp(); + using (var pool = AutoReleasePool.Create()) + { + MtlEncoder enc = _dev.BeginCommands(); + MtlPipeline? bound = null; + RecordOne(enc, recs[i], ref bound, 0, 0, 0, 0); + enc.End(); enc.CommitAndWait(); + } + double ms = (Stopwatch.GetTimestamp() - t0) / (Stopwatch.Frequency / 1e3); + string tag = recs[i].Pipe.Name; + var a = agg.TryGetValue(tag, out var v) ? v : (ms: 0.0, n: 0); + agg[tag] = (a.ms + ms, a.n + 1); + } + Console.WriteLine($"--- Metal profile ({recs.Count} dispatches, serialized) ---"); + foreach (var kv in agg.OrderByDescending(k => k.Value.ms)) + Console.WriteLine($" {kv.Key,-16} {kv.Value.ms,9:F3} ms x{kv.Value.n}"); + return; + } + Console.WriteLine($"--- Metal dispatches ({recs.Count}) ---"); + foreach (MetalRec r in recs) + Console.WriteLine($" {r.Tag,-40} {r.Pipe.Name} gx={r.Gx} gy={r.Gy} tx={r.Tx}"); + } + + public void Dispose() + { + if (_disposed) return; + _disposed = true; + _arena?.Dispose(); _in?.Dispose(); _out?.Dispose(); _part?.Dispose(); + _arena = _in = _out = _part = null; + _inMap = _outMap = null; + _model.Release(); + } +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalSession.cs b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalSession.cs new file mode 100644 index 0000000..b0e3f29 --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalSession.cs @@ -0,0 +1,324 @@ +using System.Buffers.Binary; +using System.Runtime.InteropServices; +using Sdcb.SimdPaddleOCR.OnnxSharp; + +namespace Sdcb.SimdPaddleOCR.Backends.Metal; + +/// +/// GPU-backed over Metal: the whole ONNX graph runs +/// via . InputData is a managed staging span +/// uploaded on each run; the CTC tail matches 's +/// contract (stop before the vocab MatMul, hand operands to the shared CPU +/// ArgMax). One MtlDevice is shared process-wide; weights/pipelines/schedules +/// are shared per model, while each session owns its buffers, so sessions +/// only serialize at the command-buffer commits. +/// +internal sealed class MetalSession : IOcrSession, IBatchedCtcSession +{ + private readonly MtlDevice _dev; + private readonly MetalDetGraph _graph; + private readonly CompiledModel _compiled; + private readonly Model _model; + private readonly ResizeWorkspace _resizeWorkspace = new(); + private InferenceSession? _cpuFallback; + private float[] _cpuNhwc = []; // NCHW→NHWC staging for the CPU fallback + + private int[] _shape = []; + private float[] _input = []; + private int _curVolume; // numel of the current shape; _input may be bigger + private int _hwVolume; + private bool _disposed; + private bool _gpuDead; // plan/run failure → serve everything from _cpuFallback + + // CTC projection resolution, cached per Reshape + private bool _ctcResolved; + private int _ctcMatMulIndex = -1; + private int _ctcActTensor; + private int _ctcBatch, _ctcRows, _ctcInner, _ctcColumns; + private byte[]? _ctcWBytes, _ctcBiasBytes; + private float[]? _ctcPacked; + + internal MetalSession(MtlDevice dev, CompiledModel compiled) + { + _dev = dev; + _compiled = compiled; + _model = compiled.Model; + _graph = new MetalDetGraph(dev, compiled); + } + + public TensorShape InputShape => new(_shape); + private int[]? _outputShape; + public TensorShape OutputShape => + new(_outputShape ??= _compiled.ResolveShapesFor(_shape)[checked((int)_model.GraphOutputs[0])]); + public Span InputData => _input.AsSpan(0, _curVolume); + public ResizeWorkspace ResizeWorkspace => _resizeWorkspace; + public bool InputIsNhwc => false; // convk_f32n reads NCHW fp32 directly + public int IntraOpThreads { get; set; } + public bool GpuAlive => !_gpuDead; + public int HighWaterInputVolume => _hwVolume; + public bool PlanForCtcProjection { get; set; } + public bool IsProfilingEnabled => InferenceSession.ProfilingEnabled; + + public void Reshape(ReadOnlySpan inputShape) + { + if (_disposed) throw new ObjectDisposedException(nameof(MetalSession)); + _shape = inputShape.ToArray(); + long vol = 1; + foreach (int d in _shape) vol *= Math.Max(d, 1); + _curVolume = checked((int)vol); + if (vol > _input.Length) + { + _input = new float[vol]; + _hwVolume = checked((int)Math.Min(vol, int.MaxValue)); + } + _outputShape = null; + _ctcResolved = false; + _ctcMatMulIndex = -1; + } + + public ReadOnlySpan RunInternal(ReadOnlySpan input) + { + if (_disposed) throw new ObjectDisposedException(nameof(MetalSession)); + if (_gpuDead) return Cpu().RunInternal(CpuInput(input)); + try + { + return _graph.Run(_shape, input); + } + catch (Exception ex) + { + Console.Error.WriteLine($"[metal] RunInternal fallback: {ex.GetType().Name} {ex.Message}"); + _gpuDead = true; + return Cpu().RunInternal(CpuInput(input)); + } + } + + public bool TryRunUntilCtcProjection(ReadOnlySpan input, out CtcProjectionOperands operands) + { + operands = default; + if (_disposed) throw new ObjectDisposedException(nameof(MetalSession)); + if (_gpuDead) return Cpu().TryRunUntilCtcProjection(CpuInput(input), out operands); + ResolveCtcProjection(); + if (_ctcMatMulIndex < 0) return false; + + ReadOnlySpan act; + try + { + act = _graph.Run(_shape, input, _ctcMatMulIndex, _ctcActTensor); + } + catch (Exception ex) + { + Console.Error.WriteLine($"[metal] CtcProj fallback: {ex.GetType().Name} {ex.Message}"); + _gpuDead = true; + return Cpu().TryRunUntilCtcProjection(CpuInput(input), out operands); + } + operands = new CtcProjectionOperands(act, + MemoryMarshal.Cast(_ctcWBytes), + _ctcBiasBytes is null ? [] : MemoryMarshal.Cast(_ctcBiasBytes), + _ctcPacked, _ctcBatch, _ctcRows, _ctcInner, _ctcColumns, _ctcMatMulIndex); + return true; + } + + // ---- IBatchedCtcSession: several REC shapes, one GPU submission ---- + private int[][] _many = []; + private int _manyVolume; + + public bool CanRunMany => !_gpuDead && !_disposed; + + public Span ReshapeMany(IReadOnlyList shapes) + { + if (_disposed) throw new ObjectDisposedException(nameof(MetalSession)); + _many = [.. shapes]; + long vol = 0; + foreach (int[] s in _many) + { + long v = 1; + foreach (int d in s) v *= Math.Max(d, 1); + vol += v; + } + _manyVolume = checked((int)vol); + if (vol > _input.Length) + { + _input = new float[vol]; + _hwVolume = _manyVolume; + } + // single-shape members (CTC resolution, fallback) key off the first unit + Reshape(_many[0]); + return _input.AsSpan(0, _manyVolume); + } + + public bool TryResolveManyHead(out int[] rows, out CtcHead head) + { + rows = []; head = null!; + if (!CanRunMany || _many.Length == 0) return false; + ResolveCtcProjection(); + if (_ctcMatMulIndex < 0) return false; + rows = new int[_many.Length]; + for (int i = 0; i < _many.Length; i++) + rows[i] = _compiled.ResolveShapesFor(_many[i])[_ctcActTensor][^2]; + head = new CtcHead + { + Weights = _ctcWBytes!, Bias = _ctcBiasBytes, Packed = _ctcPacked, + Inner = _ctcInner, Columns = _ctcColumns, MatMulIndex = _ctcMatMulIndex, + }; + return true; + } + + public bool TryRunManyUntilCtcProjection(CtcUnitsReady onReady) + { + if (!CanRunMany || _many.Length == 0) return false; + ResolveCtcProjection(); + if (_ctcMatMulIndex < 0) return false; + // consumer failures are the caller's, not a GPU fault: keep them out + // of the fallback catch below and rethrow once the device is drained + System.Runtime.ExceptionServices.ExceptionDispatchInfo? consumerError = null; + bool Forward(float[] acts, int[] offsets, int first, int count) + { + try { return onReady(acts, offsets, first, count); } + catch (Exception ex) + { + consumerError = System.Runtime.ExceptionServices.ExceptionDispatchInfo.Capture(ex); + return false; + } + } + bool ok; + try + { + ok = _graph.RunMany(_many, _input.AsSpan(0, _manyVolume), _ctcMatMulIndex, _ctcActTensor, Forward); + } + catch (Exception ex) + { + Console.Error.WriteLine($"[metal] RunMany fallback: {ex.GetType().Name} {ex.Message}"); + _gpuDead = true; + return false; + } + consumerError?.Throw(); + return ok; + } + + private InferenceSession Cpu() + { + _cpuFallback ??= new InferenceSession(_compiled) + { PlanForCtcProjection = PlanForCtcProjection }; + _cpuFallback.Reshape(_shape); + if (IntraOpThreads > 0) _cpuFallback.IntraOpThreads = IntraOpThreads; + return _cpuFallback; + } + + /// Graphs without the CTC tail fall back to a lazy CPU session. + public ReadOnlySpan RunInternalSkipFinalSoftmax(ReadOnlySpan input, out bool outputIsLogits) + { + if (_disposed) throw new ObjectDisposedException(nameof(MetalSession)); + return Cpu().RunInternalSkipFinalSoftmax(CpuInput(input), out outputIsLogits); + } + + /// + /// Callers always write logical NCHW (this session reports + /// = false). When the compiled model flags its + /// graph input as NHWC-physical, the CPU fallback must be fed the + /// transposed layout or it decodes garbage. + /// + private ReadOnlySpan CpuInput(ReadOnlySpan input) + { + if (!_compiled.InputIsNhwc) return input; + if (_cpuNhwc.Length < input.Length) _cpuNhwc = new float[input.Length]; + Span dest = _cpuNhwc.AsSpan(0, input.Length); + Kernels.Nhwc.NchwToNhwc(input, dest, _shape[0], _shape[1], + checked(_shape[2] * _shape[3]), Math.Max(IntraOpThreads, 1)); + return dest; + } + + public void NoteProfile(OperatorId operation, long started, int nodeIndex) + => InferenceSession.NoteProfileStatic(operation, started, nodeIndex); + + /// + /// Mirrors InferenceSession.TryResolveCtcProjection but resolves + /// shapes from instead of live + /// tensor bindings — the GPU session has no CPU workspace. + /// + private void ResolveCtcProjection() + { + if (_ctcResolved) return; + _ctcResolved = true; + _ctcMatMulIndex = -1; + + NodeRecord[] nodes = _model.Nodes; + if (nodes.Length < 2) return; + int[][] shapes = _compiled.ResolveShapesFor(_shape); + int outIdx = checked((int)_model.GraphOutputs[0]); + + NodeRecord softmax = nodes[^1]; + if (softmax.Operator != OperatorId.Softmax || + softmax.Inputs.Length != 1 || softmax.Outputs.Length != 1 || + softmax.Outputs[0] != (uint)outIdx) return; + int axis = I32(_model.GetParameters(softmax), 4); + int[] smInShape = shapes[softmax.Inputs[0]]; + if (axis < 0) axis += smInShape.Length; + if (axis != smInShape.Length - 1) return; + + int terminalIndex = nodes.Length - 2; + NodeRecord terminal = nodes[terminalIndex]; + int matMulIndex; + byte[]? biasBytes = null; + if (terminal.Operator == OperatorId.MatMul && + terminal.Outputs[0] == softmax.Inputs[0]) + { + matMulIndex = terminalIndex; + } + else if (terminal.Operator == OperatorId.Add && terminal.Inputs.Length == 2 && + terminalIndex > 0 && terminal.Outputs[0] == softmax.Inputs[0]) + { + matMulIndex = terminalIndex - 1; + NodeRecord mm = nodes[matMulIndex]; + if (mm.Operator != OperatorId.MatMul || mm.Outputs.Length != 1) return; + uint biasIndex; + if (terminal.Inputs[0] == mm.Outputs[0]) biasIndex = terminal.Inputs[1]; + else if (terminal.Inputs[1] == mm.Outputs[0]) biasIndex = terminal.Inputs[0]; + else return; + var biasMeta = _compiled.GetTensor(checked((int)biasIndex)); + int[] projShape = shapes[mm.Outputs[0]]; + if (!biasMeta.IsConstant || biasMeta.Shape.Length != 1 || + projShape.Length < 2 || biasMeta.Shape[0] != projShape[^1]) return; + biasBytes = biasMeta.Constant; + } + else return; + + NodeRecord matMul = nodes[matMulIndex]; + if (matMul.Inputs.Length != 2 || matMul.Outputs.Length != 1) return; + var aMeta = _compiled.GetTensor(checked((int)matMul.Inputs[0])); + var bMeta = _compiled.GetTensor(checked((int)matMul.Inputs[1])); + int[] aShape = shapes[matMul.Inputs[0]]; + int[] bShape = shapes[matMul.Inputs[1]]; + int[] projOut = shapes[matMul.Outputs[0]]; + if (!bMeta.IsConstant || aShape.Length < 2 || bShape.Length != 2 || + projOut.Length < 2) return; + int rows = aShape[^2], inner = aShape[^1], columns = bShape[1]; + long aCount = 1; + foreach (int d in aShape) aCount *= Math.Max(d, 1); + if (rows <= 0 || inner <= 0 || aCount % (rows * (long)inner) != 0) return; + int batch = checked((int)(aCount / (rows * (long)inner))); + if (bShape[0] != inner) return; + long projCount = 1; + foreach (int d in projOut) projCount *= Math.Max(d, 1); + if (projOut[^1] != columns || projCount != (long)batch * rows * columns) return; + _compiled.TryGetPackedMatMul(matMul.Inputs[1], out float[]? packed); + // Column count does not gate here (unlike the recognizer's ArgMax + // fusion): small-column heads like the cls [inner,2] tail resolve too. + if (columns <= 0) return; + + _ctcMatMulIndex = matMulIndex; + _ctcActTensor = checked((int)matMul.Inputs[0]); + _ctcBatch = batch; _ctcRows = rows; _ctcInner = inner; _ctcColumns = columns; + _ctcWBytes = bMeta.Constant; + _ctcBiasBytes = biasBytes; + _ctcPacked = packed; + } + + private static int I32(ReadOnlySpan p, int o) => BinaryPrimitives.ReadInt32LittleEndian(p[o..]); + + public void Dispose() + { + _disposed = true; + _graph.Dispose(); + _cpuFallback?.Dispose(); + } +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/ObjC.cs b/src/Sdcb.SimdPaddleOCR/Backends/Metal/ObjC.cs new file mode 100644 index 0000000..3876174 --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/ObjC.cs @@ -0,0 +1,229 @@ +using System.Runtime.InteropServices; + +namespace Sdcb.SimdPaddleOCR.Backends.Metal; + +// Minimal libobjc P/Invoke layer — pure net10.0 C#, no macOS workload. +// Ownership follows Cocoa rules: alloc/init/newXxx/copy return +1 (caller +// releases); everything else is autoreleased within the enclosing +// NSAutoreleasePool. Selectors and class handles are cached in static +// readonly fields — a dispatch already costs enough objc_msgSend calls that +// re-registering selector strings per call is wasted marshaling. +internal static unsafe partial class ObjC +{ + private const string LibObjC = "libobjc.dylib"; + private const string LibMetal = "/System/Library/Frameworks/Metal.framework/Metal"; + + [LibraryImport(LibObjC, EntryPoint = "objc_getClass")] + private static partial IntPtr GetClassRaw(IntPtr name); + + [LibraryImport(LibObjC, EntryPoint = "sel_registerName")] + private static partial IntPtr RegSelRaw(IntPtr name); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial IntPtr Send0(IntPtr recv, IntPtr sel); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial IntPtr Send1P(IntPtr recv, IntPtr sel, IntPtr a); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial IntPtr Send2P(IntPtr recv, IntPtr sel, IntPtr a, IntPtr b); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial IntPtr Send3P(IntPtr recv, IntPtr sel, IntPtr a, IntPtr b, IntPtr c); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial IntPtr Send1N(IntPtr recv, IntPtr sel, nuint a); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial IntPtr Send1P1N(IntPtr recv, IntPtr sel, IntPtr a, nuint b); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial IntPtr Send2P1N(IntPtr recv, IntPtr sel, IntPtr a, IntPtr b, nuint c); + + // newBufferWithBytes:length:options: / newBufferWithLength:options: + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial IntPtr Send1P2N(IntPtr recv, IntPtr sel, IntPtr a, nuint b, nuint c); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial IntPtr Send1N1N(IntPtr recv, IntPtr sel, nuint a, nuint b); + + // setBuffer:offset:atIndex: / setBytes:length:atIndex: + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial void SendV1P2N(IntPtr recv, IntPtr sel, IntPtr a, nuint b, nuint c); + + // dispatchThreadgroups:threadsPerThreadgroup: — MTLSize by value (24B) + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial void SendV2Size(IntPtr recv, IntPtr sel, MTLSize a, MTLSize b); + + // setThreadgroupMemoryLength:atIndex: + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial void SendV1N1N(IntPtr recv, IntPtr sel, nuint a, nuint b); + + // single-NSUInteger void calls (setLanguageVersion:, memoryBarrierWithScope:, …) + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial void SendV1N(IntPtr recv, IntPtr sel, nuint a); + + // single-BOOL void calls (setFastMathEnabled:…) — BOOL is 1 byte on arm64. + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial void SendV1B(IntPtr recv, IntPtr sel, byte a); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial byte SendBool(IntPtr recv, IntPtr sel); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial byte SendBool1P(IntPtr recv, IntPtr sel, IntPtr a); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial byte SendBool1N(IntPtr recv, IntPtr sel, nuint a); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial nuint SendNuint(IntPtr recv, IntPtr sel); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial double SendDouble(IntPtr recv, IntPtr sel); + + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial long SendLong(IntPtr recv, IntPtr sel); + + // supportsFamily: — (id,SEL,NSInteger) + [LibraryImport(LibObjC, EntryPoint = "objc_msgSend")] + public static partial byte SendBool1L(IntPtr recv, IntPtr sel, nint a); + + [LibraryImport(LibMetal, EntryPoint = "MTLCreateSystemDefaultDevice")] + public static partial IntPtr CreateSystemDefaultDevice(); + + [StructLayout(LayoutKind.Sequential)] + public readonly struct MTLSize + { + public readonly nuint X, Y, Z; + public MTLSize(nuint x, nuint y, nuint z) { X = x; Y = y; Z = z; } + } + + [StructLayout(LayoutKind.Sequential)] + public readonly struct NSRange + { + public readonly nuint Location, Length; + public NSRange(nuint location, nuint length) { Location = location; Length = length; } + } + + // ---- cached class/selector tables ------------------------------------ + // sel_registerName returns the same SEL for the same string, so register + // once. Classes likewise are process-wide stable. + + private static readonly IntPtr s_classNsString = Cls("NSString"); + private static readonly IntPtr s_classNsAutoreleasePool = Cls("NSAutoreleasePool"); + private static readonly IntPtr s_classMtlCompileOptions = Cls("MTLCompileOptions"); + private static readonly IntPtr s_classMtlFunctionConstantValues = Cls("MTLFunctionConstantValues"); + + public static IntPtr ClassNSString => s_classNsString; + public static IntPtr ClassNSAutoreleasePool => s_classNsAutoreleasePool; + public static IntPtr ClassMTLCompileOptions => s_classMtlCompileOptions; + public static IntPtr ClassMTLFunctionConstantValues => s_classMtlFunctionConstantValues; + + public static IntPtr Cls(string name) + { + IntPtr n = Marshal.StringToHGlobalAnsi(name); + try { return GetClassRaw(n); } finally { Marshal.FreeHGlobal(n); } + } + + public static IntPtr Sel(string name) + { + IntPtr n = Marshal.StringToHGlobalAnsi(name); + try { return RegSelRaw(n); } finally { Marshal.FreeHGlobal(n); } + } + + public static class S + { + public static readonly IntPtr Alloc = Sel("alloc"); + public static readonly IntPtr Init = Sel("init"); + public static readonly IntPtr Release = Sel("release"); + public static readonly IntPtr Retain = Sel("retain"); + public static readonly IntPtr Drain = Sel("drain"); + public static readonly IntPtr InitWithUTF8String = Sel("initWithUTF8String:"); + public static readonly IntPtr Utf8String = Sel("UTF8String"); + public static readonly IntPtr LocalizedDescription = Sel("localizedDescription"); + public static readonly IntPtr Length = Sel("length"); + + public static readonly IntPtr Name = Sel("name"); + public static readonly IntPtr RespondsToSelector = Sel("respondsToSelector:"); + public static readonly IntPtr SupportsFamily = Sel("supportsFamily:"); + public static readonly IntPtr MaxThreadgroupMemoryLength = Sel("maxThreadgroupMemoryLength"); + public static readonly IntPtr MaxThreadsPerThreadgroup = Sel("maxThreadsPerThreadgroup"); + public static readonly IntPtr MaxTotalThreadsPerThreadgroup = Sel("maxTotalThreadsPerThreadgroup"); + public static readonly IntPtr ThreadExecutionWidth = Sel("threadExecutionWidth"); + public static readonly IntPtr HasUnifiedMemory = Sel("hasUnifiedMemory"); + public static readonly IntPtr RecommendedMaxWorkingSetSize = Sel("recommendedMaxWorkingSetSize"); + public static readonly IntPtr CurrentAllocatedSize = Sel("currentAllocatedSize"); + public static readonly IntPtr RegistryID = Sel("registryID"); + + public static readonly IntPtr NewCommandQueue = Sel("newCommandQueue"); + public static readonly IntPtr NewBufferWithBytesLengthOptions = Sel("newBufferWithBytes:length:options:"); + public static readonly IntPtr NewBufferWithLengthOptions = Sel("newBufferWithLength:options:"); + public static readonly IntPtr Contents = Sel("contents"); + public static readonly IntPtr NewLibraryWithSourceOptionsError = Sel("newLibraryWithSource:options:error:"); + public static readonly IntPtr NewFunctionWithName = Sel("newFunctionWithName:"); + public static readonly IntPtr NewFunctionWithNameConstantValuesError = Sel("newFunctionWithName:constantValues:error:"); + public static readonly IntPtr NewComputePipelineStateWithFunctionError = Sel("newComputePipelineStateWithFunction:error:"); + public static readonly IntPtr FunctionName = Sel("functionName"); + public static readonly IntPtr FunctionConstantsDictionary = Sel("functionConstantsDictionary"); + public static readonly IntPtr SetConstantValueTypeAtIndex = Sel("setConstantValue:type:atIndex:"); + + public static readonly IntPtr SetFastMathEnabled = Sel("setFastMathEnabled:"); + public static readonly IntPtr SetMathMode = Sel("setMathMode:"); + public static readonly IntPtr SetLanguageVersion = Sel("setLanguageVersion:"); + + public static readonly IntPtr CommandBuffer = Sel("commandBuffer"); + public static readonly IntPtr ComputeCommandEncoder = Sel("computeCommandEncoder"); + public static readonly IntPtr ComputeCommandEncoderWithDispatchType = Sel("computeCommandEncoderWithDispatchType:"); + public static readonly IntPtr SetComputePipelineState = Sel("setComputePipelineState:"); + public static readonly IntPtr SetBufferOffsetAtIndex = Sel("setBuffer:offset:atIndex:"); + public static readonly IntPtr SetBytesLengthAtIndex = Sel("setBytes:length:atIndex:"); + public static readonly IntPtr DispatchThreadgroupsThreadsPerThreadgroup = Sel("dispatchThreadgroups:threadsPerThreadgroup:"); + public static readonly IntPtr DispatchThreadsThreadsPerThreadgroup = Sel("dispatchThreads:threadsPerThreadgroup:"); + public static readonly IntPtr SetThreadgroupMemoryLengthAtIndex = Sel("setThreadgroupMemoryLength:atIndex:"); + public static readonly IntPtr MemoryBarrierWithScope = Sel("memoryBarrierWithScope:"); + public static readonly IntPtr EndEncoding = Sel("endEncoding"); + public static readonly IntPtr Commit = Sel("commit"); + public static readonly IntPtr WaitUntilCompleted = Sel("waitUntilCompleted"); + public static readonly IntPtr Status = Sel("status"); + public static readonly IntPtr Error = Sel("error"); + public static readonly IntPtr GpuStartTime = Sel("GPUStartTime"); + public static readonly IntPtr GpuEndTime = Sel("GPUEndTime"); + public static readonly IntPtr Label = Sel("label"); + public static readonly IntPtr SetLabel = Sel("setLabel:"); + } + + // NSString* with UTF-8 — caller owns (+1 via alloc/init); pass owns:false + // to consume within the current autorelease pool. + public static IntPtr NsStr(string s) + { + IntPtr obj = Send0(s_classNsString, S.Alloc); + IntPtr cstr = Marshal.StringToHGlobalAnsi(s); + try { return Send1P(obj, S.InitWithUTF8String, cstr); } + finally { Marshal.FreeHGlobal(cstr); } + } + + public static string? NsToString(IntPtr nsString) + { + if (nsString == 0) return null; + IntPtr utf8 = Send0(nsString, S.Utf8String); + return utf8 == 0 ? null : Marshal.PtrToStringUTF8(utf8); + } + + public static void Release(IntPtr obj) { if (obj != 0) Send0(obj, S.Release); } + public static IntPtr Retain(IntPtr obj) => Send0(obj, S.Retain); +} + +/// Scoped NSAutoreleasePool — .NET worker threads have none, so every +/// public entry that allocates autoreleased ObjC objects wraps itself in one. +/// Drain at Dispose releases the pool itself. +internal readonly struct AutoReleasePool : IDisposable +{ + private readonly IntPtr _pool; + private AutoReleasePool(IntPtr pool) => _pool = pool; + + public static AutoReleasePool Create() => + new(ObjC.Send0(ObjC.Send0(ObjC.ClassNSAutoreleasePool, ObjC.S.Alloc), ObjC.S.Init)); + + public void Dispose() => ObjC.Send0(_pool, ObjC.S.Drain); +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/a_common.metal b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/a_common.metal new file mode 100644 index 0000000..d3b4180 --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/a_common.metal @@ -0,0 +1,42 @@ +// Shared helpers for the OCR kernel set — concatenated first. +#include +using namespace metal; + +// Abramowitz-Stegun 7.1.26 erf (same polynomial the GLSL kernels use — +// identical rounding keeps CPU↔GPU diffs minimal). +static float erf1(float v) +{ + float s = v < 0.0 ? -1.0 : 1.0; + v = fabs(v); + float t = 1.0 / (1.0 + 0.47047 * v); + float q = t * (0.3480242 + t * (-0.0958798 + t * 0.7478556)); + return s * (1.0 - q * exp(-v * v)); +} + +// fp16 alpha lo16 / fp16 beta hi16 packed pair (hardsigmoid/hardswish aux) +static inline float2 auxHalf2(uint aux) +{ + ushort2 u = as_type(aux); + return float2(as_type(u.x), as_type(u.y)); +} + +// Conv epilogue activation (flags bits4-6; hardswish alpha in bits16-31): +// 0=none 1=relu 2=gelu 3=hardswish(a in hi16) 4=sigmoid +static inline float actf(float v, uint act, uint flg) +{ + if (act == 1u) return max(v, 0.0f); + if (act == 2u) return 0.5f * v * (1.0f + erf1(v * 0.70710678118654752f)); + if (act == 3u) { + float a = float(as_type((ushort)(flg >> 16))); + return v * clamp(v * a + 0.5f, 0.0f, 1.0f); + } + if (act == 4u) return 1.0f / (1.0f + exp(-v)); + if (act == 5u) return v / (1.0f + exp(-v)); // swish x*sigmoid(x) + return v; +} + +static inline float4 actf4(float4 v, uint act, uint flg) +{ + return float4(actf(v.x, act, flg), actf(v.y, act, flg), + actf(v.z, act, flg), actf(v.w, act, flg)); +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/conv.metal b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/conv.metal new file mode 100644 index 0000000..937ed7a --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/conv.metal @@ -0,0 +1,358 @@ +// Port of {conv_dw, conv_dw4, conv_dw4t(+ABSORB), conv_dense, convt2s2, +// convt4, im2col}.comp — depthwise/dense/transpose convs on NHWC fp16. +// flags: bit0 bias, bit3 scalarBias(convt), bits4-6 act, bit8 absorb. +#include +using namespace metal; + +struct PcDw { + uint outW, outH, C, kH, kW, sh, sw, pt, pl, inW, inH, flags; +}; + +// ---------------- conv_dw ---------------- +// Depthwise conv, one thread per output element (scalar C). +kernel void conv_dw(device const half* x [[buffer(0)]], + device const half* w [[buffer(1)]], + device const half* bias [[buffer(2)]], + device half* o [[buffer(3)]], + constant PcDw& p [[buffer(4)]], + uint2 gid [[thread_position_in_grid]]) +{ + uint total = p.outW * p.outH * p.C; + uint g = gid.x; + if (g >= total) return; + uint c = g % p.C; + uint px = g / p.C; + uint oy = px / p.outW, ox = px % p.outW; + int iy0 = int(oy * p.sh) - int(p.pt); + int ix0 = int(ox * p.sw) - int(p.pl); + + float acc = 0.0f; + for (uint dy = 0u; dy < p.kH; dy++) { + int iy = iy0 + int(dy); + if (iy < 0 || iy >= int(p.inH)) continue; + for (uint dx = 0u; dx < p.kW; dx++) { + int ix = ix0 + int(dx); + if (ix < 0 || ix >= int(p.inW)) continue; + acc += float(w[c * p.kH * p.kW + dy * p.kW + dx]) * + float(x[(uint(iy) * p.inW + uint(ix)) * p.C + c]); + } + } + if ((p.flags & 1u) != 0u) acc += float(bias[c]); + o[g] = half(actf(acc, (p.flags >> 4) & 7u, p.flags)); +} + +// ---------------- conv_dw4 ---------------- +// 4 channels per thread (C%4==0); weights tap-major [tap*C + c] (vec4). +// flags bit8: input is unfused addps → x*se[c4]+ra[i] reads (conv_dw4a). +kernel void conv_dw4(device const half4* x [[buffer(0)]], + device const half4* w [[buffer(1)]], + device const half* bias [[buffer(2)]], + device half4* o [[buffer(3)]], + constant PcDw& p [[buffer(4)]], + uint2 gid [[thread_position_in_grid]]) +{ + uint cv = p.C >> 2; + uint total = p.outW * p.outH * cv; + if (gid.x >= total) return; + uint b = gid.y; + uint xb = b * p.inH * p.inW * cv; + uint c4 = gid.x % cv; + uint px = gid.x / cv; + uint oy = px / p.outW, ox = px % p.outW; + int iy0 = int(oy * p.sh) - int(p.pt); + int ix0 = int(ox * p.sw) - int(p.pl); + + float4 acc = float4(0.0); + for (uint dy = 0u; dy < p.kH; dy++) { + int iy = iy0 + int(dy); + if (iy < 0 || iy >= int(p.inH)) continue; + uint rowBase = xb + uint(iy) * p.inW * cv; + for (uint dx = 0u; dx < p.kW; dx++) { + int ix = ix0 + int(dx); + if (ix < 0 || ix >= int(p.inW)) continue; + uint tap = dy * p.kW + dx; + float4 wv = float4(w[tap * cv + c4]); + acc += wv * float4(x[rowBase + uint(ix) * cv + c4]); + } + } + if ((p.flags & 1u) != 0u) { + uint cb = c4 * 4u; + acc += float4(float(bias[cb]), float(bias[cb + 1u]), + float(bias[cb + 2u]), float(bias[cb + 3u])); + } + uint act = (p.flags >> 4) & 7u; + o[b * total + gid.x] = half4(actf4(acc, act, p.flags)); +} + +// ---------------- conv_dw4a ---------------- +// conv_dw4 + addps absorb (extra ra/se binds; bit8 always set on emit). +kernel void conv_dw4a(device const half4* x [[buffer(0)]], + device const half4* w [[buffer(1)]], + device const half* bias [[buffer(2)]], + device half4* o [[buffer(3)]], + device const half4* ra [[buffer(4)]], + device const half4* se [[buffer(5)]], + constant PcDw& p [[buffer(6)]], + uint2 gid [[thread_position_in_grid]]) +{ + uint cv = p.C >> 2; + uint total = p.outW * p.outH * cv; + if (gid.x >= total) return; + uint b = gid.y; + uint xb = b * p.inH * p.inW * cv; + uint c4 = gid.x % cv; + uint px = gid.x / cv; + uint oy = px / p.outW, ox = px % p.outW; + int iy0 = int(oy * p.sh) - int(p.pt); + int ix0 = int(ox * p.sw) - int(p.pl); + + float4 acc = float4(0.0); + for (uint dy = 0u; dy < p.kH; dy++) { + int iy = iy0 + int(dy); + if (iy < 0 || iy >= int(p.inH)) continue; + uint rowBase = xb + uint(iy) * p.inW * cv; + for (uint dx = 0u; dx < p.kW; dx++) { + int ix = ix0 + int(dx); + if (ix < 0 || ix >= int(p.inW)) continue; + uint tap = dy * p.kW + dx; + float4 wv = float4(w[tap * cv + c4]); + uint gi = rowBase + uint(ix) * cv + c4; + float4 xv = (p.flags & 256u) != 0u + ? float4(x[gi]) * float4(se[b * cv + c4]) + float4(ra[gi]) + : float4(x[gi]); + acc += wv * xv; + } + } + if ((p.flags & 1u) != 0u) { + uint cb = c4 * 4u; + acc += float4(float(bias[cb]), float(bias[cb + 1u]), + float(bias[cb + 2u]), float(bias[cb + 3u])); + } + uint act = (p.flags >> 4) & 7u; + o[b * total + gid.x] = half4(actf4(acc, act, p.flags)); +} + +// ---------------- conv_dw4t ---------------- +// Shared-tile depthwise (C%4==0, k<=9, s<=2): one 256-thread group per +// 16x16 output tile × 4 channels; the (16*s+k-1)² input halo is staged in +// threadgroup memory once. +// grid x = tilesX*tilesY*(C/4) groups; gy = batch. +kernel void conv_dw4t(device const half4* x [[buffer(0)]], + device const half4* w [[buffer(1)]], + device const half* bias [[buffer(2)]], + device half4* o [[buffer(3)]], + device const half4* ra [[buffer(4)]], + device const half4* se [[buffer(5)]], + constant PcDw& p [[buffer(6)]], + uint2 wg [[threadgroup_position_in_grid]], + uint2 tp2 [[thread_position_in_threadgroup]]) +{ + uint tid = tp2.x; + // worst tile: s=2,k=9 → (16*2+8)² = 40² half4 = 12.8KB + 81 weights + threadgroup half4 shTile[40 * 40]; + threadgroup half4 shW[81]; + uint cv = p.C >> 2; + uint b = wg.y; + uint xb = b * p.inH * p.inW * cv; + uint seb = b * cv; + uint tilesX = (p.outW + 15u) / 16u; + uint g = wg.x; + uint c4 = g % cv; + uint tile = g / cv; + uint tx = tile % tilesX, ty = tile / tilesX; + + int iyT = int(ty * 16u * p.sh) - int(p.pt); + int ixT = int(tx * 16u * p.sw) - int(p.pl); + uint stW = 16u * p.sw + p.kW - 1u; + uint stH = 16u * p.sh + p.kH - 1u; + uint tileElems = stW * stH; + + for (uint e = tid; e < tileElems; e += 256u) { + uint iy = e / stW, ix = e - iy * stW; + int gy = iyT + int(iy), gx2 = ixT + int(ix); + half4 v = half4(0.0); + if (gy >= 0 && uint(gy) < p.inH && gx2 >= 0 && uint(gx2) < p.inW) { + uint gi = xb + (uint(gy) * p.inW + uint(gx2)) * cv + c4; + v = x[gi]; + if ((p.flags & 256u) != 0u) + v = half4(float4(v) * float4(se[seb + c4]) + float4(ra[gi])); + } + shTile[e] = v; + } + uint taps = p.kH * p.kW; + for (uint e = tid; e < taps; e += 256u) + shW[e] = w[e * cv + c4]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint ox = tx * 16u + (tid & 15u); + uint oy = ty * 16u + (tid >> 4); + if (ox >= p.outW || oy >= p.outH) return; + + uint lx0 = (tid & 15u) * p.sw; + uint ly0 = (tid >> 4) * p.sh; + float4 acc = float4(0.0); + for (uint dy = 0u; dy < p.kH; dy++) + for (uint dx = 0u; dx < p.kW; dx++) + acc += float4(shW[dy * p.kW + dx]) * float4(shTile[(ly0 + dy) * stW + lx0 + dx]); + if ((p.flags & 1u) != 0u) { + uint cb = c4 * 4u; + acc += float4(float(bias[cb]), float(bias[cb + 1u]), + float(bias[cb + 2u]), float(bias[cb + 3u])); + } + uint act = (p.flags >> 4) & 7u; + o[b * p.outH * p.outW * cv + (oy * p.outW + ox) * cv + c4] = half4(actf4(acc, act, p.flags)); +} + +// ---------------- conv_dense ---------------- +// Small dense conv (any k/s/pad), one thread per output element. +// W OIHW flattened [Cout, Cin*kH*kW]. +struct PcDense { + uint outW, outH, Cout, Cin, kH, kW, sh, sw, pt, pl, inW, inH, flags; +}; +kernel void conv_dense(device const half* x [[buffer(0)]], + device const half* w [[buffer(1)]], + device const half* bias [[buffer(2)]], + device half* o [[buffer(3)]], + constant PcDense& p [[buffer(4)]], + uint gid [[thread_position_in_grid]]) +{ + uint total = p.outW * p.outH * p.Cout; + if (gid >= total) return; + uint co = gid % p.Cout; + uint px = gid / p.Cout; + uint oy = px / p.outW, ox = px % p.outW; + int iy0 = int(oy * p.sh) - int(p.pt); + int ix0 = int(ox * p.sw) - int(p.pl); + uint kSpan = p.kH * p.kW; + + float acc = 0.0f; + for (uint dy = 0u; dy < p.kH; dy++) { + int iy = iy0 + int(dy); + if (iy < 0 || iy >= int(p.inH)) continue; + for (uint dx = 0u; dx < p.kW; dx++) { + int ix = ix0 + int(dx); + if (ix < 0 || ix >= int(p.inW)) continue; + uint xbase = (uint(iy) * p.inW + uint(ix)) * p.Cin; + uint wbase = co * (p.Cin * kSpan) + dy * p.kW + dx; + for (uint ci = 0u; ci < p.Cin; ci++) + acc += float(w[wbase + ci * kSpan]) * float(x[xbase + ci]); + } + } + if ((p.flags & 1u) != 0u) acc += float(bias[co]); + o[gid] = half(actf(acc, (p.flags >> 4) & 7u, p.flags)); +} + +// ---------------- convt2s2 ---------------- +// ConvTranspose 2x2 s2 (PP-OCR FPN head): out px depends on exactly one input +// px (oy/2, ox/2) and the (oy%2, ox%2) tap. W [Cin, Cout, 2, 2] flattened. +// flags bit7: write fp32 graph output to `of` instead of fp16 `o`. +struct PcCt { uint outW, outH, Cout, Cin, inW, flags; }; +kernel void convt2s2(device const half* x [[buffer(0)]], + device const half* w [[buffer(1)]], + device const half* bias [[buffer(2)]], + device half* o [[buffer(3)]], + device float* of [[buffer(4)]], + constant PcCt& p [[buffer(5)]], + uint gid [[thread_position_in_grid]]) +{ + uint total = p.outW * p.outH * p.Cout; + if (gid >= total) return; + uint co = gid % p.Cout; + uint px = gid / p.Cout; + uint oy = px / p.outW, ox = px % p.outW; + uint iy = oy >> 1, ix = ox >> 1; + uint dy = oy & 1u, dx = ox & 1u; + + float acc = 0.0f; + uint xbase = (iy * p.inW + ix) * p.Cin; + uint tap = dy * 2u + dx; + for (uint ci = 0u; ci < p.Cin; ci++) + acc += float(w[(ci * p.Cout + co) * 4u + tap]) * float(x[xbase + ci]); + if ((p.flags & 1u) != 0u) acc += float(bias[(p.flags & 8u) != 0u ? 0u : co]); + float v = actf(acc, (p.flags >> 4) & 7u, p.flags); + if ((p.flags & 128u) != 0u) of[gid] = v; + else o[gid] = half(v); +} + +// ---------------- convt4 ---------------- +// vec4 convtranspose: weights repacked [tap][ci][co] fp16. +kernel void convt4(device const half4* x [[buffer(0)]], + device const half4* w [[buffer(1)]], + device const half* bias [[buffer(2)]], + device half4* o [[buffer(3)]], + constant PcCt& p [[buffer(4)]], + uint gid [[thread_position_in_grid]]) +{ + uint nv = p.Cout >> 2; + if (gid >= p.outW * p.outH * nv) return; + uint co4 = gid % nv; + uint px = gid / nv; + uint oy = px / p.outW, ox = px % p.outW; + uint iy = oy >> 1, ix = ox >> 1; + uint dy = oy & 1u, dx = ox & 1u; + + float4 acc = float4(0.0); + uint xb = (iy * p.inW + ix) * (p.Cin >> 2); + uint tap = dy * 2u + dx; + uint civ = p.Cin >> 2; + for (uint ci4 = 0u; ci4 < civ; ci4++) { + float4 xv = float4(x[xb + ci4]); + uint wbase = (tap * p.Cin + ci4 * 4u) * nv + co4; + for (uint i = 0u; i < 4u; i++) + acc += xv[i] * float4(w[wbase + i * nv]); + } + if ((p.flags & 1u) != 0u) { + if ((p.flags & 8u) != 0u) acc += float4(float(bias[0])); + else acc += float4(float(bias[co4 * 4u]), float(bias[co4 * 4u + 1u]), + float(bias[co4 * 4u + 2u]), float(bias[co4 * 4u + 3u])); + } + uint act = (p.flags >> 4) & 7u; + o[gid] = half4(actf4(acc, act, p.flags)); +} + +// ---------------- im2col ---------------- +// im2col gather, tap-major: k = (dy*kW+dx)*Cin + ci over [0,K); pad → 0. +// Each thread emits 4 consecutive k (Kp%4==0) → one half4 store. +// flags bit0: fully-scalar fallback when Cin%4 != 0. +struct PcIm2col { + uint outW, outH, Cin, kH, kW, sH, sW, pT, pL, inW, inH, K, Kp, flags; +}; +kernel void im2col(device const half4* x [[buffer(0)]], + device half4* a [[buffer(1)]], + constant PcIm2col& p [[buffer(2)]], + uint gid [[thread_position_in_grid]]) +{ + uint kv = p.Kp >> 2; + uint total = p.outW * p.outH * kv; + if (gid >= total) return; + uint k4 = (gid % kv) * 4u; + uint m = gid / kv; + uint oy = m / p.outW, ox = m % p.outW; + half4 v = half4(0.0); + if ((p.flags & 1u) == 0u) { + if (k4 < p.K) { + uint tap = k4 / p.Cin; + uint ci = k4 - tap * p.Cin; + uint dy = tap / p.kW, dx = tap % p.kW; + int iy = int(oy * p.sH + dy) - int(p.pT); + int ix = int(ox * p.sW + dx) - int(p.pL); + if (iy >= 0 && iy < int(p.inH) && ix >= 0 && ix < int(p.inW)) + v = x[(uint(iy) * p.inW + uint(ix)) * (p.Cin >> 2) + (ci >> 2)]; + } + } else { + for (uint i = 0u; i < 4u; i++) { + uint k = k4 + i; + if (k < p.K) { + uint tap = k / p.Cin, ci = k - tap * p.Cin; + uint dy = tap / p.kW, dx = tap % p.kW; + int iy = int(oy * p.sH + dy) - int(p.pT); + int ix = int(ox * p.sW + dx) - int(p.pL); + if (iy >= 0 && iy < int(p.inH) && ix >= 0 && ix < int(p.inW)) { + uint xb = (uint(iy) * p.inW + uint(ix)) * p.Cin + ci; + v[i] = x[xb >> 2][xb & 3u]; + } + } + } + } + a[gid] = v; +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/dot.metal b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/dot.metal new file mode 100644 index 0000000..3f72b2c --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/dot.metal @@ -0,0 +1,334 @@ +// Port of {conv1x1_dot, mmdot_sk, convk_dot, convk_f32n}.comp — direct-dot +// m4×n4 GEMM-ish kernels on fp16 data with k-major vec4 weights. +// flags: bit0 bias, bit1 residual, bit2 prescale(SE), bit3 dual-write +// (o2=gelu(o); bit10 1024 → o2=raw pre-act), bit8 apsorb, bits4-6 act. +#include +using namespace metal; + +struct PcDot { uint M, N, K, kv4, flags, scMod, mImg; }; + +// ---------------- conv1x1_dot ---------------- +// O[m0:m0+4, n0:n0+4] = act(bias + res + A[m,k]·W[k,n4]) — A rows vec4 strided. +kernel void conv1x1_dot(device const half4* x [[buffer(0)]], + device const half4* w [[buffer(1)]], + device const half4* bias [[buffer(2)]], + device const half4* res [[buffer(3)]], + device half4* o [[buffer(4)]], + device const half4* s [[buffer(5)]], + device half4* o2 [[buffer(6)]], + device const half4* ra [[buffer(7)]], + constant PcDot& p [[buffer(8)]], + uint gid [[thread_position_in_grid]]) +{ + uint nv = p.N >> 2; + uint mv = (p.M + 3u) >> 2; + if (gid >= mv * nv) return; + uint n4 = gid % nv; + uint m0 = (gid / nv) * 4u; + + float4 acc0 = float4(0.0), acc1 = float4(0.0), acc2 = float4(0.0), acc3 = float4(0.0); + bool preScale = (p.flags & 4u) != 0u; + bool apsorb = (p.flags & 512u) != 0u; + uint sb0 = 0u, sb1 = 0u, sb2 = 0u, sb3 = 0u; + if (p.mImg != 0u && (preScale || apsorb)) { + sb0 = (m0 / p.mImg) * p.scMod; + sb1 = ((m0 + 1u) / p.mImg) * p.scMod; + sb2 = ((m0 + 2u) / p.mImg) * p.scMod; + sb3 = ((m0 + 3u) / p.mImg) * p.scMod; + } + uint kv4t = (p.K + 3u) >> 2; + for (uint k4 = 0u; k4 < kv4t; k4++) { + float4 a0 = m0 < p.M ? float4(x[m0 * p.kv4 + k4]) : float4(0.0); + float4 a1 = m0 + 1u < p.M ? float4(x[(m0 + 1u) * p.kv4 + k4]) : float4(0.0); + float4 a2 = m0 + 2u < p.M ? float4(x[(m0 + 2u) * p.kv4 + k4]) : float4(0.0); + float4 a3 = m0 + 3u < p.M ? float4(x[(m0 + 3u) * p.kv4 + k4]) : float4(0.0); + if (apsorb) { + if (m0 < p.M) a0 = a0 * float4(s[sb0 + k4]) + float4(ra[m0 * p.kv4 + k4]); + if (m0 + 1u < p.M) a1 = a1 * float4(s[sb1 + k4]) + float4(ra[(m0 + 1u) * p.kv4 + k4]); + if (m0 + 2u < p.M) a2 = a2 * float4(s[sb2 + k4]) + float4(ra[(m0 + 2u) * p.kv4 + k4]); + if (m0 + 3u < p.M) a3 = a3 * float4(s[sb3 + k4]) + float4(ra[(m0 + 3u) * p.kv4 + k4]); + } + if (preScale) { + a0 *= float4(s[sb0 + k4 % p.scMod]); + a1 *= float4(s[sb1 + k4 % p.scMod]); + a2 *= float4(s[sb2 + k4 % p.scMod]); + a3 *= float4(s[sb3 + k4 % p.scMod]); + } + uint wRow = k4 * 4u * nv + n4; + for (uint i = 0u; i < 4u; i++) { + uint k = k4 * 4u + i; + if (k >= p.K) break; + float4 wv = float4(w[wRow + i * nv]); + acc0 += a0[i] * wv; acc1 += a1[i] * wv; + acc2 += a2[i] * wv; acc3 += a3[i] * wv; + } + } + if ((p.flags & 1u) != 0u) { + float4 bv = float4(bias[n4]); + acc0 += bv; acc1 += bv; acc2 += bv; acc3 += bv; + } + uint act = (p.flags >> 4) & 7u; + if ((p.flags & 2u) != 0u) { + if (m0 < p.M) acc0 += float4(res[(m0 * p.N >> 2) + n4]); + if (m0 + 1u < p.M) acc1 += float4(res[((m0 + 1u) * p.N >> 2) + n4]); + if (m0 + 2u < p.M) acc2 += float4(res[((m0 + 2u) * p.N >> 2) + n4]); + if (m0 + 3u < p.M) acc3 += float4(res[((m0 + 3u) * p.N >> 2) + n4]); + } + float4 av0 = actf4(acc0, act, p.flags), av1 = actf4(acc1, act, p.flags); + float4 av2 = actf4(acc2, act, p.flags), av3 = actf4(acc3, act, p.flags); + if (m0 < p.M) o[m0 * nv + n4] = half4(av0); + if (m0 + 1u < p.M) o[(m0 + 1u) * nv + n4] = half4(av1); + if (m0 + 2u < p.M) o[(m0 + 2u) * nv + n4] = half4(av2); + if (m0 + 3u < p.M) o[(m0 + 3u) * nv + n4] = half4(av3); + if ((p.flags & 8u) != 0u) { + // dual-write: o2 = gelu(o); bit10 (1024) → o2 = raw pre-act acc + float4 g0, g1, g2, g3; + if ((p.flags & 1024u) != 0u) { g0 = acc0; g1 = acc1; g2 = acc2; g3 = acc3; } + else { + g0 = 0.5f * av0 * (float4(1.0) + float4(erf1(av0.x * 0.70710678f), erf1(av0.y * 0.70710678f), erf1(av0.z * 0.70710678f), erf1(av0.w * 0.70710678f))); + g1 = 0.5f * av1 * (float4(1.0) + float4(erf1(av1.x * 0.70710678f), erf1(av1.y * 0.70710678f), erf1(av1.z * 0.70710678f), erf1(av1.w * 0.70710678f))); + g2 = 0.5f * av2 * (float4(1.0) + float4(erf1(av2.x * 0.70710678f), erf1(av2.y * 0.70710678f), erf1(av2.z * 0.70710678f), erf1(av2.w * 0.70710678f))); + g3 = 0.5f * av3 * (float4(1.0) + float4(erf1(av3.x * 0.70710678f), erf1(av3.y * 0.70710678f), erf1(av3.z * 0.70710678f), erf1(av3.w * 0.70710678f))); + } + if (m0 < p.M) o2[m0 * nv + n4] = half4(g0); + if (m0 + 1u < p.M) o2[(m0 + 1u) * nv + n4] = half4(g1); + if (m0 + 2u < p.M) o2[(m0 + 2u) * nv + n4] = half4(g2); + if (m0 + 3u < p.M) o2[(m0 + 3u) * nv + n4] = half4(g3); + } +} + +// ---------------- mmdot_sk ---------------- +// Split-K dot: 256 threads = 16 lanes (k-partitions) × 16 tiles. Adjacent +// local ids = adjacent tiles at the same k → weight reads coalesce. +// Partials go through threadgroup memory [lane][tile][16] (padded 17). +kernel void mmdot_sk(device const half4* x [[buffer(0)]], + device const half4* w [[buffer(1)]], + device const half4* bias [[buffer(2)]], + device const half4* res [[buffer(3)]], + device half4* o [[buffer(4)]], + device const half4* s [[buffer(5)]], + device half4* o2 [[buffer(6)]], + device const half4* ra [[buffer(7)]], + constant PcDot& p [[buffer(8)]], + uint lid [[thread_position_in_threadgroup]], + uint wg [[threadgroup_position_in_grid]]) +{ + // [lane][tile][16] partials, row padded to 17 floats vs bank conflicts + threadgroup float sh[256 * 17]; + uint tl = lid & 15u; + uint lane = lid >> 4; + uint nv = p.N >> 2; + uint mv = (p.M + 3u) >> 2; + uint tile = wg * 16u + tl; + bool live = tile < mv * nv; + uint n4 = live ? tile % nv : 0u; + uint m0 = live ? (tile / nv) * 4u : 0u; + + float4 acc0 = float4(0.0), acc1 = float4(0.0), acc2 = float4(0.0), acc3 = float4(0.0); + bool preScale = (p.flags & 4u) != 0u; + uint sb0 = 0u, sb1 = 0u, sb2 = 0u, sb3 = 0u; + if (p.mImg != 0u && preScale) { + sb0 = (m0 / p.mImg) * p.scMod; + sb1 = ((m0 + 1u) / p.mImg) * p.scMod; + sb2 = ((m0 + 2u) / p.mImg) * p.scMod; + sb3 = ((m0 + 3u) / p.mImg) * p.scMod; + } + if (live) { + for (uint k4 = lane; k4 < p.kv4; k4 += 16u) { + float4 a0 = m0 < p.M ? float4(x[m0 * p.kv4 + k4]) : float4(0.0); + float4 a1 = m0 + 1u < p.M ? float4(x[(m0 + 1u) * p.kv4 + k4]) : float4(0.0); + float4 a2 = m0 + 2u < p.M ? float4(x[(m0 + 2u) * p.kv4 + k4]) : float4(0.0); + float4 a3 = m0 + 3u < p.M ? float4(x[(m0 + 3u) * p.kv4 + k4]) : float4(0.0); + if (preScale) { + a0 *= float4(s[sb0 + k4 % p.scMod]); + a1 *= float4(s[sb1 + k4 % p.scMod]); + a2 *= float4(s[sb2 + k4 % p.scMod]); + a3 *= float4(s[sb3 + k4 % p.scMod]); + } + uint wRow = k4 * 4u * nv + n4; + for (uint i = 0u; i < 4u; i++) { + float4 wv = float4(w[wRow + i * nv]); + acc0 += a0[i] * wv; acc1 += a1[i] * wv; + acc2 += a2[i] * wv; acc3 += a3[i] * wv; + } + } + } + uint bb = (lane * 16u + tl) * 17u; + sh[bb + 0u] = acc0.x; sh[bb + 1u] = acc0.y; sh[bb + 2u] = acc0.z; sh[bb + 3u] = acc0.w; + sh[bb + 4u] = acc1.x; sh[bb + 5u] = acc1.y; sh[bb + 6u] = acc1.z; sh[bb + 7u] = acc1.w; + sh[bb + 8u] = acc2.x; sh[bb + 9u] = acc2.y; sh[bb + 10u] = acc2.z; sh[bb + 11u] = acc2.w; + sh[bb + 12u] = acc3.x; sh[bb + 13u] = acc3.y; sh[bb + 14u] = acc3.z; sh[bb + 15u] = acc3.w; + threadgroup_barrier(mem_flags::mem_threadgroup); + // this invocation now reduces element e = lane of tile tl over all lanes + float sum = 0.0f; + for (uint l = 0u; l < 16u; l++) sum += sh[(l * 16u + tl) * 17u + lane]; + threadgroup_barrier(mem_flags::mem_threadgroup); + sh[tl * 17u + lane] = sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (!live || lane >= 4u) return; + uint m = m0 + lane; + if (m >= p.M) return; + uint r = tl * 17u + lane * 4u; + float4 v = float4(sh[r], sh[r + 1u], sh[r + 2u], sh[r + 3u]); + if ((p.flags & 1u) != 0u) v += float4(bias[n4]); + if ((p.flags & 2u) != 0u) v += float4(res[m * nv + n4]); + uint act = (p.flags >> 4) & 7u; + o[m * nv + n4] = half4(actf4(v, act, p.flags)); +} + +struct PcK { + uint M, N, outW, inW, inH, cin4v, taps, kW, sH, sW, padT, padL, flags; + uint dstCv, dstCoff, cin, inHW; +}; + +// ---------------- convk_dot ---------------- +// Dense kxk conv as direct dot WITHOUT im2col: k = tap*cin4v + ci read +// straight from NHWC input. gy = batch. +kernel void convk_dot(device const half4* x [[buffer(0)]], + device const half4* w [[buffer(1)]], + device const half4* bias [[buffer(2)]], + device const half4* res [[buffer(3)]], + device half4* o [[buffer(4)]], + constant PcK& p [[buffer(5)]], + uint2 gp [[thread_position_in_grid]], + uint2 wg [[threadgroup_position_in_grid]]) +{ + uint gid = gp.x; + uint b = wg.y; + uint nv = p.N >> 2; + uint mv = (p.M + 3u) >> 2; + if (gid >= mv * nv) return; + uint n4 = gid % nv; + uint m0 = (gid / nv) * 4u; + uint xb = b * p.inH * p.inW * p.cin4v; + uint rb = b * (p.M * p.N >> 2); + uint dcv0 = p.dstCv != 0u ? p.dstCv : nv; + uint ob = b * p.M * dcv0; + + float4 acc0 = float4(0.0), acc1 = float4(0.0), acc2 = float4(0.0), acc3 = float4(0.0); + uint oh[4], ow[4]; + bool ok[4]; + for (uint i = 0u; i < 4u; i++) { + uint m = m0 + i; + ok[i] = m < p.M; + oh[i] = m / p.outW; + ow[i] = m - oh[i] * p.outW; + } + + for (uint t = 0u; t < p.taps; t++) { + uint ky = t / p.kW, kx = t - ky * p.kW; + for (uint ci4 = 0u; ci4 < p.cin4v; ci4++) { + float4 a0 = float4(0.0), a1 = float4(0.0), a2 = float4(0.0), a3 = float4(0.0); + for (uint i = 0u; i < 4u; i++) { + if (!ok[i]) continue; + int ih = int(oh[i] * p.sH + ky) - int(p.padT); + int iw = int(ow[i] * p.sW + kx) - int(p.padL); + if (ih < 0 || uint(ih) >= p.inH || iw < 0 || uint(iw) >= p.inW) + continue; + uint px = uint(ih) * p.inW + uint(iw); + float4 v = float4(x[xb + px * p.cin4v + ci4]); + if (i == 0u) a0 = v; else if (i == 1u) a1 = v; + else if (i == 2u) a2 = v; else a3 = v; + } + uint wRow = (t * p.cin4v + ci4) * 4u * nv + n4; + for (uint i = 0u; i < 4u; i++) { + float4 wv = float4(w[wRow + i * nv]); + acc0 += a0[i] * wv; acc1 += a1[i] * wv; + acc2 += a2[i] * wv; acc3 += a3[i] * wv; + } + } + } + if ((p.flags & 1u) != 0u) { + float4 bv = float4(bias[n4]); + acc0 += bv; acc1 += bv; acc2 += bv; acc3 += bv; + } + uint act = (p.flags >> 4) & 7u; + if ((p.flags & 2u) != 0u) { + if (m0 < p.M) acc0 += float4(res[rb + (m0 * p.N >> 2) + n4]); + if (m0 + 1u < p.M) acc1 += float4(res[rb + ((m0 + 1u) * p.N >> 2) + n4]); + if (m0 + 2u < p.M) acc2 += float4(res[rb + ((m0 + 2u) * p.N >> 2) + n4]); + if (m0 + 3u < p.M) acc3 += float4(res[rb + ((m0 + 3u) * p.N >> 2) + n4]); + } + uint dcv = dcv0; + uint obase = p.dstCoff + n4; + if (m0 < p.M) o[ob + m0 * dcv + obase] = half4(actf4(acc0, act, p.flags)); + if (m0 + 1u < p.M) o[ob + (m0 + 1u) * dcv + obase] = half4(actf4(acc1, act, p.flags)); + if (m0 + 2u < p.M) o[ob + (m0 + 2u) * dcv + obase] = half4(actf4(acc2, act, p.flags)); + if (m0 + 3u < p.M) o[ob + (m0 + 3u) * dcv + obase] = half4(actf4(acc3, act, p.flags)); +} + +// ---------------- convk_f32n ---------------- +// convk_dot whose A side reads the fp32 NCHW graph input directly (stem conv, +// cin <= 4 → single vec4 column per tap). gy = batch. +kernel void convk_f32n(device const float* x [[buffer(0)]], + device const half4* w [[buffer(1)]], + device const half4* bias [[buffer(2)]], + device const half4* res [[buffer(3)]], + device half4* o [[buffer(4)]], + constant PcK& p [[buffer(5)]], + uint2 gp [[thread_position_in_grid]], + uint2 wg [[threadgroup_position_in_grid]]) +{ + uint gid = gp.x; + uint b = wg.y; + uint nv = p.N >> 2; + uint mv = (p.M + 3u) >> 2; + if (gid >= mv * nv) return; + uint n4 = gid % nv; + uint m0 = (gid / nv) * 4u; + uint xb = b * p.cin * p.inHW; + uint rb = b * (p.M * p.N >> 2); + uint dcv0 = p.dstCv != 0u ? p.dstCv : nv; + uint ob = b * p.M * dcv0; + + float4 acc0 = float4(0.0), acc1 = float4(0.0), acc2 = float4(0.0), acc3 = float4(0.0); + uint oh[4], ow[4]; + bool ok[4]; + for (uint i = 0u; i < 4u; i++) { + uint m = m0 + i; + ok[i] = m < p.M; + oh[i] = m / p.outW; + ow[i] = m - oh[i] * p.outW; + } + + for (uint t = 0u; t < p.taps; t++) { + uint ky = t / p.kW, kx = t - ky * p.kW; + float4 a0 = float4(0.0), a1 = float4(0.0), a2 = float4(0.0), a3 = float4(0.0); + for (uint i = 0u; i < 4u; i++) { + if (!ok[i]) continue; + int ih = int(oh[i] * p.sH + ky) - int(p.padT); + int iw = int(ow[i] * p.sW + kx) - int(p.padL); + if (ih < 0 || uint(ih) >= p.inH || iw < 0 || uint(iw) >= p.inW) + continue; + uint px = uint(ih) * p.inW + uint(iw); + float4 v = float4(0.0); + for (uint ci = 0u; ci < 4u; ci++) + if (ci < p.cin) v[ci] = x[xb + ci * p.inHW + px]; + if (i == 0u) a0 = v; else if (i == 1u) a1 = v; + else if (i == 2u) a2 = v; else a3 = v; + } + uint wRow = (t * p.cin4v) * 4u * nv + n4; + for (uint i = 0u; i < 4u; i++) { + float4 wv = float4(w[wRow + i * nv]); + acc0 += a0[i] * wv; acc1 += a1[i] * wv; + acc2 += a2[i] * wv; acc3 += a3[i] * wv; + } + } + if ((p.flags & 1u) != 0u) { + float4 bv = float4(bias[n4]); + acc0 += bv; acc1 += bv; acc2 += bv; acc3 += bv; + } + uint act = (p.flags >> 4) & 7u; + if ((p.flags & 2u) != 0u) { + if (m0 < p.M) acc0 += float4(res[rb + (m0 * p.N >> 2) + n4]); + if (m0 + 1u < p.M) acc1 += float4(res[rb + ((m0 + 1u) * p.N >> 2) + n4]); + if (m0 + 2u < p.M) acc2 += float4(res[rb + ((m0 + 2u) * p.N >> 2) + n4]); + if (m0 + 3u < p.M) acc3 += float4(res[rb + ((m0 + 3u) * p.N >> 2) + n4]); + } + uint dcv = dcv0; + uint obase = p.dstCoff + n4; + if (m0 < p.M) o[ob + m0 * dcv + obase] = half4(actf4(acc0, act, p.flags)); + if (m0 + 1u < p.M) o[ob + (m0 + 1u) * dcv + obase] = half4(actf4(acc1, act, p.flags)); + if (m0 + 2u < p.M) o[ob + (m0 + 2u) * dcv + obase] = half4(actf4(acc2, act, p.flags)); + if (m0 + 3u < p.M) o[ob + (m0 + 3u) * dcv + obase] = half4(actf4(acc3, act, p.flags)); +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/elem.metal b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/elem.metal new file mode 100644 index 0000000..720d4d0 --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/elem.metal @@ -0,0 +1,175 @@ +// Port of Vulkan Shaders/{nchw2nhwc,sigmoid_out,elem,elem4,addps,affine4}.comp +// Same semantics; fp16 arena, fp32 compute. The GLSL uint-pair reads become +// native half / half4 loads (vec4 binds always sit on 8B-aligned offsets). +#include +using namespace metal; +// erf1 / auxHalf2 live in a_common.metal (compiled into the same library). + +// ---------------- nchw2nhwc ---------------- +// Graph input convert: fp32 NCHW [n,C,H,W] → fp16 NHWC [n,H*W,Cout]. +// Cout may exceed C (zero-padded channels for aligned downstream reads). +struct Pc3 { uint hw, C, Cout; }; +kernel void nchw2nhwc(device const float* x [[buffer(0)]], + device half* o [[buffer(1)]], + constant Pc3& p [[buffer(2)]], + uint2 gid [[thread_position_in_grid]]) +{ + uint total = p.hw * p.Cout; + if (gid.x >= total) return; + uint c = gid.x % p.Cout; + uint px = gid.x / p.Cout; + uint b = gid.y; + o[b * total + gid.x] = c < p.C + ? half(x[b * p.C * p.hw + c * p.hw + px]) : half(0.0); +} + +// ---------------- sigmoid_out ---------------- +// Final head: fp16 → fp32 into the readback buffer, optional sigmoid. +struct Pc2 { uint n, act; }; +kernel void sigmoid_out(device const half* x [[buffer(0)]], + device float* o [[buffer(1)]], + constant Pc2& p [[buffer(2)]], + uint i [[thread_position_in_grid]]) +{ + if (i >= p.n) return; + float v = float(x[i]); + o[i] = (p.act == 4u) ? 1.0 / (1.0 + exp(-v)) : v; +} + +// ---------------- elem / elem4 ---------------- +// Generic elementwise on NHWC fp16 activations. +// unary: o[i] = f(a[i]) (op 0-7, 14) +// binary: o[i] = a[ai(i)] OP b[bi(i)] (op 8+) +// Broadcast modes per operand (aMode/bMode): 0=full index, 1=channel vector +// (index % C), 2=scalar (index 0), 3=per-image channel vector (aux = +// elements per image). +// op unary: 0=copy 1=relu 2=gelu 3=hardsigmoid(aux) 4=sigmoid 5=erf 6=sqrt +// 14=hardswish (x*hs(x), aux) +// op binary: 8=add 9=mul 10=sub 11=div 12=max 13=pow +struct PcElem { uint n, C, op, aMode, bMode, aux; }; + +static inline uint elem_idx(uint i, uint mode, constant PcElem& p) +{ + if (mode == 1u) return i % p.C; + if (mode == 2u) return 0u; + if (mode == 3u) return (i / p.aux) * p.C + i % p.C; + return i; +} + +static inline float elem_unary(uint op, float va, uint aux) +{ + switch (op) { + case 1u: return max(va, 0.0f); + case 2u: return 0.5f * va * (1.0f + erf1(va * 0.70710678118654752f)); + case 3u: { float2 ab = auxHalf2(aux); return clamp(va * ab.x + ab.y, 0.0f, 1.0f); } + case 4u: return 1.0f / (1.0f + exp(-va)); + case 5u: return erf1(va); + case 6u: return sqrt(max(va, 0.0f)); + case 14u: { float2 ab = auxHalf2(aux); return va * clamp(va * ab.x + ab.y, 0.0f, 1.0f); } + default: return va; + } +} + +kernel void elem(device const half* a [[buffer(0)]], + device const half* b [[buffer(1)]], + device half* o [[buffer(2)]], + constant PcElem& p [[buffer(3)]], + uint i [[thread_position_in_grid]]) +{ + if (i >= p.n) return; + float va = float(a[elem_idx(i, p.aMode, p)]); + float r; + if (p.op < 8u || p.op == 14u) { + r = elem_unary(p.op, va, p.aux); + } else { + float vb = float(b[elem_idx(i, p.bMode, p)]); + switch (p.op) { + case 9u: r = va * vb; break; + case 10u: r = va - vb; break; + case 11u: r = va / vb; break; + case 12u: r = max(va, vb); break; + case 13u: r = pow(va, vb); break; + default: r = va + vb; break; + } + } + o[i] = half(r); +} + +kernel void elem4(device const half4* a [[buffer(0)]], + device const half4* b [[buffer(1)]], + device half4* o [[buffer(2)]], + constant PcElem& p [[buffer(3)]], + uint gid [[thread_position_in_grid]]) +{ + uint i0 = gid * 4u; + if (i0 >= p.n) return; + + float4 va; + { + uint base = p.aMode == 1u ? i0 % p.C + : p.aMode == 3u ? (i0 / p.aux) * p.C + i0 % p.C : i0; + va = p.aMode == 2u ? float4(float(a[0].x)) : float4(a[base >> 2]); + } + float4 r; + if (p.op < 8u || p.op == 14u) { + if (p.op == 2u || p.op == 5u) { + for (int i = 0; i < 4; i++) r[i] = elem_unary(p.op, va[i], p.aux); + } else if (p.op == 3u || p.op == 14u) { + float2 ab = auxHalf2(p.aux); + float4 hs = clamp(va * ab.x + ab.y, 0.0f, 1.0f); + r = p.op == 14u ? va * hs : hs; + } else { + for (int i = 0; i < 4; i++) r[i] = elem_unary(p.op, va[i], p.aux); + } + } else { + float4 vb; + uint base = p.bMode == 1u ? i0 % p.C + : p.bMode == 3u ? (i0 / p.aux) * p.C + i0 % p.C : i0; + vb = p.bMode == 2u ? float4(float(b[0].x)) : float4(b[base >> 2]); + if (p.op == 9u) r = va * vb; + else if (p.op == 10u) r = va - vb; + else if (p.op == 11u) r = va / vb; + else if (p.op == 12u) r = max(va, vb); + else if (p.op == 13u) r = pow(va, vb); + else r = va + vb; + } + o[gid] = half4(r); +} + +// ---------------- addps ---------------- +// Scaled residual add on NHWC fp16: o[i] = a[i] + x[i] * s[c(i)] +// s is a per-channel fp16 vector with cv4 = C/4 vec4 groups. Requires +// n%4==0 && C%4==0. (per-image SE: si is vec4-indexed within the image) +struct PcAddps { uint n, cv4, perImg4; }; +kernel void addps(device const half4* a [[buffer(0)]], + device const half4* x [[buffer(1)]], + device const half4* s [[buffer(2)]], + device half4* o [[buffer(3)]], + constant PcAddps& p [[buffer(4)]], + uint gid [[thread_position_in_grid]]) +{ + uint i0 = gid * 4u; + if (i0 >= p.n) return; + uint v = i0 >> 2; + // cv4 = vec4 count of the channel dim; perImg4 = vec4 count per image + // (0 = same numel as out → channel-only). Per-batch SE [n,C,1,1]: + // perImg4>0 → index within image's channel vector. + uint si = p.perImg4 > 0 ? (v / p.perImg4) * p.cv4 + v % p.cv4 : v % p.cv4; + o[v] = a[v] + x[v] * s[si]; +} + +// ---------------- affine4 ---------------- +// Per-channel affine on NHWC fp16 (folded BatchNormalization): +// o[i] = x[i] * s[c4] + t[c4] i = flat vec4 index, c4 = i % cv +struct PcAff { uint n, cv; }; +kernel void affine4(device const half4* x [[buffer(0)]], + device const half4* s [[buffer(1)]], + device const half4* t [[buffer(2)]], + device half4* o [[buffer(3)]], + constant PcAff& p [[buffer(4)]], + uint gid [[thread_position_in_grid]]) +{ + if (gid >= p.n) return; + uint c4 = gid % p.cv; + o[gid] = x[gid] * s[c4] + t[c4]; +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/gemm.metal b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/gemm.metal new file mode 100644 index 0000000..d89beb0 --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/gemm.metal @@ -0,0 +1,602 @@ +// Metal GEMM for the conv1x1 / im2col / MatMul "cm" paths. +// O[M,N] = act(A[M,K] · W[N,K]^T + bias + res) (+ dual-write o2 = gelu/raw) +// A: fp16 [M, K] row-major (vec4-strided NHWC slabs; M folds the batch dim) +// W: fp16 [Npad, K] n-major (ConstF16 / ConstGemmW repack) +// flags: bit0 bias, bit1 res, bits4-6 act, hi16 hsA, bit3 dual-write +// (o2 = gelu(o); bit10 → o2 = raw pre-act). +// mm_sg: simdgroup_matrix 8x8x8 fp16 MMA, group tile 64x64, K-step 16 staged +// fp16 — requires simdgroup matrix support (probed by MetalBackend). +// mm_ps_sg: same + per-batch SE prescale on A (bit2): a *= s[b*scMod + k]. +#include +using namespace metal; + +struct PcGemm { uint M, N, K, kv4, flags, mImg, scMod; }; + +// Cooperative fp16 staging for the 64x64x16 k-tile: +// A stage [64][17] halfs (row = m), B stage [64][17] halfs (row = n). +// 256 threads: A stage thread t covers m=t>>2, k4=(t&3)*4; +// B stage covers n=t>>2, k4=(t&3)*4. +static inline void gemm_stage( + threadgroup half* tA, threadgroup half* tB, + device const half4* x, device const half* w, + uint m0, uint n0, uint kb, uint tid, + uint M, uint N, uint K, uint kv4) +{ + uint am = tid >> 2, ak = (tid & 3u) * 4u; + uint gm = m0 + am, gk = kb + ak; + half4 av = half4(0.0); + if (gm < M && gk < K) + av = x[(ulong)gm * kv4 + (gk >> 2)]; // K%4==0 guaranteed by emit + *(threadgroup half4*)(tA + am * 17 + ak) = av; + + uint bn = tid >> 2, bk = (tid & 3u) * 4u; + uint gn = n0 + bn; + half4 wv = half4(0.0); + if (gk < K) + wv = *(device const half4*)(w + (ulong)gn * K + gk); // n-major [N,K] + *(threadgroup half4*)(tB + bn * 17 + bk) = wv; +} + +kernel void mm_sg(device const half4* x [[buffer(0)]], + device const half* w [[buffer(1)]], + device const half* bias [[buffer(2)]], + device const half* res [[buffer(3)]], + device half4* o [[buffer(4)]], + device half4* o2 [[buffer(5)]], + device const half4* se [[buffer(6)]], // prescale (unused here) + constant PcGemm& p [[buffer(7)]], + uint2 wg [[threadgroup_position_in_grid]], + uint2 tp2 [[thread_position_in_threadgroup]]) +{ + uint tid = tp2.x; + threadgroup half tA[64 * 17]; + threadgroup half tB[64 * 17]; + threadgroup float outst[8 * 64]; // per-sg 8x8 epilogue stage + + uint sgid = tid >> 5; + uint sm = (sgid & 3u) * 16; // 4 sgs along M (16 rows each) + uint sn = (sgid >> 2) * 32; // 2 sgs along N (32 cols each) + uint m0 = wg.y * 64u + sm; + uint n0 = wg.x * 64u + sn; + + simdgroup_float8x8 acc[2][4]; + for (uint i = 0u; i < 2u; i++) for (uint j = 0u; j < 4u; j++) + acc[i][j] = simdgroup_float8x8(0.0f); + + for (uint kb = 0u; kb < p.K; kb += 16u) { + gemm_stage(tA, tB, x, w, wg.y * 64u, wg.x * 64u, kb, tid, + p.M, p.N, p.K, p.kv4); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint kk = 0u; kk < 16u; kk += 8u) { + simdgroup_half8x8 a0, a1; + simdgroup_load(a0, tA + sm * 17 + kk, 17); + simdgroup_load(a1, tA + (sm + 8u) * 17 + kk, 17); + simdgroup_half8x8 b0, b1, b2, b3; + simdgroup_load(b0, tB + (sn + 0u) * 17 + kk, 17, ulong2(0, 0), true); + simdgroup_load(b1, tB + (sn + 8u) * 17 + kk, 17, ulong2(0, 0), true); + simdgroup_load(b2, tB + (sn + 16u) * 17 + kk, 17, ulong2(0, 0), true); + simdgroup_load(b3, tB + (sn + 24u) * 17 + kk, 17, ulong2(0, 0), true); + simdgroup_multiply_accumulate(acc[0][0], a0, b0, acc[0][0]); + simdgroup_multiply_accumulate(acc[0][1], a0, b1, acc[0][1]); + simdgroup_multiply_accumulate(acc[0][2], a0, b2, acc[0][2]); + simdgroup_multiply_accumulate(acc[0][3], a0, b3, acc[0][3]); + simdgroup_multiply_accumulate(acc[1][0], a1, b0, acc[1][0]); + simdgroup_multiply_accumulate(acc[1][1], a1, b1, acc[1][1]); + simdgroup_multiply_accumulate(acc[1][2], a1, b2, acc[1][2]); + simdgroup_multiply_accumulate(acc[1][3], a1, b3, acc[1][3]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + // Epilogue: stage each 8x8 acc to threadgroup, scalars write fp16. + // outst layout [sg][64] fp32; row r col c at stage[r*8+c]. + uint lane = tid & 31u; + uint nv = p.N >> 2; + for (uint i = 0u; i < 2u; i++) + for (uint j = 0u; j < 4u; j++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + simdgroup_store(acc[i][j], outst + sgid * 64, 8); + threadgroup_barrier(mem_flags::mem_threadgroup); + // 32 lanes × 2 elems: lane l covers (r=l>>2, c=(l&3)*2 + k) + uint r = lane >> 2, c2 = (lane & 3u) * 2u; + for (uint e = 0u; e < 2u; e++) { + uint mm = m0 + i * 8u + r; + uint nn = n0 + j * 8u + c2 + e; + if (mm >= p.M || nn >= p.N) continue; + float v = outst[sgid * 64 + r * 8 + c2 + e]; + if ((p.flags & 1u) != 0u) v += float(bias[nn]); + if ((p.flags & 2u) != 0u) v += float(res[mm * p.N + nn]); + float av = actf(v, (p.flags >> 4) & 7u, p.flags); + // o/o2 are half4-bound: scalar write via half* view + ((device half*)o)[mm * p.N + nn] = half(av); + if ((p.flags & 8u) != 0u) + ((device half*)o2)[mm * p.N + nn] = + (p.flags & 1024u) != 0u ? half(v) : half(0.5f * av * (1.0f + erf1(av * 0.70710678f))); + } + } +} + +// mm_sg_dd: same 64x64 tile / 8 simdgroup fan-out as mm_sg, but fragment loads +// read device memory directly — no threadgroup staging round-trip (this GPU's +// "threadgroup" memory is just a carve-out of the same DRAM, so staging only +// adds barriers). Emit-gated to M%64==0 && N%64==0 && K%16==0 so every tile is +// full and fragment loads never run past the operand end. +kernel void mm_sg_dd(device const half* x [[buffer(0)]], + device const half* w [[buffer(1)]], + device const half* bias [[buffer(2)]], + device const half* res [[buffer(3)]], + device half4* o [[buffer(4)]], + device half4* o2 [[buffer(5)]], + device const half4* se [[buffer(6)]], + constant PcGemm& p [[buffer(7)]], + uint2 wg [[threadgroup_position_in_grid]], + uint2 tp2 [[thread_position_in_threadgroup]]) +{ + uint tid = tp2.x; + threadgroup float outst[8 * 64]; + + uint sgid = tid >> 5; + uint sm = (sgid & 3u) * 16; + uint sn = (sgid >> 2) * 32; + uint m0 = wg.y * 64u + sm; + uint n0 = wg.x * 64u + sn; + + simdgroup_float8x8 acc[2][4]; + for (uint i = 0u; i < 2u; i++) for (uint j = 0u; j < 4u; j++) + acc[i][j] = simdgroup_float8x8(0.0f); + + for (uint kb = 0u; kb < p.K; kb += 16u) { + simdgroup_half8x8 a0, a1, a00, a01; + simdgroup_load(a0, x + (ulong)(m0 * p.K) + kb, p.K); + simdgroup_load(a1, x + (ulong)((m0 + 8u) * p.K) + kb, p.K); + simdgroup_load(a00, x + (ulong)(m0 * p.K) + kb + 8u, p.K); + simdgroup_load(a01, x + (ulong)((m0 + 8u) * p.K) + kb + 8u, p.K); + simdgroup_half8x8 b0, b1, b2, b3, b00, b01, b02, b03; + simdgroup_load(b0, w + (ulong)((n0 + 0u) * p.K) + kb, p.K, ulong2(0, 0), true); + simdgroup_load(b1, w + (ulong)((n0 + 8u) * p.K) + kb, p.K, ulong2(0, 0), true); + simdgroup_load(b2, w + (ulong)((n0 + 16u) * p.K) + kb, p.K, ulong2(0, 0), true); + simdgroup_load(b3, w + (ulong)((n0 + 24u) * p.K) + kb, p.K, ulong2(0, 0), true); + simdgroup_load(b00, w + (ulong)((n0 + 0u) * p.K) + kb + 8u, p.K, ulong2(0, 0), true); + simdgroup_load(b01, w + (ulong)((n0 + 8u) * p.K) + kb + 8u, p.K, ulong2(0, 0), true); + simdgroup_load(b02, w + (ulong)((n0 + 16u) * p.K) + kb + 8u, p.K, ulong2(0, 0), true); + simdgroup_load(b03, w + (ulong)((n0 + 24u) * p.K) + kb + 8u, p.K, ulong2(0, 0), true); + simdgroup_multiply_accumulate(acc[0][0], a0, b0, acc[0][0]); + simdgroup_multiply_accumulate(acc[0][1], a0, b1, acc[0][1]); + simdgroup_multiply_accumulate(acc[0][2], a0, b2, acc[0][2]); + simdgroup_multiply_accumulate(acc[0][3], a0, b3, acc[0][3]); + simdgroup_multiply_accumulate(acc[1][0], a1, b0, acc[1][0]); + simdgroup_multiply_accumulate(acc[1][1], a1, b1, acc[1][1]); + simdgroup_multiply_accumulate(acc[1][2], a1, b2, acc[1][2]); + simdgroup_multiply_accumulate(acc[1][3], a1, b3, acc[1][3]); + simdgroup_multiply_accumulate(acc[0][0], a00, b00, acc[0][0]); + simdgroup_multiply_accumulate(acc[0][1], a00, b01, acc[0][1]); + simdgroup_multiply_accumulate(acc[0][2], a00, b02, acc[0][2]); + simdgroup_multiply_accumulate(acc[0][3], a00, b03, acc[0][3]); + simdgroup_multiply_accumulate(acc[1][0], a01, b00, acc[1][0]); + simdgroup_multiply_accumulate(acc[1][1], a01, b01, acc[1][1]); + simdgroup_multiply_accumulate(acc[1][2], a01, b02, acc[1][2]); + simdgroup_multiply_accumulate(acc[1][3], a01, b03, acc[1][3]); + } + + uint lane = tid & 31u; + for (uint i = 0u; i < 2u; i++) + for (uint j = 0u; j < 4u; j++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + simdgroup_store(acc[i][j], outst + sgid * 64, 8); + threadgroup_barrier(mem_flags::mem_threadgroup); + uint r = lane >> 2, c2 = (lane & 3u) * 2u; + for (uint e = 0u; e < 2u; e++) { + uint mm = m0 + i * 8u + r; + uint nn = n0 + j * 8u + c2 + e; + if (mm >= p.M || nn >= p.N) continue; + float v = outst[sgid * 64 + r * 8 + c2 + e]; + if ((p.flags & 1u) != 0u) v += float(bias[nn]); + if ((p.flags & 2u) != 0u) v += float(res[mm * p.N + nn]); + float av = actf(v, (p.flags >> 4) & 7u, p.flags); + ((device half*)o)[mm * p.N + nn] = half(av); + if ((p.flags & 8u) != 0u) + ((device half*)o2)[mm * p.N + nn] = + (p.flags & 1024u) != 0u ? half(v) : half(0.5f * av * (1.0f + erf1(av * 0.70710678f))); + } + } +} + +// mm_ps_sg: prescale (bit2) variant — a[m,k] *= s[b*scMod + k]; mImg = output +// pixels per image (batch split of the flat M range). +kernel void mm_ps_sg(device const half4* x [[buffer(0)]], + device const half* w [[buffer(1)]], + device const half* bias [[buffer(2)]], + device const half* res [[buffer(3)]], + device half4* o [[buffer(4)]], + device half4* o2 [[buffer(5)]], + device const half4* se [[buffer(6)]], + constant PcGemm& p [[buffer(7)]], + uint2 wg [[threadgroup_position_in_grid]], + uint2 tp2 [[thread_position_in_threadgroup]]) +{ + uint tid = tp2.x; + threadgroup half tA[64 * 17]; + threadgroup half tB[64 * 17]; + threadgroup float outst[8 * 64]; + + uint sgid = tid >> 5; + uint sm = (sgid & 3u) * 16; + uint sn = (sgid >> 2) * 32; + uint m0 = wg.y * 64u + sm; + uint n0 = wg.x * 64u + sn; + + simdgroup_float8x8 acc[2][4]; + for (uint i = 0u; i < 2u; i++) for (uint j = 0u; j < 4u; j++) + acc[i][j] = simdgroup_float8x8(0.0f); + + for (uint kb = 0u; kb < p.K; kb += 16u) { + // stage A with prescale applied + { + uint am = tid >> 2, ak = (tid & 3u) * 4u; + uint gm = wg.y * 64u + am, gk = kb + ak; + half4 av = half4(0.0); + if (gm < p.M && gk < p.K) { + av = x[(ulong)gm * p.kv4 + (gk >> 2)]; + // se[b*scMod + k4%scMod] — half4 group per k-block + av = half4(float4(av) * + float4(se[(gm / p.mImg) * p.scMod + (gk >> 2) % p.scMod])); + } + *(threadgroup half4*)(tA + am * 17 + ak) = av; + } + { + uint bn = tid >> 2, bk = (tid & 3u) * 4u; + uint gn = wg.x * 64u + bn; + half4 wv = half4(0.0); + if (kb + bk < p.K) + wv = *(device const half4*)(w + (ulong)gn * p.K + kb + bk); + *(threadgroup half4*)(tB + bn * 17 + bk) = wv; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint kk = 0u; kk < 16u; kk += 8u) { + simdgroup_half8x8 a0, a1; + simdgroup_load(a0, tA + sm * 17 + kk, 17); + simdgroup_load(a1, tA + (sm + 8u) * 17 + kk, 17); + simdgroup_half8x8 b0, b1, b2, b3; + simdgroup_load(b0, tB + (sn + 0u) * 17 + kk, 17, ulong2(0, 0), true); + simdgroup_load(b1, tB + (sn + 8u) * 17 + kk, 17, ulong2(0, 0), true); + simdgroup_load(b2, tB + (sn + 16u) * 17 + kk, 17, ulong2(0, 0), true); + simdgroup_load(b3, tB + (sn + 24u) * 17 + kk, 17, ulong2(0, 0), true); + simdgroup_multiply_accumulate(acc[0][0], a0, b0, acc[0][0]); + simdgroup_multiply_accumulate(acc[0][1], a0, b1, acc[0][1]); + simdgroup_multiply_accumulate(acc[0][2], a0, b2, acc[0][2]); + simdgroup_multiply_accumulate(acc[0][3], a0, b3, acc[0][3]); + simdgroup_multiply_accumulate(acc[1][0], a1, b0, acc[1][0]); + simdgroup_multiply_accumulate(acc[1][1], a1, b1, acc[1][1]); + simdgroup_multiply_accumulate(acc[1][2], a1, b2, acc[1][2]); + simdgroup_multiply_accumulate(acc[1][3], a1, b3, acc[1][3]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + uint lane = tid & 31u; + for (uint i = 0u; i < 2u; i++) + for (uint j = 0u; j < 4u; j++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + simdgroup_store(acc[i][j], outst + sgid * 64, 8); + threadgroup_barrier(mem_flags::mem_threadgroup); + uint r = lane >> 2, c2 = (lane & 3u) * 2u; + for (uint e = 0u; e < 2u; e++) { + uint mm = m0 + i * 8u + r; + uint nn = n0 + j * 8u + c2 + e; + if (mm >= p.M || nn >= p.N) continue; + float v = outst[sgid * 64 + r * 8 + c2 + e]; + if ((p.flags & 1u) != 0u) v += float(bias[nn]); + if ((p.flags & 2u) != 0u) v += float(res[mm * p.N + nn]); + float av = actf(v, (p.flags >> 4) & 7u, p.flags); + ((device half*)o)[mm * p.N + nn] = half(av); + if ((p.flags & 8u) != 0u) + ((device half*)o2)[mm * p.N + nn] = + (p.flags & 1024u) != 0u ? half(v) : half(0.5f * av * (1.0f + erf1(av * 0.70710678f))); + } + } +} + +// ---------------- mm_dot (portable fallback) ---------------- +// Same contract as mm_sg but no simdgroup builtins: thread per (m4,n4) tile, +// W still n-major [N,K] — each (m,ncol) acc accumulates dot(a_m, w_n4..). +// Used when the simdgroup-matrix probe fails (kernel still correct at any +// M/N/K; conv1x1-dot-style gating keeps it off the biggest shapes anyway). +kernel void mm_dot(device const half4* x [[buffer(0)]], + device const half4* w [[buffer(1)]], + device const half4* bias [[buffer(2)]], + device const half4* res [[buffer(3)]], + device half4* o [[buffer(4)]], + device half4* o2 [[buffer(5)]], + device const half4* se [[buffer(6)]], + constant PcGemm& p [[buffer(7)]], + uint gid [[thread_position_in_grid]]) +{ + uint nv = p.N >> 2; + uint mv = (p.M + 3u) >> 2; + if (gid >= mv * nv) return; + uint n4 = gid % nv; + uint m0 = (gid / nv) * 4u; + bool preScale = (p.flags & 4u) != 0u; + uint sb0 = 0u, sb1 = 0u, sb2 = 0u, sb3 = 0u; + if (preScale && p.mImg != 0u) { + sb0 = (m0 / p.mImg) * p.scMod; + sb1 = ((m0 + 1u) / p.mImg) * p.scMod; + sb2 = ((m0 + 2u) / p.mImg) * p.scMod; + sb3 = ((m0 + 3u) / p.mImg) * p.scMod; + } + float4 acc0 = float4(0.0), acc1 = float4(0.0), acc2 = float4(0.0), acc3 = float4(0.0); + for (uint k4 = 0u; k4 < p.kv4; k4++) { + float4 a0 = m0 < p.M ? float4(x[m0 * p.kv4 + k4]) : float4(0.0); + float4 a1 = m0 + 1u < p.M ? float4(x[(m0 + 1u) * p.kv4 + k4]) : float4(0.0); + float4 a2 = m0 + 2u < p.M ? float4(x[(m0 + 2u) * p.kv4 + k4]) : float4(0.0); + float4 a3 = m0 + 3u < p.M ? float4(x[(m0 + 3u) * p.kv4 + k4]) : float4(0.0); + if (preScale) { + a0 *= float4(se[sb0 + k4 % p.scMod]); + a1 *= float4(se[sb1 + k4 % p.scMod]); + a2 *= float4(se[sb2 + k4 % p.scMod]); + a3 *= float4(se[sb3 + k4 % p.scMod]); + } + // W n-major: w[n*K + k] — acc_i[nj] += dot(a_i, w4(nj)) per k4 + float4 wv0 = float4(w[(n4 * 4u + 0u) * p.kv4 + k4]); + float4 wv1 = float4(w[(n4 * 4u + 1u) * p.kv4 + k4]); + float4 wv2 = float4(w[(n4 * 4u + 2u) * p.kv4 + k4]); + float4 wv3 = float4(w[(n4 * 4u + 3u) * p.kv4 + k4]); + acc0.x += dot(a0, wv0); acc0.y += dot(a0, wv1); + acc0.z += dot(a0, wv2); acc0.w += dot(a0, wv3); + acc1.x += dot(a1, wv0); acc1.y += dot(a1, wv1); + acc1.z += dot(a1, wv2); acc1.w += dot(a1, wv3); + acc2.x += dot(a2, wv0); acc2.y += dot(a2, wv1); + acc2.z += dot(a2, wv2); acc2.w += dot(a2, wv3); + acc3.x += dot(a3, wv0); acc3.y += dot(a3, wv1); + acc3.z += dot(a3, wv2); acc3.w += dot(a3, wv3); + } + if ((p.flags & 1u) != 0u) { + float4 bv = float4(bias[n4]); + acc0 += bv; acc1 += bv; acc2 += bv; acc3 += bv; + } + if ((p.flags & 2u) != 0u) { + if (m0 < p.M) acc0 += float4(res[m0 * nv + n4]); + if (m0 + 1u < p.M) acc1 += float4(res[(m0 + 1u) * nv + n4]); + if (m0 + 2u < p.M) acc2 += float4(res[(m0 + 2u) * nv + n4]); + if (m0 + 3u < p.M) acc3 += float4(res[(m0 + 3u) * nv + n4]); + } + uint act = (p.flags >> 4) & 7u; + float4 av0 = actf4(acc0, act, p.flags), av1 = actf4(acc1, act, p.flags); + float4 av2 = actf4(acc2, act, p.flags), av3 = actf4(acc3, act, p.flags); + if (m0 < p.M) o[m0 * nv + n4] = half4(av0); + if (m0 + 1u < p.M) o[(m0 + 1u) * nv + n4] = half4(av1); + if (m0 + 2u < p.M) o[(m0 + 2u) * nv + n4] = half4(av2); + if (m0 + 3u < p.M) o[(m0 + 3u) * nv + n4] = half4(av3); + if ((p.flags & 8u) != 0u) { + float4 g0, g1, g2, g3; + if ((p.flags & 1024u) != 0u) { g0 = acc0; g1 = acc1; g2 = acc2; g3 = acc3; } + else { + g0 = 0.5f * av0 * (float4(1.0) + float4(erf1(av0.x * 0.70710678f), erf1(av0.y * 0.70710678f), erf1(av0.z * 0.70710678f), erf1(av0.w * 0.70710678f))); + g1 = 0.5f * av1 * (float4(1.0) + float4(erf1(av1.x * 0.70710678f), erf1(av1.y * 0.70710678f), erf1(av1.z * 0.70710678f), erf1(av1.w * 0.70710678f))); + g2 = 0.5f * av2 * (float4(1.0) + float4(erf1(av2.x * 0.70710678f), erf1(av2.y * 0.70710678f), erf1(av2.z * 0.70710678f), erf1(av2.w * 0.70710678f))); + g3 = 0.5f * av3 * (float4(1.0) + float4(erf1(av3.x * 0.70710678f), erf1(av3.y * 0.70710678f), erf1(av3.z * 0.70710678f), erf1(av3.w * 0.70710678f))); + } + if (m0 < p.M) o2[m0 * nv + n4] = half4(g0); + if (m0 + 1u < p.M) o2[(m0 + 1u) * nv + n4] = half4(g1); + if (m0 + 2u < p.M) o2[(m0 + 2u) * nv + n4] = half4(g2); + if (m0 + 3u < p.M) o2[(m0 + 3u) * nv + n4] = half4(g3); + } +} + +// mm_ic_sg: implicit-GEMM kxk conv — mm_sg's staged 64x64x16 pipeline, but the +// A stage gathers input pixels through tap-major im2col addressing instead of +// reading a materialized [M,K] matrix. One staged half4 never crosses a tap +// (cinE%4==0 gate on the emit side). k in [K,Kp) pads to zero via the +// iy/ix bounds check. Batch-free like the im2col path it replaces. +struct PcGemmIc { + uint M, N, K, kv4, flags, mImg, scMod; // K = Kp (padded, W stride) + uint outW, inW, inH, sH, sW, pT, pL, kW, cinE; +}; + +kernel void mm_ic_sg(device const half4* x [[buffer(0)]], + device const half* w [[buffer(1)]], + device const half* bias [[buffer(2)]], + device const half* res [[buffer(3)]], + device half4* o [[buffer(4)]], + device half4* o2 [[buffer(5)]], + device const half4* se [[buffer(6)]], + constant PcGemmIc& p [[buffer(7)]], + uint2 wg [[threadgroup_position_in_grid]], + uint2 tp2 [[thread_position_in_threadgroup]]) +{ + uint tid = tp2.x; + threadgroup half tA[64 * 17]; + threadgroup half tB[64 * 17]; + threadgroup float outst[8 * 64]; + + uint sgid = tid >> 5; + uint sm = (sgid & 3u) * 16; + uint sn = (sgid >> 2) * 32; + uint m0 = wg.y * 64u + sm; + uint n0 = wg.x * 64u + sn; + + simdgroup_float8x8 acc[2][4]; + for (uint i = 0u; i < 2u; i++) for (uint j = 0u; j < 4u; j++) + acc[i][j] = simdgroup_float8x8(0.0f); + + uint am = tid >> 2, ak = (tid & 3u) * 4u; + uint gm0 = wg.y * 64u + am; + uint bn = tid >> 2, bk = (tid & 3u) * 4u; + uint gn0 = wg.x * 64u + bn; + + for (uint kb = 0u; kb < p.K; kb += 16u) { + // A stage with inline im2col gather: element k = tap*cinE + ci, + // tap = dy*kW + dx. + uint gk = kb + ak; + half4 av = half4(0.0); + if (gm0 < p.M && gk < p.K) { + uint tap = gk / p.cinE, ci = gk - tap * p.cinE; + uint dy = tap / p.kW, dx = tap - dy * p.kW; + uint oy = gm0 / p.outW, ox = gm0 - oy * p.outW; + int iy = int(oy * p.sH + dy) - int(p.pT); + int ix = int(ox * p.sW + dx) - int(p.pL); + if (iy >= 0 && iy < int(p.inH) && ix >= 0 && ix < int(p.inW)) + av = x[(uint(iy) * p.inW + uint(ix)) * (p.cinE >> 2) + (ci >> 2)]; + } + *(threadgroup half4*)(tA + am * 17 + ak) = av; + + uint gkB = kb + bk; + half4 wv = half4(0.0); + if (gkB < p.K) + wv = *(device const half4*)(w + (ulong)gn0 * p.K + gkB); + *(threadgroup half4*)(tB + bn * 17 + bk) = wv; + + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint kk = 0u; kk < 16u; kk += 8u) { + simdgroup_half8x8 a0, a1; + simdgroup_load(a0, tA + sm * 17 + kk, 17); + simdgroup_load(a1, tA + (sm + 8u) * 17 + kk, 17); + simdgroup_half8x8 b0, b1, b2, b3; + simdgroup_load(b0, tB + (sn + 0u) * 17 + kk, 17, ulong2(0, 0), true); + simdgroup_load(b1, tB + (sn + 8u) * 17 + kk, 17, ulong2(0, 0), true); + simdgroup_load(b2, tB + (sn + 16u) * 17 + kk, 17, ulong2(0, 0), true); + simdgroup_load(b3, tB + (sn + 24u) * 17 + kk, 17, ulong2(0, 0), true); + simdgroup_multiply_accumulate(acc[0][0], a0, b0, acc[0][0]); + simdgroup_multiply_accumulate(acc[0][1], a0, b1, acc[0][1]); + simdgroup_multiply_accumulate(acc[0][2], a0, b2, acc[0][2]); + simdgroup_multiply_accumulate(acc[0][3], a0, b3, acc[0][3]); + simdgroup_multiply_accumulate(acc[1][0], a1, b0, acc[1][0]); + simdgroup_multiply_accumulate(acc[1][1], a1, b1, acc[1][1]); + simdgroup_multiply_accumulate(acc[1][2], a1, b2, acc[1][2]); + simdgroup_multiply_accumulate(acc[1][3], a1, b3, acc[1][3]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + uint lane = tid & 31u; + for (uint i = 0u; i < 2u; i++) + for (uint j = 0u; j < 4u; j++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + simdgroup_store(acc[i][j], outst + sgid * 64, 8); + threadgroup_barrier(mem_flags::mem_threadgroup); + uint r = lane >> 2, c2 = (lane & 3u) * 2u; + for (uint e = 0u; e < 2u; e++) { + uint mm = m0 + i * 8u + r; + uint nn = n0 + j * 8u + c2 + e; + if (mm >= p.M || nn >= p.N) continue; + float v = outst[sgid * 64 + r * 8 + c2 + e]; + if ((p.flags & 1u) != 0u) v += float(bias[nn]); + if ((p.flags & 2u) != 0u) v += float(res[mm * p.N + nn]); + float av = actf(v, (p.flags >> 4) & 7u, p.flags); + ((device half*)o)[mm * p.N + nn] = half(av); + if ((p.flags & 8u) != 0u) + ((device half*)o2)[mm * p.N + nn] = + (p.flags & 1024u) != 0u ? half(v) : half(0.5f * av * (1.0f + erf1(av * 0.70710678f))); + } + } +} + +// mm_ic_sg32: implicit-GEMM kxk conv — mm_sg's staged 64x64x16 pipeline, but the +// A stage gathers input pixels through tap-major im2col addressing instead of +// reading a materialized [M,K] matrix. One staged half4 never crosses a tap +// (cinE%4==0 gate on the emit side). k in [K,Kp) pads to zero via the +// iy/ix bounds check. Batch-free like the im2col path it replaces. +struct PcGemmIc32 { + uint M, N, K, kv4, flags, mImg, scMod; // K = Kp (padded, W stride) + uint outW, inW, inH, sH, sW, pT, pL, kW, cinE; +}; + +kernel void mm_ic_sg32(device const half4* x [[buffer(0)]], + device const half* w [[buffer(1)]], + device const half* bias [[buffer(2)]], + device const half* res [[buffer(3)]], + device half4* o [[buffer(4)]], + device half4* o2 [[buffer(5)]], + device const half4* se [[buffer(6)]], + constant PcGemmIc32& p [[buffer(7)]], + uint2 wg [[threadgroup_position_in_grid]], + uint2 tp2 [[thread_position_in_threadgroup]]) +{ + uint tid = tp2.x; + threadgroup half tA[64 * 34]; + threadgroup half tB[64 * 34]; + threadgroup float outst[8 * 64]; + + uint sgid = tid >> 5; + uint sm = (sgid & 3u) * 16; + uint sn = (sgid >> 2) * 32; + uint m0 = wg.y * 64u + sm; + uint n0 = wg.x * 64u + sn; + + simdgroup_float8x8 acc[2][4]; + for (uint i = 0u; i < 2u; i++) for (uint j = 0u; j < 4u; j++) + acc[i][j] = simdgroup_float8x8(0.0f); + + uint am = tid >> 2, ak = (tid & 3u) * 4u; + uint gm0 = wg.y * 64u + am; + uint bn = tid >> 2, bk = (tid & 3u) * 4u; + uint gn0 = wg.x * 64u + bn; + + for (uint kb = 0u; kb < p.K; kb += 32u) { + // A stage with inline im2col gather: element k = tap*cinE + ci, + // tap = dy*kW + dx. + uint oy = gm0 / p.outW, ox = gm0 - oy * p.outW; + for (uint h = 0u; h < 2u; h++) { + uint gk = kb + ak + h * 16u; + half4 av = half4(0.0); + if (gm0 < p.M && gk < p.K) { + uint tap = gk / p.cinE, ci = gk - tap * p.cinE; + uint dy = tap / p.kW, dx = tap - dy * p.kW; + int iy = int(oy * p.sH + dy) - int(p.pT); + int ix = int(ox * p.sW + dx) - int(p.pL); + if (iy >= 0 && iy < int(p.inH) && ix >= 0 && ix < int(p.inW)) + av = x[(uint(iy) * p.inW + uint(ix)) * (p.cinE >> 2) + (ci >> 2)]; + } + *(threadgroup half4*)(tA + am * 34 + ak + h * 16u) = av; + + uint gkB = kb + bk + h * 16u; + half4 wv = half4(0.0); + if (gkB < p.K) + wv = *(device const half4*)(w + (ulong)gn0 * p.K + gkB); + *(threadgroup half4*)(tB + bn * 34 + bk + h * 16u) = wv; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint kk = 0u; kk < 32u; kk += 8u) { + simdgroup_half8x8 a0, a1; + simdgroup_load(a0, tA + sm * 34 + kk, 34); + simdgroup_load(a1, tA + (sm + 8u) * 34 + kk, 34); + simdgroup_half8x8 b0, b1, b2, b3; + simdgroup_load(b0, tB + (sn + 0u) * 34 + kk, 34, ulong2(0, 0), true); + simdgroup_load(b1, tB + (sn + 8u) * 34 + kk, 34, ulong2(0, 0), true); + simdgroup_load(b2, tB + (sn + 16u) * 34 + kk, 34, ulong2(0, 0), true); + simdgroup_load(b3, tB + (sn + 24u) * 34 + kk, 34, ulong2(0, 0), true); + simdgroup_multiply_accumulate(acc[0][0], a0, b0, acc[0][0]); + simdgroup_multiply_accumulate(acc[0][1], a0, b1, acc[0][1]); + simdgroup_multiply_accumulate(acc[0][2], a0, b2, acc[0][2]); + simdgroup_multiply_accumulate(acc[0][3], a0, b3, acc[0][3]); + simdgroup_multiply_accumulate(acc[1][0], a1, b0, acc[1][0]); + simdgroup_multiply_accumulate(acc[1][1], a1, b1, acc[1][1]); + simdgroup_multiply_accumulate(acc[1][2], a1, b2, acc[1][2]); + simdgroup_multiply_accumulate(acc[1][3], a1, b3, acc[1][3]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + uint lane = tid & 31u; + for (uint i = 0u; i < 2u; i++) + for (uint j = 0u; j < 4u; j++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + simdgroup_store(acc[i][j], outst + sgid * 64, 8); + threadgroup_barrier(mem_flags::mem_threadgroup); + uint r = lane >> 2, c2 = (lane & 3u) * 2u; + for (uint e = 0u; e < 2u; e++) { + uint mm = m0 + i * 8u + r; + uint nn = n0 + j * 8u + c2 + e; + if (mm >= p.M || nn >= p.N) continue; + float v = outst[sgid * 64 + r * 8 + c2 + e]; + if ((p.flags & 1u) != 0u) v += float(bias[nn]); + if ((p.flags & 2u) != 0u) v += float(res[mm * p.N + nn]); + float av = actf(v, (p.flags >> 4) & 7u, p.flags); + ((device half*)o)[mm * p.N + nn] = half(av); + if ((p.flags & 8u) != 0u) + ((device half*)o2)[mm * p.N + nn] = + (p.flags & 1024u) != 0u ? half(v) : half(0.5f * av * (1.0f + erf1(av * 0.70710678f))); + } + } +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/space.metal b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/space.metal new file mode 100644 index 0000000..2651cb2 --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/space.metal @@ -0,0 +1,350 @@ +// Port of {concat_c, concat4, resize_nn, resize4, resize4add, catresize, +// maxpool2e, maxpool4, avgpool4, reduce_hw, reduce_hw4, reduce_hw4b}.comp +// NHWC fp16; gid.y (thread_position_in_grid.y) = batch where applicable. +#include +using namespace metal; + +// ---------------- concat_c ---------------- +// NHWC channel-concat slice copy: o[px, cOff + c] = x[px, c], c < Cin. +struct PcConcat { uint hw, Cin, Cout, cOff; }; +kernel void concat_c(device const half* x [[buffer(0)]], + device half* o [[buffer(1)]], + constant PcConcat& p [[buffer(2)]], + uint gid [[thread_position_in_grid]]) +{ + uint total = p.hw * p.Cin; + if (gid >= total) return; + uint c = gid % p.Cin; + uint px = gid / p.Cin; + o[px * p.Cout + p.cOff + c] = x[gid]; +} + +// ---------------- concat4 ---------------- +// vec4 concat slice (Cin%4==0, cOff%4==0); flags bit8: x is unfused addps +// → x[i]*se[c4] + ra[i]. +struct PcConcat4 { uint hw, Cin, Cout, cOff, flags; }; +kernel void concat4(device const half4* x [[buffer(0)]], + device half4* o [[buffer(1)]], + device const half4* ra [[buffer(2)]], + device const half4* se [[buffer(3)]], + constant PcConcat4& p [[buffer(4)]], + uint gid [[thread_position_in_grid]]) +{ + uint cv = p.Cin >> 2; + if (gid >= p.hw * cv) return; + uint c4 = gid % cv; + uint px = gid / cv; + float4 v = float4(x[gid]); + if ((p.flags & 256u) != 0u) v = v * float4(se[c4]) + float4(ra[gid]); + o[px * (p.Cout >> 2) + (p.cOff >> 2) + c4] = half4(v); +} + +// ---------------- resize_nn ---------------- +// Nearest-neighbour integer-factor upscale on NHWC fp16. +struct PcRs { uint outW, outH, C, fh, fw; }; +kernel void resize_nn(device const half* x [[buffer(0)]], + device half* o [[buffer(1)]], + constant PcRs& p [[buffer(2)]], + uint gid [[thread_position_in_grid]]) +{ + uint total = p.outW * p.outH * p.C; + if (gid >= total) return; + uint c = gid % p.C; + uint px = gid / p.C; + uint oy = px / p.outW, ox = px % p.outW; + uint inW = p.outW / p.fw; + o[gid] = x[((oy / p.fh) * inW + (ox / p.fw)) * p.C + c]; +} + +// ---------------- resize4 ---------------- +// vec4 variant; dstCv != 0 → write into a wider-C tensor (concat slice). +// flags bit8: input is unfused addps → x[i]*se[c4] + ra[i]. +struct PcRs4 { uint outW, outH, C, fh, fw, dstCv, dstCoff, flags; }; +kernel void resize4(device const half4* x [[buffer(0)]], + device half4* o [[buffer(1)]], + device const half4* ra [[buffer(2)]], + device const half4* se [[buffer(3)]], + constant PcRs4& p [[buffer(4)]], + uint gid [[thread_position_in_grid]]) +{ + uint cv = p.C >> 2; + if (gid >= p.outW * p.outH * cv) return; + uint c4 = gid % cv; + uint px = gid / cv; + uint oy = px / p.outW, ox = px % p.outW; + uint inW = p.outW / p.fw; + uint dcv = p.dstCv != 0u ? p.dstCv : cv; + uint idx = ((oy / p.fh) * inW + (ox / p.fw)) * cv + c4; + float4 v = float4(x[idx]); + if ((p.flags & 256u) != 0u) v = v * float4(se[c4]) + float4(ra[idx]); + o[px * dcv + p.dstCoff + c4] = half4(v); +} + +// ---------------- resize4add ---------------- +// Fused upscale + elementwise add (FPN merge): o = a + resize(x). +// flags bit8: x unfused addps; bit9: a unfused addps. +struct PcRs4A { uint outW, outH, C, fh, fw, flags; }; +kernel void resize4add(device const half4* x [[buffer(0)]], + device const half4* a [[buffer(1)]], + device half4* o [[buffer(2)]], + device const half4* ra [[buffer(3)]], + device const half4* se [[buffer(4)]], + constant PcRs4A& p [[buffer(5)]], + uint gid [[thread_position_in_grid]]) +{ + uint cv = p.C >> 2; + if (gid >= p.outW * p.outH * cv) return; + uint c4 = gid % cv; + uint px = gid / cv; + uint oy = px / p.outW, ox = px % p.outW; + uint inW = p.outW / p.fw; + uint idx = ((oy / p.fh) * inW + (ox / p.fw)) * cv + c4; + float4 xv = float4(x[idx]), av = float4(a[gid]); + if ((p.flags & 256u) != 0u) xv = xv * float4(se[c4]) + float4(ra[idx]); + if ((p.flags & 512u) != 0u) av = av * float4(se[c4]) + float4(ra[gid]); + o[gid] = half4(av + xv); +} + +// ---------------- catresize ---------------- +// Fused channel-concat: each input slice is a (nearest) resize of a possibly +// addps-absorbed source. Up to 4 slices; slice s occupies output channel-quads +// [off4_s, off4_s + cq4_s). Per-slice flags bit0 = addps absorb +// (v = x[i]*se[q] + ra[i]). +struct PcCatRes { + uint outW, outH, Cout; + uint off4a, inWa, fha, fla; + uint off4b, inWb, fhb, flb; + uint off4c, inWc, fhc, flc; + uint off4d, inWd, fhd, fld; +}; +struct CatSlice { device const half4 *x, *r, *s; }; +static inline float4 cat_src(device const half4* x, device const half4* r, + device const half4* s, uint idx, uint q, uint fl) +{ + float4 v = float4(x[idx]); + if ((fl & 1u) != 0u) v = v * float4(s[q]) + float4(r[idx]); + return v; +} +kernel void catresize(device const half4* x0 [[buffer(0)]], + device const half4* r0 [[buffer(1)]], + device const half4* s0 [[buffer(2)]], + device const half4* x1 [[buffer(3)]], + device const half4* r1 [[buffer(4)]], + device const half4* s1 [[buffer(5)]], + device const half4* x2 [[buffer(6)]], + device const half4* r2 [[buffer(7)]], + device const half4* s2 [[buffer(8)]], + device const half4* x3 [[buffer(9)]], + device const half4* r3 [[buffer(10)]], + device const half4* s3 [[buffer(11)]], + device half4* o [[buffer(12)]], + constant PcCatRes& p [[buffer(13)]], + uint gid [[thread_position_in_grid]]) +{ + uint cv4 = p.Cout >> 2; + if (gid >= p.outW * p.outH * cv4) return; + uint c4 = gid % cv4; + uint px = gid / cv4; + uint oy = px / p.outW, ox = px % p.outW; + + for (uint s = 0u; s < 4u; s++) { + uint offc, inW, fhfw, fl; + switch (s) { + case 0u: offc = p.off4a; inW = p.inWa; fhfw = p.fha; fl = p.fla; break; + case 1u: offc = p.off4b; inW = p.inWb; fhfw = p.fhb; fl = p.flb; break; + case 2u: offc = p.off4c; inW = p.inWc; fhfw = p.fhc; fl = p.flc; break; + default: offc = p.off4d; inW = p.inWd; fhfw = p.fhd; fl = p.fld; break; + } + uint off4 = offc & 0xffffu, cq4 = offc >> 16; + if (cq4 == 0u || c4 < off4 || c4 >= off4 + cq4) continue; + uint q = c4 - off4; + uint sy = oy / (fhfw & 0xffffu), sx = ox / (fhfw >> 16); + uint idx = (sy * inW + sx) * cq4 + q; + float4 v; + switch (s) { + case 0u: v = cat_src(x0, r0, s0, idx, q, fl); break; + case 1u: v = cat_src(x1, r1, s1, idx, q, fl); break; + case 2u: v = cat_src(x2, r2, s2, idx, q, fl); break; + default: v = cat_src(x3, r3, s3, idx, q, fl); break; + } + o[gid] = half4(v); + return; + } +} + +// ---------------- maxpool2e ---------------- +// MaxPool 2x2 stride 1, pad-end, NHWC fp16 — same HxW out. +struct PcPool { uint W, H, C; }; +kernel void maxpool2e(device const half* x [[buffer(0)]], + device half* o [[buffer(1)]], + constant PcPool& p [[buffer(2)]], + uint gid [[thread_position_in_grid]]) +{ + uint total = p.W * p.H * p.C; + if (gid >= total) return; + uint c = gid % p.C; + uint px = gid / p.C; + uint y = px / p.W, xx = px % p.W; + float m = -1.0e30f; + for (uint dy = 0u; dy < 2u; dy++) { + uint iy = y + dy; + if (iy >= p.H) continue; + for (uint dx = 0u; dx < 2u; dx++) { + uint ix = xx + dx; + if (ix >= p.W) continue; + m = max(m, float(x[(iy * p.W + ix) * p.C + c])); + } + } + o[gid] = half(m); +} + +// ---------------- maxpool4 ---------------- +// vec4 variant (C%4==0); dstCv != 0 → concat-slice write. gid.y = batch. +struct PcPool4 { uint W, H, C, dstCv, dstCoff; }; +kernel void maxpool4(device const half4* x [[buffer(0)]], + device half4* o [[buffer(1)]], + constant PcPool4& p [[buffer(2)]], + uint2 gid [[thread_position_in_grid]]) +{ + uint cv = p.C >> 2; + if (gid.x >= p.W * p.H * cv) return; + uint c4 = gid.x % cv; + uint px = gid.x / cv; + uint y = px / p.W, xx = px % p.W; + uint b = gid.y; + uint xb = b * p.W * p.H * cv; + + half4 m = half4(half(-1.0e4)); + for (uint dy = 0u; dy < 2u; dy++) { + uint iy = y + dy; + if (iy >= p.H) continue; + for (uint dx = 0u; dx < 2u; dx++) { + uint ix = xx + dx; + if (ix >= p.W) continue; + half4 v = x[xb + (iy * p.W + ix) * cv + c4]; + m = max(m, v); + } + } + uint dcv = p.dstCv != 0u ? p.dstCv : cv; + o[(b * p.W * p.H + px) * dcv + p.dstCoff + c4] = m; +} + +// ---------------- avgpool4 ---------------- +// AveragePool NHWC fp16 vec4 (C%4==0); OOB excluded (count_include_pad=0). +// gid.y = batch. PC: outW, outH, C, kH, kW, sh, sw, pt, pl, inW, inH. +struct PcAvg { + uint outW, outH, C, kH, kW, sh, sw, pt, pl, inW, inH; +}; +kernel void avgpool4(device const half4* x [[buffer(0)]], + device half4* o [[buffer(1)]], + constant PcAvg& p [[buffer(2)]], + uint2 gid [[thread_position_in_grid]]) +{ + uint b = gid.y; + uint cv = p.C >> 2; + uint total = p.outW * p.outH * cv; + if (gid.x >= total) return; + uint c4 = gid.x % cv; + uint px = gid.x / cv; + uint oy = px / p.outW, ox = px % p.outW; + int iy0 = int(oy * p.sh) - int(p.pt); + int ix0 = int(ox * p.sw) - int(p.pl); + uint xb = b * p.inH * p.inW * cv; + + float4 acc = float4(0.0); + uint cnt = 0u; + for (uint dy = 0u; dy < p.kH; dy++) { + int iy = iy0 + int(dy); + if (iy < 0 || iy >= int(p.inH)) continue; + uint rowBase = xb + uint(iy) * p.inW * cv; + for (uint dx = 0u; dx < p.kW; dx++) { + int ix = ix0 + int(dx); + if (ix < 0 || ix >= int(p.inW)) continue; + acc += float4(x[rowBase + uint(ix) * cv + c4]); + cnt++; + } + } + if (cnt == 0u) cnt = 1u; + o[b * total + gid.x] = half4(acc / float(cnt)); +} + +// ---------------- reduce_hw ---------------- +// Per-channel spatial mean: out[c] = mean_px x[px, c]. One threadgroup per +// channel, 256 threads strided-reduce into shared, fp32 acc. gid.y = batch. +struct PcRed { uint hw, C; }; +kernel void reduce_hw(device const half* x [[buffer(0)]], + device half* o [[buffer(1)]], + constant PcRed& p [[buffer(2)]], + uint2 wg [[threadgroup_position_in_grid]], + uint2 lp2 [[thread_position_in_threadgroup]]) +{ + threadgroup float sm[256]; + uint lane = lp2.x; + uint c = wg.x; + uint b = wg.y; + float acc = 0.0f; + for (uint px = lane; px < p.hw; px += 256u) + acc += float(x[b * p.hw * p.C + px * p.C + c]); + sm[lane] = acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = 128u; s > 0u; s >>= 1) { + if (lane < s) sm[lane] += sm[lane + s]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lane == 0u) o[b * p.C + c] = half(sm[0] / float(p.hw)); +} + +// ---------------- reduce_hw4 (phase 1) ---------------- +// Partition partial sums: group s covers pixels [s*P, s*P+P); 256 lanes laid +// out (strip = lane/cv, cLane = lane%cv) → coalesced vec4 reads. Writes fp32 +// partials. gid.y = batch. +struct PcRed4 { uint hw, C, S, P; }; +kernel void reduce_hw4(device const half4* x [[buffer(0)]], + device float* part [[buffer(1)]], + constant PcRed4& p [[buffer(2)]], + uint2 wg [[threadgroup_position_in_grid]], + uint2 lp2 [[thread_position_in_threadgroup]]) +{ + threadgroup float4 sm[256]; + uint lane = lp2.x; + uint s = wg.x; + uint b = wg.y; + uint cv = p.C >> 2; + uint pb = b * p.S * p.C; + uint xb = b * p.hw * cv; + uint strips = 256u / cv; + uint strip = lane / cv; + uint cLane = lane - strip * cv; + float4 acc = float4(0.0); + uint beg = s * p.P; + uint end = min(beg + p.P, p.hw); + for (uint px = beg + strip; px < end; px += strips) + acc += float4(x[xb + px * cv + cLane]); + sm[lane] = acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint st = strips >> 1; st > 0u; st >>= 1) { + if (strip < st) sm[lane] += sm[lane + st * cv]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (strip == 0u) { + uint ob = pb + s * p.C + cLane * 4u; + part[ob] = sm[lane].x; part[ob + 1u] = sm[lane].y; + part[ob + 2u] = sm[lane].z; part[ob + 3u] = sm[lane].w; + } +} + +// ---------------- reduce_hw4b (phase 2) ---------------- +// out[c] = (sum_s part[s][c]) / hw. One thread per channel; gid.y = batch. +struct PcRed4b { uint S, C, hw; }; +kernel void reduce_hw4b(device const float* part [[buffer(0)]], + device half* o [[buffer(1)]], + constant PcRed4b& p [[buffer(2)]], + uint2 gid [[thread_position_in_grid]]) +{ + uint c = gid.x; + uint b = gid.y; + if (c >= p.C) return; + float acc = 0.0f; + for (uint s = 0u; s < p.S; s++) acc += part[b * p.S * p.C + s * p.C + c]; + o[b * p.C + c] = half(acc / float(p.hw)); +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/xform.metal b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/xform.metal new file mode 100644 index 0000000..8d522c7 --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/xform.metal @@ -0,0 +1,236 @@ +// Port of {se_fused, layernorm, softmax, attn}.comp — SE block, row +// LayerNorm, row softmax, fused MHSA. All fp16 activations, fp32 stats. +#include +using namespace metal; + +// ---------------- se_fused ---------------- +// Single-pass SE: S threadgroups write per-partition partial sums; the last +// group to finish (atomic ticket) runs finalize inline: mean → fc1+relu → +// fc2 → hardsigmoid. +struct PcSe { uint hw, C, S, P, cvP, R; float alpha, beta; }; +// SE split into two dispatches: this paravirt GPU only supports +// memory_order_relaxed device atomics, so the Vulkan-style ticket +// (last-WG-does-stage2) cannot order partial writes vs reads. Emit +// se_part for all S workgroups, rely on the inter-dispatch buffer +// barrier, then se_join (one WG per batch) does stage2. +kernel void se_part(device const half4* x [[buffer(0)]], + device float* part [[buffer(1)]], + constant PcSe& p [[buffer(2)]], + uint2 wg [[threadgroup_position_in_grid]], + uint2 lp2 [[thread_position_in_threadgroup]]) +{ + uint lane = lp2.x; + threadgroup float4 sm[256]; + + uint s = wg.x; + uint b = wg.y; + uint cv = p.C >> 2; + uint pb = b * p.S * p.C; + uint xb = b * p.hw * cv; + uint strips = 256u / p.cvP; + uint strip = lane / p.cvP; + uint cLane = lane - strip * p.cvP; + float4 acc = float4(0.0); + uint beg = s * p.P; + uint end = min(beg + p.P, p.hw); + if (cLane < cv) + for (uint px = beg + strip; px < end; px += strips) + acc += float4(x[xb + px * cv + cLane]); + sm[lane] = acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint st = strips >> 1; st > 0u; st >>= 1) { + if (strip < st) sm[lane] += sm[lane + st * p.cvP]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (strip == 0u && cLane < cv) { + uint ob = pb + s * p.C + cLane * 4u; + part[ob] = sm[lane].x; part[ob + 1u] = sm[lane].y; + part[ob + 2u] = sm[lane].z; part[ob + 3u] = sm[lane].w; + } +} + +kernel void se_join(device const float* part [[buffer(0)]], + device const float* w1 [[buffer(1)]], + device const float* b1 [[buffer(2)]], + device const float* w2 [[buffer(3)]], + device const float* b2 [[buffer(4)]], + device half4* o [[buffer(5)]], + constant PcSe& p [[buffer(6)]], + uint2 wg [[threadgroup_position_in_grid]], + uint2 lp2 [[thread_position_in_threadgroup]]) +{ + uint lane = lp2.x; + threadgroup float zf[256]; + threadgroup float hf[64]; + + uint b = wg.x; + uint cv = p.C >> 2; + uint pb = b * p.S * p.C; + + if (lane < p.C) { + float a = 0.0f; + for (uint si = 0u; si < p.S; si++) a += part[pb + si * p.C + lane]; + zf[lane] = a / float(p.hw); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (lane < p.R) { + float acc2 = b1[lane]; + for (uint c = 0u; c < p.C; c++) + acc2 += zf[c] * w1[lane * p.C + c]; + hf[lane] = max(acc2, 0.0f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (lane < cv) { + half v[4]; + for (uint i = 0u; i < 4u; i++) { + uint c = lane * 4u + i; + float acc3 = b2[c]; + for (uint r = 0u; r < p.R; r++) + acc3 += hf[r] * w2[c * p.R + r]; + v[i] = half(clamp(p.alpha * acc3 + p.beta, 0.0f, 1.0f)); + } + o[b * cv + lane] = half4(v[0], v[1], v[2], v[3]); + } +} + +// ---------------- layernorm ---------------- +// Row LayerNorm over the contiguous last axis: +// y[r,c] = (x[r,c] - mean_r)/sqrt(var_r+eps) * g[c] + b[c] +// One 64-lane group per row, fp32 stats; C <= 64*16. +struct PcLn { uint rows, C; uint epsBits; }; +kernel void layernorm(device const half* x [[buffer(0)]], + device const float* g [[buffer(1)]], + device const float* bt [[buffer(2)]], + device half* o [[buffer(3)]], + constant PcLn& p [[buffer(4)]], + uint wg [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]]) +{ + threadgroup float red[64]; + uint r = wg; + if (r >= p.rows) return; + uint base = r * p.C; + float v[16]; + float s = 0.0f; + for (uint i = 0u; i < 16u; i++) { + uint c = t + i * 64u; + v[i] = c < p.C ? float(x[base + c]) : 0.0f; + s += v[i]; + } + // wgSum + red[t] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint k = 32u; k > 0u; k >>= 1) { + if (t < k) red[t] += red[t + k]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + float mean = red[0] / float(p.C); + threadgroup_barrier(mem_flags::mem_threadgroup); + + float q = 0.0f; + for (uint i = 0u; i < 16u; i++) { + uint c = t + i * 64u; + float d = c < p.C ? v[i] - mean : 0.0f; + q += d * d; + } + red[t] = q; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint k = 32u; k > 0u; k >>= 1) { + if (t < k) red[t] += red[t + k]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + float var = red[0] / float(p.C); + float inv = 1.0f / sqrt(var + as_type(p.epsBits)); + for (uint i = 0u; i < 16u; i++) { + uint c = t + i * 64u; + if (c < p.C) o[base + c] = half((v[i] - mean) * inv * g[c] + bt[c]); + } +} + +// ---------------- softmax ---------------- +// Row softmax on fp16: y[r,:] = softmax(x[r,:]) over the last axis. +struct PcSm { uint rows, cols; }; +kernel void softmax(device const half* x [[buffer(0)]], + device half* y [[buffer(1)]], + constant PcSm& p [[buffer(2)]], + uint r [[thread_position_in_grid]]) +{ + if (r >= p.rows) return; + uint b = r * p.cols; + float m = -3.4e38f; + for (uint j = 0u; j < p.cols; j++) m = max(m, float(x[b + j])); + float s = 0.0f; + for (uint j = 0u; j < p.cols; j++) s += exp(float(x[b + j]) - m); + float inv = 1.0f / s; + for (uint j = 0u; j < p.cols; j++) y[b + j] = half(exp(float(x[b + j]) - m) * inv); +} + +// ---------------- attn ---------------- +// Fused multi-head self-attention on packed qkv [n, T, 3, H, D] fp16: +// o[n, i, h, :] = sum_j softmax_j(scale * q_i . k_j) * v_j +// written as [n, T, H, D]. Online softmax, fp32 math; D <= 32 (zero-pad). +// Group = 64 queries of one (n, h); K/V staged through threadgroup tiles. +struct PcAttn { uint T, H, D; uint scaleBits; }; +kernel void attn(device const half* qkv [[buffer(0)]], + device half* o [[buffer(1)]], + constant PcAttn& p [[buffer(2)]], + uint2 lp2 [[thread_position_in_threadgroup]], + uint2 wg [[threadgroup_position_in_grid]]) +{ + uint lid = lp2.x; + threadgroup float ks[64 * 32]; + threadgroup float vs[64 * 32]; + + uint nh = wg.y; + uint n = nh / p.H, h = nh % p.H; + uint i = wg.x * 64u + lid; + uint hd = p.H * p.D; + uint rowStride = 3u * hd; + bool valid = i < p.T; + float sc = as_type(p.scaleBits); + + float q[32]; + float acc[32]; + uint qb = (n * p.T + (valid ? i : 0u)) * rowStride + h * p.D; + for (uint d = 0u; d < 32u; d++) { + q[d] = (valid && d < p.D) ? float(qkv[qb + d]) * sc : 0.0f; + acc[d] = 0.0f; + } + float m = -3.0e38f, l = 0.0f; + + for (uint j0 = 0u; j0 < p.T; j0 += 64u) { + for (uint e = lid; e < 64u * 32u; e += 64u) { + uint jr = e >> 5, d = e & 31u; + uint j = j0 + jr; + float kv = 0.0f, vv = 0.0f; + if (j < p.T && d < p.D) { + uint kb = (n * p.T + j) * rowStride + hd + h * p.D + d; + kv = float(qkv[kb]); + vv = float(qkv[kb + hd]); + } + ks[e] = kv; + vs[e] = vv; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint jn = min(64u, p.T - j0); + for (uint jr = 0u; jr < jn; jr++) { + float s = 0.0f; + for (uint d = 0u; d < 32u; d++) s += q[d] * ks[jr * 32u + d]; + float mn = max(m, s); + float corr = exp(m - mn); + float pe = exp(s - mn); + l = l * corr + pe; + for (uint d = 0u; d < 32u; d++) acc[d] = acc[d] * corr + pe * vs[jr * 32u + d]; + m = mn; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (valid) { + uint ob = (n * p.T + i) * hd + h * p.D; + float inv = 1.0f / l; + for (uint d = 0u; d < 32u; d++) + if (d < p.D) o[ob + d] = half(acc[d] * inv); + } +} diff --git a/src/Sdcb.SimdPaddleOCR/OcrBackend.cs b/src/Sdcb.SimdPaddleOCR/OcrBackend.cs index e75a882..9c2dcd1 100644 --- a/src/Sdcb.SimdPaddleOCR/OcrBackend.cs +++ b/src/Sdcb.SimdPaddleOCR/OcrBackend.cs @@ -9,4 +9,6 @@ public enum OcrBackend Cpu, /// Run the graph on Vulkan; falls back to CPU when no usable device exists. Vulkan, + /// Run the graph on Metal (macOS); falls back to CPU when no usable device exists. + Metal, } diff --git a/src/Sdcb.SimdPaddleOCR/OnnxSharp/OcrSessionFactory.cs b/src/Sdcb.SimdPaddleOCR/OnnxSharp/OcrSessionFactory.cs index 9127b00..c82b7ba 100644 --- a/src/Sdcb.SimdPaddleOCR/OnnxSharp/OcrSessionFactory.cs +++ b/src/Sdcb.SimdPaddleOCR/OnnxSharp/OcrSessionFactory.cs @@ -2,14 +2,17 @@ namespace Sdcb.SimdPaddleOCR.OnnxSharp; /// /// Backend-neutral session factory. On net10.0+ it can hand out a GPU-backed -/// session; on netstandard2.0 (or when the Vulkan device is absent/unusable) -/// it always yields the CPU interpreter. +/// session — Metal on macOS (unless overridden), Vulkan elsewhere; on +/// netstandard2.0 (or when no GPU device is usable) it always yields the CPU +/// interpreter. /// internal static class OcrSessionFactory { internal static IOcrSession Create(CompiledModel compiled, OcrBackend backend) { #if NET10_0_OR_GREATER + if (Backends.Metal.MetalBackend.IsMetalSelected(backend)) + return Backends.Metal.MetalBackend.CreateSession(compiled, backend); return Backends.Vulkan.GpuBackend.CreateSession(compiled, backend); #else _ = backend; @@ -21,6 +24,8 @@ internal static IOcrSession Create(CompiledModel compiled, OcrBackend backend) internal static bool IsGpuBackend(OcrBackend backend) { #if NET10_0_OR_GREATER + if (Backends.Metal.MetalBackend.IsMetalSelected(backend)) + return Backends.Metal.MetalBackend.TryGetDevice() is not null; return Backends.Vulkan.GpuBackend.UsesGpu(backend); #else _ = backend; diff --git a/src/Sdcb.SimdPaddleOCR/Sdcb.SimdPaddleOCR.csproj b/src/Sdcb.SimdPaddleOCR/Sdcb.SimdPaddleOCR.csproj index ae3688a..19d27c3 100644 --- a/src/Sdcb.SimdPaddleOCR/Sdcb.SimdPaddleOCR.csproj +++ b/src/Sdcb.SimdPaddleOCR/Sdcb.SimdPaddleOCR.csproj @@ -21,6 +21,8 @@ + + diff --git a/test/Sdcb.SimdPaddleOCR.GpuBench/PipelineBench.cs b/test/Sdcb.SimdPaddleOCR.GpuBench/PipelineBench.cs index b25f3d5..0ba2633 100644 --- a/test/Sdcb.SimdPaddleOCR.GpuBench/PipelineBench.cs +++ b/test/Sdcb.SimdPaddleOCR.GpuBench/PipelineBench.cs @@ -68,13 +68,10 @@ internal static int Run(string[] args) : Sdcb.SimdPaddleOCR.OnnxSharp.Model.Load(File.ReadAllBytes(clsPath)); using var ocr = new Sdcb.SimdPaddleOCR.PaddleOcrAll(detM, clsM, recM, File.ReadAllBytes(keysPath), opts); - using var bmp = new System.Drawing.Bitmap(imgPath); - var bd = bmp.LockBits(new System.Drawing.Rectangle(0, 0, bmp.Width, bmp.Height), - System.Drawing.Imaging.ImageLockMode.ReadOnly, - System.Drawing.Imaging.PixelFormat.Format24bppRgb); - byte[] bgr = new byte[bmp.Height * bd.Stride]; - Marshal.Copy(bd.Scan0, bgr, 0, bgr.Length); - bmp.UnlockBits(bd); + using var _im = SixLabors.ImageSharp.Image.Load(imgPath); + byte[] bgr = new byte[_im.Width * _im.Height * 3]; + _im.CopyPixelDataTo(bgr); + int imgW = _im.Width, imgH = _im.Height, imgStride = imgW * 3; bool prof = Environment.GetEnvironmentVariable("SIMD_OCR_PROF") == "1"; Sdcb.SimdPaddleOCR.PipelineProfiler.Enable(prof); @@ -82,7 +79,7 @@ internal static int Run(string[] args) { if (prof) Sdcb.SimdPaddleOCR.PipelineProfiler.Enable(true); // reset per rep var t0 = System.Diagnostics.Stopwatch.GetTimestamp(); - var r = ocr.Run(bgr, bmp.Width, bmp.Height, bd.Stride, + var r = ocr.Run(bgr, imgW, imgH, imgStride, Sdcb.SimdPaddleOCR.ImagePixelFormat.Bgr24); double ms = (System.Diagnostics.Stopwatch.GetTimestamp() - t0) / (double)System.Diagnostics.Stopwatch.Frequency * 1000; diff --git a/test/Sdcb.SimdPaddleOCR.Tests/IBenchEngine.cs b/test/Sdcb.SimdPaddleOCR.Tests/IBenchEngine.cs index 3567b8e..54a18ed 100644 --- a/test/Sdcb.SimdPaddleOCR.Tests/IBenchEngine.cs +++ b/test/Sdcb.SimdPaddleOCR.Tests/IBenchEngine.cs @@ -49,9 +49,10 @@ static class BenchEngines // the factory pick (Vulkan when a usable device exists). "sharp" => new SharpEngine(modelType, workers, OcrBackend.Cpu), "vulkan" => new SharpEngine(modelType, workers, OcrBackend.Vulkan), + "metal" => new SharpEngine(modelType, workers, OcrBackend.Metal), "auto" => new SharpEngine(modelType, workers, OcrBackend.Auto), "c" => new CEngine(cAssetsDir, workers, modelType), - _ => throw new ArgumentException("--engine must be sharp, vulkan, auto, or c"), + _ => throw new ArgumentException("--engine must be sharp, vulkan, metal, auto, or c"), }; public static PaddleOcrModelBundle Bundle(string modelType) => modelType switch diff --git a/test/Sdcb.SimdPaddleOCR.Tests/Program.cs b/test/Sdcb.SimdPaddleOCR.Tests/Program.cs index 97a0d65..3c42ad4 100644 --- a/test/Sdcb.SimdPaddleOCR.Tests/Program.cs +++ b/test/Sdcb.SimdPaddleOCR.Tests/Program.cs @@ -16,7 +16,7 @@ if (args.Length == 0 || args[0] is "-h" or "--help") { - Console.WriteLine("usage: --workers 1..16 --model tiny|small|medium --input --out [--engine sharp|vulkan|auto|c] [--smoke|--benchmark] [--benchmark-kind simd|engine] [--count N] [--warmup N] [--case-id ID] [--replica N] [--c-assets ]"); + Console.WriteLine("usage: --workers 1..16 --model tiny|small|medium --input --out [--engine sharp|vulkan|metal|auto|c] [--smoke|--benchmark] [--benchmark-kind simd|engine] [--count N] [--warmup N] [--case-id ID] [--replica N] [--c-assets ]"); Console.WriteLine(" --summarize [--input ] [--out-md ]"); return args.Length == 0 ? 2 : 0; } @@ -64,8 +64,8 @@ throw new ArgumentException("--workers must be 1..16"); if (modelType is not ("tiny" or "small" or "medium")) throw new ArgumentException("--model must be tiny, small, or medium"); -if (engineName is not ("sharp" or "vulkan" or "auto" or "c")) - throw new ArgumentException("--engine must be sharp, vulkan, auto, or c"); +if (engineName is not ("sharp" or "vulkan" or "metal" or "auto" or "c")) + throw new ArgumentException("--engine must be sharp, vulkan, metal, auto, or c"); if (benchmarkKind is not ("simd" or "engine")) throw new ArgumentException("--benchmark-kind must be simd or engine"); if (count is < 1 or > 100) From 5f5c5968345072d4b53292a97bdb15bc0ab1bd55 Mon Sep 17 00:00:00 2001 From: ZHOU Jie Date: Mon, 28 Sep 2026 04:48:44 -0700 Subject: [PATCH 2/8] =?UTF-8?q?docs:=20README=20=E8=AE=BE=E5=A4=87?= =?UTF-8?q?=E8=A1=8C=E5=88=97=E5=87=BA=20Vulkan/Metal=20GPU=20=E5=90=8E?= =?UTF-8?q?=E7=AB=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 65d24dc..85313a7 100644 --- a/README.md +++ b/README.md @@ -202,7 +202,7 @@ ImageSharp 默认分配器会把像素拆成 4MB 块，大图上拿不到一整 | CI 架构 | Windows x64 / x86 / ARM64，Linux x64 / ARM64，macOS x64 / ARM64 | | SIMD | .NET 10 运行时探测 AVX → AVX2 → AVX-512 / VNNI；无对应指令集或 ARM 时用 Vector/标量 | | 输入 | 交错像素内存（默认 BGR24，也可 RGB24 / BGRA32 / RGBA32）；无图片路径、文件或图片库 API | -| 设备 | 默认 CPU；`net10.0` 另有可选 Vulkan 后端（`OcrBackend`），直连系统加载器：Windows `vulkan-1.dll`、Linux `libvulkan.so.1`、Android `libvulkan.so`；`netstandard2.0` 只有 CPU | +| 设备 | 默认 CPU；`net10.0` 另有可选 GPU 后端（`OcrBackend`）：Vulkan 直连系统加载器（Windows `vulkan-1.dll`、Linux `libvulkan.so.1`、Android `libvulkan.so`），Metal 走 macOS 系统框架；`netstandard2.0` 只有 CPU | | 安卓 | 目前通过开发用宿主 `test/Sdcb.SimdPaddleOCR.AndroidBench`（`net10.0-android`，引用 `net10.0` 库，adb 驱动）在骁龙 8 Gen 3 上跑 CPU 与 Vulkan，见 [`docs/vulkan-8gen3.md`](docs/vulkan-8gen3.md)；桌面的 Vulkan 路由和 shader 没有变 | | NativeAOT | 裁剪发布时请保留核心程序集和所用模型程序集 | From 72ddcd42dc053bcebc140937a5904ee368388a43 Mon Sep 17 00:00:00 2001 From: ZHOU Jie Date: Mon, 28 Sep 2026 05:35:44 -0700 Subject: [PATCH 3/8] Address PR #21 review: fix icconv resT/dualT binds + scalarBias gate, device-level MSL lib/PSO cache, concurrent interleave encoder, GpuSessionBase dedupe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - mm_ic_* emit sites bound residual/dual tensors as (arena,0)/(arena,0,2) while flags could carry bit1/bit3 — bind the real slots, and gate the second site on scalarBias==0 (the scalar bias channel is repurposed as the dual-write flag in the mm kernels). - Auto no longer selects Metal on Intel Macs: Auto now requires Arm64+Apple7; explicit OcrBackend.Metal still allowed elsewhere via UsesGpu (probe validates simdgroup matrix live). - Move the embedded MSL library + ~40 PSOs to MtlDevice (OcrLibrary / OcrPipe): one compile per process instead of per MetalGraphModel, released on device Dispose. - Interleave now records through a concurrent compute encoder with one buffer barrier per level — under a serial encoder the per-rec barrier was a no-op and the level overlap unreachable. Serial path drops the dead per-rec barrier. - Extract GpuSessionBase + IOcrGraphRunner: GpuSession/MetalSession (identical 324-line bodies) become thin constructors over the shared session logic; GpuUnitsReady/MetalUnitsReady unify on CtcUnitsReady. - ProbeSimdMatrix now verifies PSO geometry (execWidth=32, maxTg>=256) and compares a full 64x64 tile against a reference GEMM. - Comment/doc nits: correct DWA/NOIC direction comments, mm_sg_dd gate comment, se two-phase explanation (relaxed atomics are universal MSL, not paravirt-specific), mm_ic pad-k semantics, MtlDataType UInt=30 fix (was Int2), MTLCreateSystemDefaultDevice +1 rule, doc CER metric and tool references. --- docs/metal-m4.md | 14 +- .../Backends/GpuSessionBase.cs | 324 ++++++++++++++++++ .../Backends/IOcrGraphRunner.cs | 32 ++ .../Backends/Metal/Metal.cs | 55 ++- .../Backends/Metal/MetalBackend.cs | 25 +- .../Backends/Metal/MetalGraph.cs | 181 +++++----- .../Backends/Metal/MetalRunner.cs | 30 +- .../Backends/Metal/MetalSession.cs | 320 +---------------- .../Backends/Metal/Shaders/gemm.metal | 18 +- .../Backends/Metal/Shaders/xform.metal | 15 +- .../Backends/Vulkan/GpuRunner.cs | 7 +- .../Backends/Vulkan/GpuSession.cs | 321 +---------------- .../OnnxSharp/OcrSessionFactory.cs | 2 +- 13 files changed, 564 insertions(+), 780 deletions(-) create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/GpuSessionBase.cs create mode 100644 src/Sdcb.SimdPaddleOCR/Backends/IOcrGraphRunner.cs diff --git a/docs/metal-m4.md b/docs/metal-m4.md index 429ac45..455da70 100644 --- a/docs/metal-m4.md +++ b/docs/metal-m4.md @@ -1,6 +1,6 @@ # SimdPaddleOCR — Apple Silicon Metal 后端实测（Devin VM / M4 paravirt） -实测机：macOS VM（Apple M4 paravirt GPU，8 vCPU），SDK .NET 10，`test/Sdcb.SimdPaddleOCR.Tests` bench（`--workers 4 --benchmark-kind simd --engine sharp|metal`，warmup=1，n=99，dataset/ 100 张固定图、张张不同尺寸）。commit：working tree @ feature/2.0 + metal-backend 工作区（未提交）。 +实测机：macOS VM（Apple M4 paravirt GPU，8 vCPU），SDK .NET 10，`test/Sdcb.SimdPaddleOCR.Tests` bench（`--workers 4 --benchmark-kind simd --engine sharp|metal`，warmup=1，n=99，dataset/ 100 张固定图、张张不同尺寸）。commit：`metal-backend` 分支（PR #21），性能数据采集于 `680bf38`（其后的 review 修订不动热路径）。 **结论：三档 Metal 全部净胜 CPU（tiny 1.9× / small 4.7× / medium 5.2×），检测行数与 CPU 完全一致，文本差异仅 fp16 级噪声（3 图各 1 字符）。受 paravirt 虚拟化所限绝对数偏保守——实测 MMA≈fp32≈3TFLOPS、copy ~84GB/s、dispatch ~45µs，均低于真机 M4，真机数字预期更好。** @@ -28,7 +28,7 @@ | medium | cls_graph | 11.9 | 2.61 | 4.6× | | medium | rec_graph | 2322.2 | 112.6 | 20.6× | -精度：metal/sharp 检出 lines 完全一致（1016/1009/1013）；逐图 hash 3 张差（fp16 阈值临界像素，与 Vulkan 同型已知限制）；文本 3 图各差 1 字符（CER 0.67% vs CPU 0.14%，在 fp16 噪声范围内）。`--conc`（4 线程 × 72 runs/模型）0 mismatch；`--dtr`（det medium 228 节点 × 8 跑）逐节点 bitwise 一致。 +精度：metal/sharp 检出 lines 完全一致（1016/1009/1013）；逐图 hash 3 张差（fp16 阈值临界像素，与 Vulkan 同型已知限制）；文本 3 图各差 1 字符（CER 口径：对 dataset 真值标注的字符错误率——Metal 0.67% vs CPU 0.14%，在 fp16 噪声范围内）。并发正确性：`--conc`（4 线程 × 72 runs/模型）0 mismatch；det medium 228 节点逐节点对拍（本地 probe 工具，未提交）× 8 跑 bitwise 一致。 ## 内存与分配 @@ -40,7 +40,7 @@ GPU 侧 arena 是共享 grow-only device buffer + 托管 schedule（MetalSchedule 纯数据 LRU=512，逐 shape 录制代价为零），与 Vulkan 的 plan/arena 形态同构。托管堆分配 ~1.4–2.5MB/图。 -## GPU 能力实测（paravirt VM，--mprobe） +## GPU 能力实测（paravirt VM microbench） | 项 | 实测 | |---|---| @@ -59,7 +59,7 @@ GPU 侧 arena 是共享 grow-only device buffer + 托管 schedule（MetalSchedul |---|---|---| | `mm_sg_dd` | 64×64 MMA tile 直接从 device 内存 simdgroup_load，省掉 threadgroup staging 往返；emit gate `K%16==0`（边界 tile 的越界读被 arena ≥1MB 尾 slack + W /128 pad + epilogue 掩码兜住，数学不变） | 5760×512×1024 GEMM 2.63→2.31ms；det tiny 8.6→6.8ms、small 20.8→17.7ms | | `mm_ic_sg`（隐式 GEMM 卷积） | kxk conv 不再物化 `[M,Kp]` im2col 矩阵：A staging 内联 tap 寻址 gather；同时接管 nb==1 的 convk_dot 分支 | det medium 960² 86.2→**80.0ms**（-7%）；消掉 ~1.28GB 无谓读写 | -| SE 两阶段定稿 | `se_part`（S WG 写 partial，零原子）+ buffer barrier + `se_join`（1 WG/batch）| 修掉跨 WG ticket 竞态（paravirt MSL 无 acq_rel 原子，握手不可能）；`--dtr` 全绿 | +| SE 两阶段定稿 | `se_part`（S WG 写 partial，零原子）+ buffer barrier + `se_join`（1 WG/batch）| 修掉跨 WG ticket 竞态（MSL 原子仅 relaxed 序，跨 WG 握手不可能）；det 逐节点对拍全绿 | | conv_dw4a flat | 平铺 depthwise 省 threadgroup staging | **未采纳**：实测比 conv_dw4t 慢（staging 省 halo 重读仍有收益）→ 留 `SIMD_OCR_DWA` 开关 | | `mm_ic_sg32`（BK=32） | staging barrier 减半 | **未采纳**：medium −1ms 但 tiny/small 更差 → 留 `SIMD_OCR_IC32` 开关 | | NODOT（conv1x1 全 MMA） | 关 dot kernel | 与现状持平 → 阈值保持 B580 标定值 | @@ -68,6 +68,6 @@ GPU 侧 arena 是共享 grow-only device buffer + 托管 schedule（MetalSchedul ## 已知限制 / 未验证项 - DET bs>1 resize 不支持（同 Vulkan）。 -- paravirt 独有：MSL `memory_order` 只有 relaxed → 跨 WG 定序走两阶段 split + `memoryBarrierWithScope:Buffers`；`char16/uchar16` 不可用。 -- 序列化 `--prof` 有 ~0.4–0.5ms/dispatch 地板，绝对值偏大；真实比例看 `--time` 墙钟与 `--mgemm` N-reps-in-one-CB。 -- 无 CI Metal 依赖；单元测试在非 macOS 上 SKIP。 +- MSL 通用限制：device `memory_order` 只有 relaxed（非 paravirt 独有）→ 跨 WG 定序走两阶段 split + `memoryBarrierWithScope:Buffers`；paravirt 另限 `char16/uchar16` 不可用。 +- 序列化逐 dispatch profile（`SIMD_OCR_GPU_PROF`）有 ~0.4–0.5ms/dispatch 地板，绝对值偏大；真实比例看 `SIMD_OCR_GPU_TIME` 的墙钟分解。 +- 无 CI Metal 依赖：Metal 路径只在 macOS 上被选中，其他 OS 的测试和 bench 不会走到它。 diff --git a/src/Sdcb.SimdPaddleOCR/Backends/GpuSessionBase.cs b/src/Sdcb.SimdPaddleOCR/Backends/GpuSessionBase.cs new file mode 100644 index 0000000..3bfce31 --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/GpuSessionBase.cs @@ -0,0 +1,324 @@ +using System.Buffers.Binary; +using System.Runtime.InteropServices; +using Sdcb.SimdPaddleOCR.OnnxSharp; + +namespace Sdcb.SimdPaddleOCR.Backends; + +/// +/// GPU-backed shared by the Vulkan and Metal +/// backends: the whole ONNX graph runs through an . +/// InputData is a managed staging span uploaded on each run; the CTC tail +/// matches 's contract (stop before the vocab +/// MatMul, hand operands to the shared CPU ArgMax). One device is shared +/// process-wide; weights/pipelines/schedules are shared per model, while each +/// session owns its buffers, so sessions only serialize at the submissions. +/// +internal abstract class GpuSessionBase : IOcrSession, IBatchedCtcSession +{ + private readonly IOcrGraphRunner _runner; + private readonly string _logTag; + private readonly CompiledModel _compiled; + private readonly Model _model; + private readonly ResizeWorkspace _resizeWorkspace = new(); + private InferenceSession? _cpuFallback; + private float[] _cpuNhwc = []; // NCHW→NHWC staging for the CPU fallback + + private int[] _shape = []; + private float[] _input = []; + private int _curVolume; // numel of the current shape; _input may be bigger + private int _hwVolume; + private bool _disposed; + private bool _gpuDead; // plan/run failure → serve everything from _cpuFallback + + // CTC projection resolution, cached per Reshape + private bool _ctcResolved; + private int _ctcMatMulIndex = -1; + private int _ctcActTensor; + private int _ctcBatch, _ctcRows, _ctcInner, _ctcColumns; + private byte[]? _ctcWBytes, _ctcBiasBytes; + private float[]? _ctcPacked; + + private protected GpuSessionBase(IOcrGraphRunner runner, CompiledModel compiled, string logTag) + { + _runner = runner; + _compiled = compiled; + _model = compiled.Model; + _logTag = logTag; + } + + public TensorShape InputShape => new(_shape); + private int[]? _outputShape; + public TensorShape OutputShape => + new(_outputShape ??= _compiled.ResolveShapesFor(_shape)[checked((int)_model.GraphOutputs[0])]); + public Span InputData => _input.AsSpan(0, _curVolume); + public ResizeWorkspace ResizeWorkspace => _resizeWorkspace; + public bool InputIsNhwc => false; // GPU conv kernels read NCHW fp32 directly + public int IntraOpThreads { get; set; } + public bool GpuAlive => !_gpuDead; + public int HighWaterInputVolume => _hwVolume; + public bool PlanForCtcProjection { get; set; } + public bool IsProfilingEnabled => InferenceSession.ProfilingEnabled; + + public void Reshape(ReadOnlySpan inputShape) + { + if (_disposed) throw new ObjectDisposedException(GetType().Name); + _shape = inputShape.ToArray(); + long vol = 1; + foreach (int d in _shape) vol *= Math.Max(d, 1); + _curVolume = checked((int)vol); + if (vol > _input.Length) + { + _input = new float[vol]; + _hwVolume = checked((int)Math.Min(vol, int.MaxValue)); + } + _outputShape = null; + _ctcResolved = false; + _ctcMatMulIndex = -1; + } + + public ReadOnlySpan RunInternal(ReadOnlySpan input) + { + if (_disposed) throw new ObjectDisposedException(GetType().Name); + if (_gpuDead) return Cpu().RunInternal(CpuInput(input)); + try + { + return _runner.Run(_shape, input); + } + catch (Exception ex) + { + Console.Error.WriteLine($"[{_logTag}] RunInternal fallback: {ex.GetType().Name} {ex.Message}"); + _gpuDead = true; + return Cpu().RunInternal(CpuInput(input)); + } + } + + public bool TryRunUntilCtcProjection(ReadOnlySpan input, out CtcProjectionOperands operands) + { + operands = default; + if (_disposed) throw new ObjectDisposedException(GetType().Name); + if (_gpuDead) return Cpu().TryRunUntilCtcProjection(CpuInput(input), out operands); + ResolveCtcProjection(); + if (_ctcMatMulIndex < 0) return false; + + ReadOnlySpan act; + try + { + act = _runner.Run(_shape, input, _ctcMatMulIndex, _ctcActTensor); + } + catch (Exception ex) + { + Console.Error.WriteLine($"[{_logTag}] CtcProj fallback: {ex.GetType().Name} {ex.Message}"); + _gpuDead = true; + return Cpu().TryRunUntilCtcProjection(CpuInput(input), out operands); + } + operands = new CtcProjectionOperands(act, + MemoryMarshal.Cast(_ctcWBytes), + _ctcBiasBytes is null ? [] : MemoryMarshal.Cast(_ctcBiasBytes), + _ctcPacked, _ctcBatch, _ctcRows, _ctcInner, _ctcColumns, _ctcMatMulIndex); + return true; + } + + // ---- IBatchedCtcSession: several REC shapes, one GPU submission ---- + private int[][] _many = []; + private int _manyVolume; + + public bool CanRunMany => !_gpuDead && !_disposed; + + public Span ReshapeMany(IReadOnlyList shapes) + { + if (_disposed) throw new ObjectDisposedException(GetType().Name); + _many = [.. shapes]; + long vol = 0; + foreach (int[] s in _many) + { + long v = 1; + foreach (int d in s) v *= Math.Max(d, 1); + vol += v; + } + _manyVolume = checked((int)vol); + if (vol > _input.Length) + { + _input = new float[vol]; + _hwVolume = _manyVolume; + } + // single-shape members (CTC resolution, fallback) key off the first unit + Reshape(_many[0]); + return _input.AsSpan(0, _manyVolume); + } + + public bool TryResolveManyHead(out int[] rows, out CtcHead head) + { + rows = []; head = null!; + if (!CanRunMany || _many.Length == 0) return false; + ResolveCtcProjection(); + if (_ctcMatMulIndex < 0) return false; + rows = new int[_many.Length]; + for (int i = 0; i < _many.Length; i++) + rows[i] = _compiled.ResolveShapesFor(_many[i])[_ctcActTensor][^2]; + head = new CtcHead + { + Weights = _ctcWBytes!, Bias = _ctcBiasBytes, Packed = _ctcPacked, + Inner = _ctcInner, Columns = _ctcColumns, MatMulIndex = _ctcMatMulIndex, + }; + return true; + } + + public bool TryRunManyUntilCtcProjection(CtcUnitsReady onReady) + { + if (!CanRunMany || _many.Length == 0) return false; + ResolveCtcProjection(); + if (_ctcMatMulIndex < 0) return false; + // consumer failures are the caller's, not a GPU fault: keep them out + // of the fallback catch below and rethrow once the device is drained + System.Runtime.ExceptionServices.ExceptionDispatchInfo? consumerError = null; + bool Forward(float[] acts, int[] offsets, int first, int count) + { + try { return onReady(acts, offsets, first, count); } + catch (Exception ex) + { + consumerError = System.Runtime.ExceptionServices.ExceptionDispatchInfo.Capture(ex); + return false; + } + } + bool ok; + try + { + ok = _runner.RunMany(_many, _input.AsSpan(0, _manyVolume), _ctcMatMulIndex, _ctcActTensor, Forward); + } + catch (Exception ex) + { + Console.Error.WriteLine($"[{_logTag}] RunMany fallback: {ex.GetType().Name} {ex.Message}"); + _gpuDead = true; + return false; + } + consumerError?.Throw(); + return ok; + } + + private InferenceSession Cpu() + { + _cpuFallback ??= new InferenceSession(_compiled) + { PlanForCtcProjection = PlanForCtcProjection }; + _cpuFallback.Reshape(_shape); + if (IntraOpThreads > 0) _cpuFallback.IntraOpThreads = IntraOpThreads; + return _cpuFallback; + } + + /// Graphs without the CTC tail fall back to a lazy CPU session. + public ReadOnlySpan RunInternalSkipFinalSoftmax(ReadOnlySpan input, out bool outputIsLogits) + { + if (_disposed) throw new ObjectDisposedException(GetType().Name); + return Cpu().RunInternalSkipFinalSoftmax(CpuInput(input), out outputIsLogits); + } + + /// + /// Callers always write logical NCHW (this session reports + /// = false). When the compiled model flags its + /// graph input as NHWC-physical, the CPU fallback must be fed the + /// transposed layout or it decodes garbage. + /// + private ReadOnlySpan CpuInput(ReadOnlySpan input) + { + if (!_compiled.InputIsNhwc) return input; + if (_cpuNhwc.Length < input.Length) _cpuNhwc = new float[input.Length]; + Span dest = _cpuNhwc.AsSpan(0, input.Length); + Kernels.Nhwc.NchwToNhwc(input, dest, _shape[0], _shape[1], + checked(_shape[2] * _shape[3]), Math.Max(IntraOpThreads, 1)); + return dest; + } + + public void NoteProfile(OperatorId operation, long started, int nodeIndex) + => InferenceSession.NoteProfileStatic(operation, started, nodeIndex); + + /// + /// Mirrors InferenceSession.TryResolveCtcProjection but resolves + /// shapes from instead of live + /// tensor bindings — the GPU session has no CPU workspace. + /// + private void ResolveCtcProjection() + { + if (_ctcResolved) return; + _ctcResolved = true; + _ctcMatMulIndex = -1; + + NodeRecord[] nodes = _model.Nodes; + if (nodes.Length < 2) return; + int[][] shapes = _compiled.ResolveShapesFor(_shape); + int outIdx = checked((int)_model.GraphOutputs[0]); + + NodeRecord softmax = nodes[^1]; + if (softmax.Operator != OperatorId.Softmax || + softmax.Inputs.Length != 1 || softmax.Outputs.Length != 1 || + softmax.Outputs[0] != (uint)outIdx) return; + int axis = I32(_model.GetParameters(softmax), 4); + int[] smInShape = shapes[softmax.Inputs[0]]; + if (axis < 0) axis += smInShape.Length; + if (axis != smInShape.Length - 1) return; + + int terminalIndex = nodes.Length - 2; + NodeRecord terminal = nodes[terminalIndex]; + int matMulIndex; + byte[]? biasBytes = null; + if (terminal.Operator == OperatorId.MatMul && + terminal.Outputs[0] == softmax.Inputs[0]) + { + matMulIndex = terminalIndex; + } + else if (terminal.Operator == OperatorId.Add && terminal.Inputs.Length == 2 && + terminalIndex > 0 && terminal.Outputs[0] == softmax.Inputs[0]) + { + matMulIndex = terminalIndex - 1; + NodeRecord mm = nodes[matMulIndex]; + if (mm.Operator != OperatorId.MatMul || mm.Outputs.Length != 1) return; + uint biasIndex; + if (terminal.Inputs[0] == mm.Outputs[0]) biasIndex = terminal.Inputs[1]; + else if (terminal.Inputs[1] == mm.Outputs[0]) biasIndex = terminal.Inputs[0]; + else return; + var biasMeta = _compiled.GetTensor(checked((int)biasIndex)); + int[] projShape = shapes[mm.Outputs[0]]; + if (!biasMeta.IsConstant || biasMeta.Shape.Length != 1 || + projShape.Length < 2 || biasMeta.Shape[0] != projShape[^1]) return; + biasBytes = biasMeta.Constant; + } + else return; + + NodeRecord matMul = nodes[matMulIndex]; + if (matMul.Inputs.Length != 2 || matMul.Outputs.Length != 1) return; + var aMeta = _compiled.GetTensor(checked((int)matMul.Inputs[0])); + var bMeta = _compiled.GetTensor(checked((int)matMul.Inputs[1])); + int[] aShape = shapes[matMul.Inputs[0]]; + int[] bShape = shapes[matMul.Inputs[1]]; + int[] projOut = shapes[matMul.Outputs[0]]; + if (!bMeta.IsConstant || aShape.Length < 2 || bShape.Length != 2 || + projOut.Length < 2) return; + int rows = aShape[^2], inner = aShape[^1], columns = bShape[1]; + long aCount = 1; + foreach (int d in aShape) aCount *= Math.Max(d, 1); + if (rows <= 0 || inner <= 0 || aCount % (rows * (long)inner) != 0) return; + int batch = checked((int)(aCount / (rows * (long)inner))); + if (bShape[0] != inner) return; + long projCount = 1; + foreach (int d in projOut) projCount *= Math.Max(d, 1); + if (projOut[^1] != columns || projCount != (long)batch * rows * columns) return; + _compiled.TryGetPackedMatMul(matMul.Inputs[1], out float[]? packed); + // Column count does not gate here (unlike the recognizer's ArgMax + // fusion): small-column heads like the cls [inner,2] tail resolve too. + if (columns <= 0) return; + + _ctcMatMulIndex = matMulIndex; + _ctcActTensor = checked((int)matMul.Inputs[0]); + _ctcBatch = batch; _ctcRows = rows; _ctcInner = inner; _ctcColumns = columns; + _ctcWBytes = bMeta.Constant; + _ctcBiasBytes = biasBytes; + _ctcPacked = packed; + } + + private static int I32(ReadOnlySpan p, int o) => BinaryPrimitives.ReadInt32LittleEndian(p[o..]); + + public void Dispose() + { + _disposed = true; + _runner.Dispose(); + _cpuFallback?.Dispose(); + } +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/IOcrGraphRunner.cs b/src/Sdcb.SimdPaddleOCR/Backends/IOcrGraphRunner.cs new file mode 100644 index 0000000..21e2f63 --- /dev/null +++ b/src/Sdcb.SimdPaddleOCR/Backends/IOcrGraphRunner.cs @@ -0,0 +1,32 @@ +using Sdcb.SimdPaddleOCR.OnnxSharp; + +namespace Sdcb.SimdPaddleOCR.Backends; + +/// +/// Per-session GPU executor behind : one runner +/// instance owns the session's device buffers and turns the shared, +/// model-level schedules into submissions. Implemented by the Vulkan +/// GpuDetGraph and the Metal MetalDetGraph. +/// +internal interface IOcrGraphRunner : IDisposable +{ + /// + /// Run the graph: input fp32 NCHW [n,C,H,W] → fp32 [graph output]. + /// truncates the emit loop (nodes beyond it + /// are not dispatched) and overrides the + /// readback tensor — used to stop the REC graph before the vocab + /// projection (activations readback) instead of the graph output. + /// + ReadOnlySpan Run(int[] inputShape, ReadOnlySpan input, + int nodeLimit = int.MaxValue, int outTensor = -1); + + /// + /// Runs several shapes of the same graph (inputs back to back in + /// , unit i = numel(shapes[i]) floats) and hands + /// the activations to in unit order as they + /// land. Returns false when declined a batch; + /// every submission has completed by the time this returns or throws. + /// + bool RunMany(IReadOnlyList shapes, ReadOnlySpan input, + int nodeLimit, int outTensor, CtcUnitsReady onReady); +} diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Metal.cs b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Metal.cs index f444c0a..94f2c10 100644 --- a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Metal.cs +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Metal.cs @@ -12,19 +12,23 @@ namespace Sdcb.SimdPaddleOCR.Backends.Metal; // is CPU- and GPU-visible without didModifyRange/coherence calls. internal static class MtlStorageMode { public const nuint Shared = 0; } -// MTLDataType values (MTLDataType.h) +// MTLDataType values (MTLDataType.h) — only what NewPipeline's intConstants +// path uses. internal static class MtlDataType { - public const nuint Float = 3; public const nuint Int = 29; - public const nuint UInt = 30; - public const nuint Bool = 53; } internal sealed class MtlDevice : IDisposable { public readonly IntPtr H; private readonly IntPtr _queue; + // OCR shader library + PSO cache, shared by every MetalGraphModel on this + // device (each PaddleOcrAll carries det/cls/rec models — per-model + // libraries would recompile ~2.5k lines of MSL three times and leak the + // whole set on every release). Disposed with the device. + private MtlLibrary? _ocrLib; + private readonly Dictionary _ocrPipes = new(); private MtlDevice(IntPtr h) { @@ -67,10 +71,41 @@ private MtlDevice(IntPtr h) { if (!OperatingSystem.IsMacOS()) return null; using var pool = AutoReleasePool.Create(); + // MTLCreateSystemDefaultDevice follows the Create rule: already +1. IntPtr h = ObjC.CreateSystemDefaultDevice(); - if (h == 0) return null; - ObjC.Retain(h); // C-function result is autoreleased; promote to +1. - return new MtlDevice(h); + return h == 0 ? null : new MtlDevice(h); + } + + /// Pipeline for a kernel in the embedded OCR shader library, + /// compiled once and cached per device. + internal MtlPipeline OcrPipe(string name) + { + lock (_ocrPipes) + { + if (!_ocrPipes.TryGetValue(name, out MtlPipeline? p)) + _ocrPipes[name] = p = NewPipeline(OcrLibrary(), name); + return p; + } + } + + private MtlLibrary OcrLibrary() + { + if (_ocrLib is not null) return _ocrLib; + var asm = typeof(MtlDevice).Assembly; + var names = asm.GetManifestResourceNames() + .Where(n => n.StartsWith("Sdcb.SimdPaddleOCR.Backends.Metal.Shaders.") + && n.EndsWith(".metal")) + .OrderBy(n => n).ToArray(); + if (names.Length == 0) + throw new FileNotFoundException("Metal: no embedded .metal shaders"); + var sb = new System.Text.StringBuilder(); + foreach (var n in names) + { + using Stream s = asm.GetManifestResourceStream(n)!; + using var r = new StreamReader(s); + sb.AppendLine(r.ReadToEnd()); + } + return _ocrLib = NewLibrary(sb.ToString()); } public MtlBuffer NewBuffer(nuint bytes) @@ -182,6 +217,12 @@ public MtlEncoder BeginCommands(bool concurrent) public void Dispose() { + lock (_ocrPipes) + { + foreach (MtlPipeline p in _ocrPipes.Values) p.Dispose(); + _ocrPipes.Clear(); + } + _ocrLib?.Dispose(); ObjC.Release(_queue); ObjC.Release(H); } diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalBackend.cs b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalBackend.cs index 050d75e..16c69db 100644 --- a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalBackend.cs +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalBackend.cs @@ -1,3 +1,4 @@ +using System.Runtime.InteropServices; using Sdcb.SimdPaddleOCR.OnnxSharp; namespace Sdcb.SimdPaddleOCR.Backends.Metal; @@ -38,15 +39,33 @@ internal static class MetalBackend { } s when s.Equals("metal", StringComparison.OrdinalIgnoreCase) => true, { } s when s.Equals("cpu", StringComparison.OrdinalIgnoreCase) || s.Equals("vulkan", StringComparison.OrdinalIgnoreCase) => false, - // Auto: Metal is the GPU of choice on macOS only - _ => OperatingSystem.IsMacOS(), + // Auto: Apple Silicon only — the GEMM kernels are built around + // 32-lane simdgroups and were never validated on Intel/AMD GPUs. + _ => OperatingSystem.IsMacOS() + && RuntimeInformation.ProcessArchitecture == Architecture.Arm64, }, }; + /// Metal asked for by name: the option, or Auto with + /// SIMD_OCR_BACKEND=metal. + private static bool IsMetalExplicit(OcrBackend backend) => + backend == OcrBackend.Metal + || backend == OcrBackend.Auto + && string.Equals(Environment.GetEnvironmentVariable("SIMD_OCR_BACKEND"), + "metal", StringComparison.OrdinalIgnoreCase); + + /// Whether a session for this option runs on Metal. Without an + /// Apple GPU family only an explicit choice gets Metal: the simdgroup + /// path is unproven off Apple Silicon, so Auto stays on CPU there + /// (mirrors ). + internal static bool UsesGpu(OcrBackend backend) => + IsMetalSelected(backend) && TryGetDevice() is { } dev + && (dev.SupportsApple7 || IsMetalExplicit(backend)); + /// Creates a session on Metal; CPU on any failure. internal static IOcrSession CreateSession(CompiledModel compiled, OcrBackend backend) { - if (IsMetalSelected(backend) && TryGetDevice() is { } dev) + if (UsesGpu(backend) && TryGetDevice() is { } dev) { try { return new MetalSession(dev, compiled); } catch { /* fall through to CPU */ } diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalGraph.cs b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalGraph.cs index dae67a8..336c01a 100644 --- a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalGraph.cs +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalGraph.cs @@ -6,10 +6,9 @@ namespace Sdcb.SimdPaddleOCR.Backends.Metal; /// -/// -/// Metal counterpart of the Metal MetalGraphModel: shared per (device, -/// compiled model); owns the embedded-MSL library pipelines, fp16 weight -/// uploads and the graph compiler that emits a +/// Metal counterpart of the Vulkan GpuGraphModel: shared per (device, +/// compiled model); owns the fp16 weight uploads and the graph compiler that +/// emits a /// — pure managed metadata (pipeline, arena-relative binding offsets, push /// constants, grid), no Metal object per shape. Sessions /// () own the grow-only arena/IO buffers and @@ -70,8 +69,9 @@ internal void Release() _pResize, _pResize4, _pResize4Add, _pConcat, _pConcat4, _pNchw, _pOut, _pIm2col, _pAddPs, _pSePart, _pSeJoin, _pConvD, _pConvDF32, _pCatRes, _pAvg4, _pAffine, _pLn, _pAttn, _pDotSk; - // Flat (unstaged) depthwise — threadgroup staging buys nothing on this - // unified-memory GPU; SIMD_OCR_NODWA=1 keeps the tiled variant for A/B. + // Flat (unstaged) depthwise — opt-in only: measured slower than + // conv_dw4t on this GPU (staging still pays off), SIMD_OCR_DWA enables it + // for A/B. private readonly MtlPipeline? _pDw4A; // fp16 copies of constant tensors (weights, biases, scalars), lazy @@ -180,60 +180,58 @@ private MetalGraphModel(MtlDevice dev, CompiledModel compiled) _dev = dev; _compiled = compiled; _model = compiled.Model; - _lib = LoadLibrary(); // one probe per capability tier: the simdgroup-matrix kernels get a // PSO + smoke dispatch; failure downgrades the cm paths to the // portable dot kernels (correct, just slower). _sgm = ProbeSimdMatrix(); - _pMm = Pipe("mm_sg"); - _pMmPs = Pipe("mm_ps_sg"); - _pMmDot = Pipe("mm_dot"); // portable cm-path fallback - _pMmDd = Pipe("mm_sg_dd"); // direct fragment loads, no staging - _pMmIc = Pipe("mm_ic_sg"); // implicit-GEMM kxk conv (tap gather in A stage) + _pMm = _dev.OcrPipe("mm_sg"); + _pMmPs = _dev.OcrPipe("mm_ps_sg"); + _pMmDot = _dev.OcrPipe("mm_dot"); // portable cm-path fallback + _pMmDd = _dev.OcrPipe("mm_sg_dd"); // direct fragment loads, no staging + _pMmIc = _dev.OcrPipe("mm_ic_sg"); // implicit-GEMM kxk conv (tap gather in A stage) _pMmIc32 = Environment.GetEnvironmentVariable("SIMD_OCR_IC32") != null - ? Pipe("mm_ic_sg32") : _pMmIc; // measured: BK=32 loses except det-medium + ? _dev.OcrPipe("mm_ic_sg32") : _pMmIc; // measured: BK=32 loses except det-medium if (!_sgm) { _pMm = _pMmDot; _pMmPs = _pMmDot; _pMmDd = _pMmDot; _pMmIc = _pMmDot; _pMmIc32 = _pMmDot; } - _pDot = Pipe("conv1x1_dot"); - _pDw = Pipe("conv_dw"); - _pDw4 = Pipe("conv_dw4"); - _pDw4T = Pipe("conv_dw4t"); + _pDot = _dev.OcrPipe("conv1x1_dot"); + _pDw = _dev.OcrPipe("conv_dw"); + _pDw4 = _dev.OcrPipe("conv_dw4"); + _pDw4T = _dev.OcrPipe("conv_dw4t"); _pDw4A = Environment.GetEnvironmentVariable("SIMD_OCR_DWA") != null - ? Pipe("conv_dw4a") : null; // measured: flat loses to conv_dw4t - _pDense = Pipe("conv_dense"); - _pIm2col = Pipe("im2col"); - _pConvT = Pipe("convt2s2"); - _pConvT4 = Pipe("convt4"); - _pReduce4 = Pipe("reduce_hw4"); - _pReduce4b = Pipe("reduce_hw4b"); - _pPool4 = Pipe("maxpool4"); - _pResize4 = Pipe("resize4"); - _pResize4Add = Pipe("resize4add"); - _pConcat4 = Pipe("concat4"); - _pCatRes = Pipe("catresize"); - _pElem = Pipe("elem"); - _pElem4 = Pipe("elem4"); - _pAddPs = Pipe("addps"); - _pSePart = Pipe("se_part"); - _pSeJoin = Pipe("se_join"); - _pConvD = Pipe("convk_dot"); - _pConvDF32 = Pipe("convk_f32n"); - _pReduce = Pipe("reduce_hw"); - _pPool = Pipe("maxpool2e"); - _pResize = Pipe("resize_nn"); - _pConcat = Pipe("concat_c"); - _pNchw = Pipe("nchw2nhwc"); - _pOut = Pipe("sigmoid_out"); - _pAvg4 = Pipe("avgpool4"); - _pAffine = Pipe("affine4"); - _pSoftmax = Pipe("softmax"); - _pLn = Pipe("layernorm"); - _pAttn = Pipe("attn"); - _pDotSk = Pipe("mmdot_sk"); + ? _dev.OcrPipe("conv_dw4a") : null; // measured: flat loses to conv_dw4t + _pDense = _dev.OcrPipe("conv_dense"); + _pIm2col = _dev.OcrPipe("im2col"); + _pConvT = _dev.OcrPipe("convt2s2"); + _pConvT4 = _dev.OcrPipe("convt4"); + _pReduce4 = _dev.OcrPipe("reduce_hw4"); + _pReduce4b = _dev.OcrPipe("reduce_hw4b"); + _pPool4 = _dev.OcrPipe("maxpool4"); + _pResize4 = _dev.OcrPipe("resize4"); + _pResize4Add = _dev.OcrPipe("resize4add"); + _pConcat4 = _dev.OcrPipe("concat4"); + _pCatRes = _dev.OcrPipe("catresize"); + _pElem = _dev.OcrPipe("elem"); + _pElem4 = _dev.OcrPipe("elem4"); + _pAddPs = _dev.OcrPipe("addps"); + _pSePart = _dev.OcrPipe("se_part"); + _pSeJoin = _dev.OcrPipe("se_join"); + _pConvD = _dev.OcrPipe("convk_dot"); + _pConvDF32 = _dev.OcrPipe("convk_f32n"); + _pReduce = _dev.OcrPipe("reduce_hw"); + _pPool = _dev.OcrPipe("maxpool2e"); + _pResize = _dev.OcrPipe("resize_nn"); + _pConcat = _dev.OcrPipe("concat_c"); + _pNchw = _dev.OcrPipe("nchw2nhwc"); + _pOut = _dev.OcrPipe("sigmoid_out"); + _pAvg4 = _dev.OcrPipe("avgpool4"); + _pAffine = _dev.OcrPipe("affine4"); + _pSoftmax = _dev.OcrPipe("softmax"); + _pLn = _dev.OcrPipe("layernorm"); + _pAttn = _dev.OcrPipe("attn"); + _pDotSk = _dev.OcrPipe("mmdot_sk"); } - private readonly MtlLibrary _lib; private readonly bool _sgm; // simdgroup_matrix 8x8 fp16 MMA verified live // sg32-family fields kept for emit parity with the Vulkan schedule // (DwTiled uses the sg32 kernel cap; PackSlabs runs unconditionally here) @@ -241,36 +239,19 @@ private MetalGraphModel(MtlDevice dev, CompiledModel compiled) private readonly bool _lite = false; private readonly MtlPipeline _pMm, _pMmPs, _pMmDot, _pMmDd, _pMmIc, _pMmIc32; - // Concatenated .metal sources embedded as manifest resources, sorted so - // a_common.metal's helpers come first. - private MtlLibrary LoadLibrary() - { - var asm = typeof(MetalGraphModel).Assembly; - var names = asm.GetManifestResourceNames() - .Where(n => n.StartsWith("Sdcb.SimdPaddleOCR.Backends.Metal.Shaders.") - && n.EndsWith(".metal")) - .OrderBy(n => n).ToArray(); - if (names.Length == 0) - throw new FileNotFoundException("Metal: no embedded .metal shaders"); - var sb = new System.Text.StringBuilder(); - foreach (var n in names) - { - using Stream s = asm.GetManifestResourceStream(n)!; - using var r = new StreamReader(s); - sb.AppendLine(r.ReadToEnd()); - } - return _dev.NewLibrary(sb.ToString()); - } - /// PSO + live correctness smoke for simdgroup_matrix kernels: /// a PSO that links is not proof the driver executes it — the probe runs - /// a 64x64 tile GEMM against a canned input and compares one element. + /// a 64x64 tile GEMM against a canned input and compares the whole tile. + /// The kernel assumes 8 32-lane simdgroups (256 threads), so the PSO + /// geometry is gated too. /// private unsafe bool ProbeSimdMatrix() { try { - MtlPipeline p = _dev.NewPipeline(_lib, "mm_sg"); + MtlPipeline p = _dev.OcrPipe("mm_sg"); + if (p.ThreadExecutionWidth != 32 || p.MaxTotalThreadsPerThreadgroup < 256) + return false; const int M = 64, N = 64, K = 16; var a = new Half[M * K]; var w = new Half[N * K]; var rnd = new Random(1234); @@ -296,26 +277,21 @@ private unsafe bool ProbeSimdMatrix() e.End(); e.CommitAndWait(); } } - // verify a corner element: o[0,0] = sum_k a[0,k]*w[0,k] - float expect = 0; - for (int k = 0; k < K; k++) expect += (float)a[k] * (float)w[k]; - float got = (float)*(Half*)ob.Contents; - return MathF.Abs(got - expect) <= 0.05f * MathF.Max(1, MathF.Abs(expect)); + // whole-tile compare: o[m,n] = sum_k a[m,k]*w[n,k] (w is n-major) + Half* got = (Half*)ob.Contents; + for (int m = 0; m < M; m++) + for (int n = 0; n < N; n++) + { + float expect = 0; + for (int k = 0; k < K; k++) expect += (float)a[m * K + k] * (float)w[n * K + k]; + if (MathF.Abs((float)got[m * N + n] - expect) > 0.05f * MathF.Max(1, MathF.Abs(expect))) + return false; + } + return true; } catch { return false; } } - private readonly Dictionary _pipes = new(); - private MtlPipeline Pipe(string name) - { - if (!_pipes.TryGetValue(name, out MtlPipeline? p)) - { - p = _dev.NewPipeline(_lib, name); - _pipes[name] = p; - } - return p; - } - // constant upload helpers (Shared-mode buffers: contents is a plain // pointer on unified memory — the repack writes straight in) private unsafe MtlBuffer Up16(Half[] h) @@ -1436,21 +1412,15 @@ bool Absorbable(int cn, int srcPhys) // im2col-free direct conv: tap addressing inline in the // dot kernel; skips materializing the K-expanded matrix. // batched graphs always take it: the im2col fallback is batch-free - // sg32: implicit-GEMM coopmat conv instead (same tap - // addressing, tensor cores; any K, batched or not) - // Metal has no implicit-GEMM kxk kernel: the dot - // path covers every case (batched or K<=1024); the - // rest falls to im2col + mm below. - bool convkCm = false; - if (convkCm || (cout % 4 == 0 && scalarBias == 0 && cinIn % 4 == 0 + if (cout % 4 == 0 && scalarBias == 0 && cinIn % 4 == 0 && (nb > 1 || Kp <= (Environment.GetEnvironmentVariable( "SIMD_OCR_CONVD_KMAX") is string km ? int.Parse(km) : 1024)) - && Environment.GetEnvironmentVariable("SIMD_OCR_NODCONV") == null)) + && Environment.GetEnvironmentVariable("SIMD_OCR_NODCONV") == null) { // stem conv on the fp32 NCHW input: skip nchw2nhwc // when this is the sole consumer of the graph input - bool f32In = !convkCm && inT == inIdx && cinIn == 4 && cin <= 4 + bool f32In = inT == inIdx && cinIn == 4 && cin <= 4 && Environment.GetEnvironmentVariable( "SIMD_OCR_NOF32IN") == null; if (f32In && (consumers[inIdx]?.Count ?? 0) == 1) @@ -1492,9 +1462,11 @@ bool Absorbable(int cn, int srcPhys) Emit(_pMmIc32, $"icconv n{ni} {M}x{Kp}x{cout}", [(arena, SlotOf(node.Inputs[0]), 2), (wic, 0, 2), - (biasBuf, 0, 2), (arena, 0, 2), + (biasBuf, 0, 2), + (arena, resT != uint.MaxValue ? SlotOf(resT) : 0, 2), (arena, off[outPhys], 2), - (arena, 0, 2), (arena, 0, 2)], + (arena, dualT != uint.MaxValue ? SlotOf(dualT) : 0, 2), + (arena, 0, 2)], [M, (uint)cout, (uint)Kp, (uint)(Kp / 4), flags, 0u, 0u, (uint)outW, (uint)inW, (uint)inH, @@ -1525,7 +1497,10 @@ bool Absorbable(int cn, int srcPhys) Div256(((M + 3) / 4) * (long)(cout / 4)), (uint)nb); } - else if (_sgm && !_noIc && cinIn % 4 == 0 && nb == 1) + // scalarBias shares flags bit3 with the mm kernels' + // dual-write — the ic epilogue would write o2 garbage. + else if (_sgm && !_noIc && scalarBias == 0 + && cinIn % 4 == 0 && nb == 1) { // implicit-GEMM conv: mm_ic_sg gathers A through // tap-major im2col addressing inside the staging @@ -1535,9 +1510,11 @@ bool Absorbable(int cn, int srcPhys) Emit(_pMmIc32, $"icconv n{ni} {M}x{Kp}x{cout}", [(arena, SlotOf(node.Inputs[0]), 2), (wic, 0, 2), - (biasBuf, 0, 2), (arena, 0, 2), + (biasBuf, 0, 2), + (arena, resT != uint.MaxValue ? SlotOf(resT) : 0, 2), (arena, off[outPhys], 2), - (arena, 0, 2), (arena, 0, 2)], + (arena, dualT != uint.MaxValue ? SlotOf(dualT) : 0, 2), + (arena, 0, 2)], [M, (uint)cout, (uint)Kp, (uint)(Kp / 4), flags, 0u, 0u, (uint)outW, (uint)inW, (uint)inH, diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalRunner.cs b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalRunner.cs index 6b3e330..064c8b3 100644 --- a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalRunner.cs +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalRunner.cs @@ -40,18 +40,18 @@ internal sealed class MetalSchedule /// Streamed-run callback: units [first, first+count) of result /// (unit i at offsets[i]) are complete. Return false to stop receiving /// further batches. -internal delegate bool MetalUnitsReady(float[] result, int[] offsets, int first, int count); /// /// Per-session Metal runtime over a shared : /// owns the fp16 arena, fp32 input/output buffers and the partials scratch — /// all Shared-mode (unified memory: the CPU writes/reads the same mapping, /// no staging copies). Each run re-records one command buffer per wave from -/// the pure-managed schedule; dispatches go through a serial compute encoder -/// with a memory barrier after each rec (the schedule's correctness relies -/// on strict ordering just like the Vulkan per-rec barrier). +/// the pure-managed schedule; single-unit runs go through a serial compute +/// encoder (implicit ordering, no barriers needed), while interleaved +/// multi-unit runs use a concurrent encoder with one buffer barrier per +/// level — the Vulkan runner's per-rec pipeline barrier equivalent. /// -internal sealed unsafe class MetalDetGraph : IDisposable +internal sealed unsafe class MetalDetGraph : IOcrGraphRunner { private readonly MtlDevice _dev; private readonly MetalGraphModel _model; @@ -105,7 +105,7 @@ public ReadOnlySpan Run(int[] inputShape, ReadOnlySpan input, /// submission has completed by the time this returns or throws. /// public bool RunMany(IReadOnlyList shapes, ReadOnlySpan input, - int nodeLimit, int outTensor, MetalUnitsReady onReady) + int nodeLimit, int outTensor, CtcUnitsReady onReady) { ObjectDisposedException.ThrowIf(_disposed, this); var sched = new MetalSchedule[shapes.Count]; @@ -121,7 +121,7 @@ public bool RunMany(IReadOnlyList shapes, ReadOnlySpan input, private static long Align64(long floats) => (floats + 63) & ~63L; private float[] RunCore(MetalSchedule[] sched, ReadOnlySpan input, - out int[] outOffsets, MetalUnitsReady? sink = null) + out int[] outOffsets, CtcUnitsReady? sink = null) { long t0 = Stopwatch.GetTimestamp(); int n = sched.Length; @@ -190,9 +190,12 @@ private float[] RunCore(MetalSchedule[] sched, ReadOnlySpan input, IntPtr cmd; using (var pool = AutoReleasePool.Create()) { - // Serial encoder: dispatches run in order; one device-scope - // barrier per level/rec mirrors the Vulkan pipeline barrier. - MtlEncoder enc = _dev.BeginCommands(); + // Interleave needs a concurrent encoder — under a serial + // encoder dispatches already run in order, so a per-rec + // barrier would be a no-op and the level overlap the whole + // point of interleaving is unreachable. Concurrent + + // one buffer barrier per level restores both. + MtlEncoder enc = _dev.BeginCommands(concurrent: interleave); cmd = enc.Cmd; // +1 on the command buffer so it survives the pool drain // while we wait below. @@ -311,8 +314,7 @@ private void EnsurePart(int units) int alloc = Math.Max(units, _partUnits * 3 / 2); _part?.Dispose(); _part = null; - // newBuffer is zero-initialised on Metal: se_fused ticket counters - // start at 0 and self-reset after each run. + // se partial-sum scratch; contents are fully rewritten per run. _part = _dev.NewBuffer((nuint)alloc * MetalGraphModel.PartBytes); _partUnits = alloc; } @@ -365,11 +367,9 @@ private IReadOnlyList Selected(MetalSchedule s) private void RecordRecs(MtlEncoder enc, IReadOnlyList recs, ref MtlPipeline? bound, ulong inBase, ulong outBase) { + // serial encoder: dispatch order is the dependency order. foreach (MetalRec r in recs) - { RecordOne(enc, r, ref bound, inBase, outBase, 0, 0); - if (!s_noBar) enc.BarrierBuffers(); - } } private void RecordOne(MtlEncoder enc, MetalRec r, ref MtlPipeline? bound, diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalSession.cs b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalSession.cs index b0e3f29..4211484 100644 --- a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalSession.cs +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalSession.cs @@ -1,324 +1,18 @@ -using System.Buffers.Binary; -using System.Runtime.InteropServices; using Sdcb.SimdPaddleOCR.OnnxSharp; namespace Sdcb.SimdPaddleOCR.Backends.Metal; /// -/// GPU-backed over Metal: the whole ONNX graph runs -/// via . InputData is a managed staging span -/// uploaded on each run; the CTC tail matches 's -/// contract (stop before the vocab MatMul, hand operands to the shared CPU -/// ArgMax). One MtlDevice is shared process-wide; weights/pipelines/schedules -/// are shared per model, while each session owns its buffers, so sessions -/// only serialize at the command-buffer commits. +/// Metal-backed : the whole ONNX graph runs via +/// under the shared +/// contract (staging input, CPU CTC tail, permanent CPU fallback on any +/// plan/run failure). One MtlDevice is shared process-wide; each session owns +/// its buffers, so sessions only serialize at the command-buffer commits. /// -internal sealed class MetalSession : IOcrSession, IBatchedCtcSession +internal sealed class MetalSession : GpuSessionBase { - private readonly MtlDevice _dev; - private readonly MetalDetGraph _graph; - private readonly CompiledModel _compiled; - private readonly Model _model; - private readonly ResizeWorkspace _resizeWorkspace = new(); - private InferenceSession? _cpuFallback; - private float[] _cpuNhwc = []; // NCHW→NHWC staging for the CPU fallback - - private int[] _shape = []; - private float[] _input = []; - private int _curVolume; // numel of the current shape; _input may be bigger - private int _hwVolume; - private bool _disposed; - private bool _gpuDead; // plan/run failure → serve everything from _cpuFallback - - // CTC projection resolution, cached per Reshape - private bool _ctcResolved; - private int _ctcMatMulIndex = -1; - private int _ctcActTensor; - private int _ctcBatch, _ctcRows, _ctcInner, _ctcColumns; - private byte[]? _ctcWBytes, _ctcBiasBytes; - private float[]? _ctcPacked; - internal MetalSession(MtlDevice dev, CompiledModel compiled) + : base(new MetalDetGraph(dev, compiled), compiled, "metal") { - _dev = dev; - _compiled = compiled; - _model = compiled.Model; - _graph = new MetalDetGraph(dev, compiled); - } - - public TensorShape InputShape => new(_shape); - private int[]? _outputShape; - public TensorShape OutputShape => - new(_outputShape ??= _compiled.ResolveShapesFor(_shape)[checked((int)_model.GraphOutputs[0])]); - public Span InputData => _input.AsSpan(0, _curVolume); - public ResizeWorkspace ResizeWorkspace => _resizeWorkspace; - public bool InputIsNhwc => false; // convk_f32n reads NCHW fp32 directly - public int IntraOpThreads { get; set; } - public bool GpuAlive => !_gpuDead; - public int HighWaterInputVolume => _hwVolume; - public bool PlanForCtcProjection { get; set; } - public bool IsProfilingEnabled => InferenceSession.ProfilingEnabled; - - public void Reshape(ReadOnlySpan inputShape) - { - if (_disposed) throw new ObjectDisposedException(nameof(MetalSession)); - _shape = inputShape.ToArray(); - long vol = 1; - foreach (int d in _shape) vol *= Math.Max(d, 1); - _curVolume = checked((int)vol); - if (vol > _input.Length) - { - _input = new float[vol]; - _hwVolume = checked((int)Math.Min(vol, int.MaxValue)); - } - _outputShape = null; - _ctcResolved = false; - _ctcMatMulIndex = -1; - } - - public ReadOnlySpan RunInternal(ReadOnlySpan input) - { - if (_disposed) throw new ObjectDisposedException(nameof(MetalSession)); - if (_gpuDead) return Cpu().RunInternal(CpuInput(input)); - try - { - return _graph.Run(_shape, input); - } - catch (Exception ex) - { - Console.Error.WriteLine($"[metal] RunInternal fallback: {ex.GetType().Name} {ex.Message}"); - _gpuDead = true; - return Cpu().RunInternal(CpuInput(input)); - } - } - - public bool TryRunUntilCtcProjection(ReadOnlySpan input, out CtcProjectionOperands operands) - { - operands = default; - if (_disposed) throw new ObjectDisposedException(nameof(MetalSession)); - if (_gpuDead) return Cpu().TryRunUntilCtcProjection(CpuInput(input), out operands); - ResolveCtcProjection(); - if (_ctcMatMulIndex < 0) return false; - - ReadOnlySpan act; - try - { - act = _graph.Run(_shape, input, _ctcMatMulIndex, _ctcActTensor); - } - catch (Exception ex) - { - Console.Error.WriteLine($"[metal] CtcProj fallback: {ex.GetType().Name} {ex.Message}"); - _gpuDead = true; - return Cpu().TryRunUntilCtcProjection(CpuInput(input), out operands); - } - operands = new CtcProjectionOperands(act, - MemoryMarshal.Cast(_ctcWBytes), - _ctcBiasBytes is null ? [] : MemoryMarshal.Cast(_ctcBiasBytes), - _ctcPacked, _ctcBatch, _ctcRows, _ctcInner, _ctcColumns, _ctcMatMulIndex); - return true; - } - - // ---- IBatchedCtcSession: several REC shapes, one GPU submission ---- - private int[][] _many = []; - private int _manyVolume; - - public bool CanRunMany => !_gpuDead && !_disposed; - - public Span ReshapeMany(IReadOnlyList shapes) - { - if (_disposed) throw new ObjectDisposedException(nameof(MetalSession)); - _many = [.. shapes]; - long vol = 0; - foreach (int[] s in _many) - { - long v = 1; - foreach (int d in s) v *= Math.Max(d, 1); - vol += v; - } - _manyVolume = checked((int)vol); - if (vol > _input.Length) - { - _input = new float[vol]; - _hwVolume = _manyVolume; - } - // single-shape members (CTC resolution, fallback) key off the first unit - Reshape(_many[0]); - return _input.AsSpan(0, _manyVolume); - } - - public bool TryResolveManyHead(out int[] rows, out CtcHead head) - { - rows = []; head = null!; - if (!CanRunMany || _many.Length == 0) return false; - ResolveCtcProjection(); - if (_ctcMatMulIndex < 0) return false; - rows = new int[_many.Length]; - for (int i = 0; i < _many.Length; i++) - rows[i] = _compiled.ResolveShapesFor(_many[i])[_ctcActTensor][^2]; - head = new CtcHead - { - Weights = _ctcWBytes!, Bias = _ctcBiasBytes, Packed = _ctcPacked, - Inner = _ctcInner, Columns = _ctcColumns, MatMulIndex = _ctcMatMulIndex, - }; - return true; - } - - public bool TryRunManyUntilCtcProjection(CtcUnitsReady onReady) - { - if (!CanRunMany || _many.Length == 0) return false; - ResolveCtcProjection(); - if (_ctcMatMulIndex < 0) return false; - // consumer failures are the caller's, not a GPU fault: keep them out - // of the fallback catch below and rethrow once the device is drained - System.Runtime.ExceptionServices.ExceptionDispatchInfo? consumerError = null; - bool Forward(float[] acts, int[] offsets, int first, int count) - { - try { return onReady(acts, offsets, first, count); } - catch (Exception ex) - { - consumerError = System.Runtime.ExceptionServices.ExceptionDispatchInfo.Capture(ex); - return false; - } - } - bool ok; - try - { - ok = _graph.RunMany(_many, _input.AsSpan(0, _manyVolume), _ctcMatMulIndex, _ctcActTensor, Forward); - } - catch (Exception ex) - { - Console.Error.WriteLine($"[metal] RunMany fallback: {ex.GetType().Name} {ex.Message}"); - _gpuDead = true; - return false; - } - consumerError?.Throw(); - return ok; - } - - private InferenceSession Cpu() - { - _cpuFallback ??= new InferenceSession(_compiled) - { PlanForCtcProjection = PlanForCtcProjection }; - _cpuFallback.Reshape(_shape); - if (IntraOpThreads > 0) _cpuFallback.IntraOpThreads = IntraOpThreads; - return _cpuFallback; - } - - /// Graphs without the CTC tail fall back to a lazy CPU session. - public ReadOnlySpan RunInternalSkipFinalSoftmax(ReadOnlySpan input, out bool outputIsLogits) - { - if (_disposed) throw new ObjectDisposedException(nameof(MetalSession)); - return Cpu().RunInternalSkipFinalSoftmax(CpuInput(input), out outputIsLogits); - } - - /// - /// Callers always write logical NCHW (this session reports - /// = false). When the compiled model flags its - /// graph input as NHWC-physical, the CPU fallback must be fed the - /// transposed layout or it decodes garbage. - /// - private ReadOnlySpan CpuInput(ReadOnlySpan input) - { - if (!_compiled.InputIsNhwc) return input; - if (_cpuNhwc.Length < input.Length) _cpuNhwc = new float[input.Length]; - Span dest = _cpuNhwc.AsSpan(0, input.Length); - Kernels.Nhwc.NchwToNhwc(input, dest, _shape[0], _shape[1], - checked(_shape[2] * _shape[3]), Math.Max(IntraOpThreads, 1)); - return dest; - } - - public void NoteProfile(OperatorId operation, long started, int nodeIndex) - => InferenceSession.NoteProfileStatic(operation, started, nodeIndex); - - /// - /// Mirrors InferenceSession.TryResolveCtcProjection but resolves - /// shapes from instead of live - /// tensor bindings — the GPU session has no CPU workspace. - /// - private void ResolveCtcProjection() - { - if (_ctcResolved) return; - _ctcResolved = true; - _ctcMatMulIndex = -1; - - NodeRecord[] nodes = _model.Nodes; - if (nodes.Length < 2) return; - int[][] shapes = _compiled.ResolveShapesFor(_shape); - int outIdx = checked((int)_model.GraphOutputs[0]); - - NodeRecord softmax = nodes[^1]; - if (softmax.Operator != OperatorId.Softmax || - softmax.Inputs.Length != 1 || softmax.Outputs.Length != 1 || - softmax.Outputs[0] != (uint)outIdx) return; - int axis = I32(_model.GetParameters(softmax), 4); - int[] smInShape = shapes[softmax.Inputs[0]]; - if (axis < 0) axis += smInShape.Length; - if (axis != smInShape.Length - 1) return; - - int terminalIndex = nodes.Length - 2; - NodeRecord terminal = nodes[terminalIndex]; - int matMulIndex; - byte[]? biasBytes = null; - if (terminal.Operator == OperatorId.MatMul && - terminal.Outputs[0] == softmax.Inputs[0]) - { - matMulIndex = terminalIndex; - } - else if (terminal.Operator == OperatorId.Add && terminal.Inputs.Length == 2 && - terminalIndex > 0 && terminal.Outputs[0] == softmax.Inputs[0]) - { - matMulIndex = terminalIndex - 1; - NodeRecord mm = nodes[matMulIndex]; - if (mm.Operator != OperatorId.MatMul || mm.Outputs.Length != 1) return; - uint biasIndex; - if (terminal.Inputs[0] == mm.Outputs[0]) biasIndex = terminal.Inputs[1]; - else if (terminal.Inputs[1] == mm.Outputs[0]) biasIndex = terminal.Inputs[0]; - else return; - var biasMeta = _compiled.GetTensor(checked((int)biasIndex)); - int[] projShape = shapes[mm.Outputs[0]]; - if (!biasMeta.IsConstant || biasMeta.Shape.Length != 1 || - projShape.Length < 2 || biasMeta.Shape[0] != projShape[^1]) return; - biasBytes = biasMeta.Constant; - } - else return; - - NodeRecord matMul = nodes[matMulIndex]; - if (matMul.Inputs.Length != 2 || matMul.Outputs.Length != 1) return; - var aMeta = _compiled.GetTensor(checked((int)matMul.Inputs[0])); - var bMeta = _compiled.GetTensor(checked((int)matMul.Inputs[1])); - int[] aShape = shapes[matMul.Inputs[0]]; - int[] bShape = shapes[matMul.Inputs[1]]; - int[] projOut = shapes[matMul.Outputs[0]]; - if (!bMeta.IsConstant || aShape.Length < 2 || bShape.Length != 2 || - projOut.Length < 2) return; - int rows = aShape[^2], inner = aShape[^1], columns = bShape[1]; - long aCount = 1; - foreach (int d in aShape) aCount *= Math.Max(d, 1); - if (rows <= 0 || inner <= 0 || aCount % (rows * (long)inner) != 0) return; - int batch = checked((int)(aCount / (rows * (long)inner))); - if (bShape[0] != inner) return; - long projCount = 1; - foreach (int d in projOut) projCount *= Math.Max(d, 1); - if (projOut[^1] != columns || projCount != (long)batch * rows * columns) return; - _compiled.TryGetPackedMatMul(matMul.Inputs[1], out float[]? packed); - // Column count does not gate here (unlike the recognizer's ArgMax - // fusion): small-column heads like the cls [inner,2] tail resolve too. - if (columns <= 0) return; - - _ctcMatMulIndex = matMulIndex; - _ctcActTensor = checked((int)matMul.Inputs[0]); - _ctcBatch = batch; _ctcRows = rows; _ctcInner = inner; _ctcColumns = columns; - _ctcWBytes = bMeta.Constant; - _ctcBiasBytes = biasBytes; - _ctcPacked = packed; - } - - private static int I32(ReadOnlySpan p, int o) => BinaryPrimitives.ReadInt32LittleEndian(p[o..]); - - public void Dispose() - { - _disposed = true; - _graph.Dispose(); - _cpuFallback?.Dispose(); } } diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/gemm.metal b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/gemm.metal index d89beb0..fa4c75d 100644 --- a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/gemm.metal +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/gemm.metal @@ -119,8 +119,10 @@ kernel void mm_sg(device const half4* x [[buffer(0)]], // mm_sg_dd: same 64x64 tile / 8 simdgroup fan-out as mm_sg, but fragment loads // read device memory directly — no threadgroup staging round-trip (this GPU's // "threadgroup" memory is just a carve-out of the same DRAM, so staging only -// adds barriers). Emit-gated to M%64==0 && N%64==0 && K%16==0 so every tile is -// full and fragment loads never run past the operand end. +// adds barriers). Emit-gated to K%16==0 only: edge tiles are partial — fragment +// loads do read past the logical M/N end, which is safe because A overread stays +// inside the arena's tail slack, W rows are padded, and the epilogue masks +// mm using namespace metal; -// ---------------- se_fused ---------------- -// Single-pass SE: S threadgroups write per-partition partial sums; the last -// group to finish (atomic ticket) runs finalize inline: mean → fc1+relu → -// fc2 → hardsigmoid. +// ---------------- se_part + se_join ---------------- struct PcSe { uint hw, C, S, P, cvP, R; float alpha, beta; }; -// SE split into two dispatches: this paravirt GPU only supports -// memory_order_relaxed device atomics, so the Vulkan-style ticket -// (last-WG-does-stage2) cannot order partial writes vs reads. Emit -// se_part for all S workgroups, rely on the inter-dispatch buffer -// barrier, then se_join (one WG per batch) does stage2. +// SE split into two dispatches: MSL device atomics are relaxed-only, so the +// Vulkan-style ticket (last-WG-does-stage2) cannot order partial writes vs +// reads. Emit se_part for all S workgroups, rely on the inter-dispatch buffer +// barrier, then se_join (one WG per batch) does stage2: mean → fc1+relu → +// fc2 → hardsigmoid. kernel void se_part(device const half4* x [[buffer(0)]], device float* part [[buffer(1)]], constant PcSe& p [[buffer(2)]], diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Vulkan/GpuRunner.cs b/src/Sdcb.SimdPaddleOCR/Backends/Vulkan/GpuRunner.cs index 2491dba..56ecf6d 100644 --- a/src/Sdcb.SimdPaddleOCR/Backends/Vulkan/GpuRunner.cs +++ b/src/Sdcb.SimdPaddleOCR/Backends/Vulkan/GpuRunner.cs @@ -39,7 +39,6 @@ internal sealed class GpuSchedule /// Streamed-run callback: units [first, first+count) of result /// (unit i at offsets[i]) are complete. Return false to stop receiving /// further batches. -internal delegate bool GpuUnitsReady(float[] result, int[] offsets, int first, int count); /// /// Per-session GPU runtime over a shared : owns @@ -50,7 +49,7 @@ internal sealed class GpuSchedule /// contend beyond the queue-submit lock — CPU work of concurrent sessions /// (upload, record, readback) overlaps freely. /// -internal sealed unsafe class GpuDetGraph : IDisposable +internal sealed unsafe class GpuDetGraph : IOcrGraphRunner { private readonly VkDevice _dev; private readonly GpuGraphModel _model; @@ -128,7 +127,7 @@ public ReadOnlySpan Run(int[] inputShape, ReadOnlySpan input, /// submission has completed by the time this returns or throws. /// public bool RunMany(IReadOnlyList shapes, ReadOnlySpan input, - int nodeLimit, int outTensor, GpuUnitsReady onReady) + int nodeLimit, int outTensor, CtcUnitsReady onReady) { ObjectDisposedException.ThrowIf(_disposed, this); var sched = new GpuSchedule[shapes.Count]; @@ -145,7 +144,7 @@ public bool RunMany(IReadOnlyList shapes, ReadOnlySpan input, private static long Align64(long floats) => (floats + 63) & ~63L; private float[] RunCore(GpuSchedule[] sched, ReadOnlySpan input, out int[] outOffsets, - GpuUnitsReady? sink = null) + CtcUnitsReady? sink = null) { long t0 = Stopwatch.GetTimestamp(); int n = sched.Length; diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Vulkan/GpuSession.cs b/src/Sdcb.SimdPaddleOCR/Backends/Vulkan/GpuSession.cs index 39161b0..66869e1 100644 --- a/src/Sdcb.SimdPaddleOCR/Backends/Vulkan/GpuSession.cs +++ b/src/Sdcb.SimdPaddleOCR/Backends/Vulkan/GpuSession.cs @@ -1,324 +1,19 @@ -using System.Buffers.Binary; -using System.Runtime.InteropServices; using Sdcb.SimdPaddleOCR.OnnxSharp; namespace Sdcb.SimdPaddleOCR.Backends.Vulkan; /// -/// GPU-backed : the whole ONNX graph runs on Vulkan -/// via . InputData is a managed staging span uploaded -/// on each run; the CTC tail matches 's contract -/// (stop before the vocab MatMul, hand operands to the shared CPU ArgMax). -/// One VkDevice is shared process-wide; weights/pipelines/schedules are shared -/// per model, while each session owns its buffers and command buffer, so -/// sessions only serialize on the queue submit itself. +/// Vulkan-backed : the whole ONNX graph runs via +/// under the shared +/// contract (staging input, CPU CTC tail, permanent CPU fallback on any +/// plan/run failure). One VkDevice is shared process-wide; each session owns +/// its buffers and command buffer, so sessions only serialize on the queue +/// submit itself. /// -internal sealed class GpuSession : IOcrSession, IBatchedCtcSession +internal sealed class GpuSession : GpuSessionBase { - private readonly VkDevice _dev; - private readonly GpuDetGraph _graph; - private readonly CompiledModel _compiled; - private readonly Model _model; - private readonly ResizeWorkspace _resizeWorkspace = new(); - private InferenceSession? _cpuFallback; - private float[] _cpuNhwc = []; // NCHW→NHWC staging for the CPU fallback - - private int[] _shape = []; - private float[] _input = []; - private int _curVolume; // numel of the current shape; _input may be bigger - private int _hwVolume; - private bool _disposed; - private bool _gpuDead; // plan/run failure → serve everything from _cpuFallback - - // CTC projection resolution, cached per Reshape - private bool _ctcResolved; - private int _ctcMatMulIndex = -1; - private int _ctcActTensor; - private int _ctcBatch, _ctcRows, _ctcInner, _ctcColumns; - private byte[]? _ctcWBytes, _ctcBiasBytes; - private float[]? _ctcPacked; - internal GpuSession(VkDevice dev, CompiledModel compiled) + : base(new GpuDetGraph(dev, compiled), compiled, "gpu") { - _dev = dev; - _compiled = compiled; - _model = compiled.Model; - _graph = new GpuDetGraph(dev, compiled); - } - - public TensorShape InputShape => new(_shape); - private int[]? _outputShape; - public TensorShape OutputShape => - new(_outputShape ??= _compiled.ResolveShapesFor(_shape)[checked((int)_model.GraphOutputs[0])]); - public Span InputData => _input.AsSpan(0, _curVolume); - public ResizeWorkspace ResizeWorkspace => _resizeWorkspace; - public bool InputIsNhwc => false; // convk_f32n reads NCHW fp32 directly - public int IntraOpThreads { get; set; } - public bool GpuAlive => !_gpuDead; - public int HighWaterInputVolume => _hwVolume; - public bool PlanForCtcProjection { get; set; } - public bool IsProfilingEnabled => InferenceSession.ProfilingEnabled; - - public void Reshape(ReadOnlySpan inputShape) - { - if (_disposed) throw new ObjectDisposedException(nameof(GpuSession)); - _shape = inputShape.ToArray(); - long vol = 1; - foreach (int d in _shape) vol *= Math.Max(d, 1); - _curVolume = checked((int)vol); - if (vol > _input.Length) - { - _input = new float[vol]; - _hwVolume = checked((int)Math.Min(vol, int.MaxValue)); - } - _outputShape = null; - _ctcResolved = false; - _ctcMatMulIndex = -1; - } - - public ReadOnlySpan RunInternal(ReadOnlySpan input) - { - if (_disposed) throw new ObjectDisposedException(nameof(GpuSession)); - if (_gpuDead) return Cpu().RunInternal(CpuInput(input)); - try - { - return _graph.Run(_shape, input); - } - catch (Exception ex) - { - Console.Error.WriteLine($"[gpu] RunInternal fallback: {ex.GetType().Name} {ex.Message}"); - _gpuDead = true; - return Cpu().RunInternal(CpuInput(input)); - } - } - - public bool TryRunUntilCtcProjection(ReadOnlySpan input, out CtcProjectionOperands operands) - { - operands = default; - if (_disposed) throw new ObjectDisposedException(nameof(GpuSession)); - if (_gpuDead) return Cpu().TryRunUntilCtcProjection(CpuInput(input), out operands); - ResolveCtcProjection(); - if (_ctcMatMulIndex < 0) return false; - - ReadOnlySpan act; - try - { - act = _graph.Run(_shape, input, _ctcMatMulIndex, _ctcActTensor); - } - catch (Exception ex) - { - Console.Error.WriteLine($"[gpu] CtcProj fallback: {ex.GetType().Name} {ex.Message}"); - _gpuDead = true; - return Cpu().TryRunUntilCtcProjection(CpuInput(input), out operands); - } - operands = new CtcProjectionOperands(act, - MemoryMarshal.Cast(_ctcWBytes), - _ctcBiasBytes is null ? [] : MemoryMarshal.Cast(_ctcBiasBytes), - _ctcPacked, _ctcBatch, _ctcRows, _ctcInner, _ctcColumns, _ctcMatMulIndex); - return true; - } - - // ---- IBatchedCtcSession: several REC shapes, one GPU submission ---- - private int[][] _many = []; - private int _manyVolume; - - public bool CanRunMany => !_gpuDead && !_disposed; - - public Span ReshapeMany(IReadOnlyList shapes) - { - if (_disposed) throw new ObjectDisposedException(nameof(GpuSession)); - _many = [.. shapes]; - long vol = 0; - foreach (int[] s in _many) - { - long v = 1; - foreach (int d in s) v *= Math.Max(d, 1); - vol += v; - } - _manyVolume = checked((int)vol); - if (vol > _input.Length) - { - _input = new float[vol]; - _hwVolume = _manyVolume; - } - // single-shape members (CTC resolution, fallback) key off the first unit - Reshape(_many[0]); - return _input.AsSpan(0, _manyVolume); - } - - public bool TryResolveManyHead(out int[] rows, out CtcHead head) - { - rows = []; head = null!; - if (!CanRunMany || _many.Length == 0) return false; - ResolveCtcProjection(); - if (_ctcMatMulIndex < 0) return false; - rows = new int[_many.Length]; - for (int i = 0; i < _many.Length; i++) - rows[i] = _compiled.ResolveShapesFor(_many[i])[_ctcActTensor][^2]; - head = new CtcHead - { - Weights = _ctcWBytes!, Bias = _ctcBiasBytes, Packed = _ctcPacked, - Inner = _ctcInner, Columns = _ctcColumns, MatMulIndex = _ctcMatMulIndex, - }; - return true; - } - - public bool TryRunManyUntilCtcProjection(CtcUnitsReady onReady) - { - if (!CanRunMany || _many.Length == 0) return false; - ResolveCtcProjection(); - if (_ctcMatMulIndex < 0) return false; - // consumer failures are the caller's, not a GPU fault: keep them out - // of the fallback catch below and rethrow once the device is drained - System.Runtime.ExceptionServices.ExceptionDispatchInfo? consumerError = null; - bool Forward(float[] acts, int[] offsets, int first, int count) - { - try { return onReady(acts, offsets, first, count); } - catch (Exception ex) - { - consumerError = System.Runtime.ExceptionServices.ExceptionDispatchInfo.Capture(ex); - return false; - } - } - bool ok; - try - { - ok = _graph.RunMany(_many, _input.AsSpan(0, _manyVolume), _ctcMatMulIndex, _ctcActTensor, Forward); - } - catch (Exception ex) - { - Console.Error.WriteLine($"[gpu] RunMany fallback: {ex.GetType().Name} {ex.Message}"); - _gpuDead = true; - return false; - } - consumerError?.Throw(); - return ok; - } - - private InferenceSession Cpu() - { - _cpuFallback ??= new InferenceSession(_compiled) - { PlanForCtcProjection = PlanForCtcProjection }; - _cpuFallback.Reshape(_shape); - if (IntraOpThreads > 0) _cpuFallback.IntraOpThreads = IntraOpThreads; - return _cpuFallback; - } - - /// Graphs without the CTC tail fall back to a lazy CPU session. - public ReadOnlySpan RunInternalSkipFinalSoftmax(ReadOnlySpan input, out bool outputIsLogits) - { - if (_disposed) throw new ObjectDisposedException(nameof(GpuSession)); - return Cpu().RunInternalSkipFinalSoftmax(CpuInput(input), out outputIsLogits); - } - - /// - /// Callers always write logical NCHW (this session reports - /// = false). When the compiled model flags its - /// graph input as NHWC-physical, the CPU fallback must be fed the - /// transposed layout or it decodes garbage. - /// - private ReadOnlySpan CpuInput(ReadOnlySpan input) - { - if (!_compiled.InputIsNhwc) return input; - if (_cpuNhwc.Length < input.Length) _cpuNhwc = new float[input.Length]; - Span dest = _cpuNhwc.AsSpan(0, input.Length); - Kernels.Nhwc.NchwToNhwc(input, dest, _shape[0], _shape[1], - checked(_shape[2] * _shape[3]), Math.Max(IntraOpThreads, 1)); - return dest; - } - - public void NoteProfile(OperatorId operation, long started, int nodeIndex) - => InferenceSession.NoteProfileStatic(operation, started, nodeIndex); - - /// - /// Mirrors InferenceSession.TryResolveCtcProjection but resolves - /// shapes from instead of live - /// tensor bindings — the GPU session has no CPU workspace. - /// - private void ResolveCtcProjection() - { - if (_ctcResolved) return; - _ctcResolved = true; - _ctcMatMulIndex = -1; - - NodeRecord[] nodes = _model.Nodes; - if (nodes.Length < 2) return; - int[][] shapes = _compiled.ResolveShapesFor(_shape); - int outIdx = checked((int)_model.GraphOutputs[0]); - - NodeRecord softmax = nodes[^1]; - if (softmax.Operator != OperatorId.Softmax || - softmax.Inputs.Length != 1 || softmax.Outputs.Length != 1 || - softmax.Outputs[0] != (uint)outIdx) return; - int axis = I32(_model.GetParameters(softmax), 4); - int[] smInShape = shapes[softmax.Inputs[0]]; - if (axis < 0) axis += smInShape.Length; - if (axis != smInShape.Length - 1) return; - - int terminalIndex = nodes.Length - 2; - NodeRecord terminal = nodes[terminalIndex]; - int matMulIndex; - byte[]? biasBytes = null; - if (terminal.Operator == OperatorId.MatMul && - terminal.Outputs[0] == softmax.Inputs[0]) - { - matMulIndex = terminalIndex; - } - else if (terminal.Operator == OperatorId.Add && terminal.Inputs.Length == 2 && - terminalIndex > 0 && terminal.Outputs[0] == softmax.Inputs[0]) - { - matMulIndex = terminalIndex - 1; - NodeRecord mm = nodes[matMulIndex]; - if (mm.Operator != OperatorId.MatMul || mm.Outputs.Length != 1) return; - uint biasIndex; - if (terminal.Inputs[0] == mm.Outputs[0]) biasIndex = terminal.Inputs[1]; - else if (terminal.Inputs[1] == mm.Outputs[0]) biasIndex = terminal.Inputs[0]; - else return; - var biasMeta = _compiled.GetTensor(checked((int)biasIndex)); - int[] projShape = shapes[mm.Outputs[0]]; - if (!biasMeta.IsConstant || biasMeta.Shape.Length != 1 || - projShape.Length < 2 || biasMeta.Shape[0] != projShape[^1]) return; - biasBytes = biasMeta.Constant; - } - else return; - - NodeRecord matMul = nodes[matMulIndex]; - if (matMul.Inputs.Length != 2 || matMul.Outputs.Length != 1) return; - var aMeta = _compiled.GetTensor(checked((int)matMul.Inputs[0])); - var bMeta = _compiled.GetTensor(checked((int)matMul.Inputs[1])); - int[] aShape = shapes[matMul.Inputs[0]]; - int[] bShape = shapes[matMul.Inputs[1]]; - int[] projOut = shapes[matMul.Outputs[0]]; - if (!bMeta.IsConstant || aShape.Length < 2 || bShape.Length != 2 || - projOut.Length < 2) return; - int rows = aShape[^2], inner = aShape[^1], columns = bShape[1]; - long aCount = 1; - foreach (int d in aShape) aCount *= Math.Max(d, 1); - if (rows <= 0 || inner <= 0 || aCount % (rows * (long)inner) != 0) return; - int batch = checked((int)(aCount / (rows * (long)inner))); - if (bShape[0] != inner) return; - long projCount = 1; - foreach (int d in projOut) projCount *= Math.Max(d, 1); - if (projOut[^1] != columns || projCount != (long)batch * rows * columns) return; - _compiled.TryGetPackedMatMul(matMul.Inputs[1], out float[]? packed); - // Column count does not gate here (unlike the recognizer's ArgMax - // fusion): small-column heads like the cls [inner,2] tail resolve too. - if (columns <= 0) return; - - _ctcMatMulIndex = matMulIndex; - _ctcActTensor = checked((int)matMul.Inputs[0]); - _ctcBatch = batch; _ctcRows = rows; _ctcInner = inner; _ctcColumns = columns; - _ctcWBytes = bMeta.Constant; - _ctcBiasBytes = biasBytes; - _ctcPacked = packed; - } - - private static int I32(ReadOnlySpan p, int o) => BinaryPrimitives.ReadInt32LittleEndian(p[o..]); - - public void Dispose() - { - _disposed = true; - _graph.Dispose(); - _cpuFallback?.Dispose(); } } diff --git a/src/Sdcb.SimdPaddleOCR/OnnxSharp/OcrSessionFactory.cs b/src/Sdcb.SimdPaddleOCR/OnnxSharp/OcrSessionFactory.cs index c82b7ba..200fc4d 100644 --- a/src/Sdcb.SimdPaddleOCR/OnnxSharp/OcrSessionFactory.cs +++ b/src/Sdcb.SimdPaddleOCR/OnnxSharp/OcrSessionFactory.cs @@ -25,7 +25,7 @@ internal static bool IsGpuBackend(OcrBackend backend) { #if NET10_0_OR_GREATER if (Backends.Metal.MetalBackend.IsMetalSelected(backend)) - return Backends.Metal.MetalBackend.TryGetDevice() is not null; + return Backends.Metal.MetalBackend.UsesGpu(backend); return Backends.Vulkan.GpuBackend.UsesGpu(backend); #else _ = backend; From 043f98ac5c63ddcf885f22f58dcc1fbd3e05854c Mon Sep 17 00:00:00 2001 From: ZHOU Jie Date: Mon, 28 Sep 2026 05:39:58 -0700 Subject: [PATCH 4/8] [metal] device breadcrumb on session create; clarify mm_ic pad-k NaN edge --- .../Backends/Metal/MetalBackend.cs | 10 +++++++++- .../Backends/Metal/Shaders/gemm.metal | 16 ++++++++++------ 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalBackend.cs b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalBackend.cs index 16c69db..bb1287b 100644 --- a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalBackend.cs +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalBackend.cs @@ -13,6 +13,7 @@ internal static class MetalBackend private static readonly object s_probeLock = new(); private static MtlDevice? s_device; private static bool s_probed; + private static int s_announced; /// The shared device, or null when Metal is unavailable (not /// macOS, no GPU, or the ObjC bridge can't be reached). @@ -67,7 +68,14 @@ internal static IOcrSession CreateSession(CompiledModel compiled, OcrBackend bac { if (UsesGpu(backend) && TryGetDevice() is { } dev) { - try { return new MetalSession(dev, compiled); } + try + { + // one stderr breadcrumb per process so CI/users can tell the + // Metal path actually engaged (a failed probe is silent). + if (System.Threading.Interlocked.Exchange(ref s_announced, 1) == 0) + Console.Error.WriteLine($"[metal] device: {dev.Name}"); + return new MetalSession(dev, compiled); + } catch { /* fall through to CPU */ } } return compiled.CreateRequest(); diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/gemm.metal b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/gemm.metal index fa4c75d..ba0c818 100644 --- a/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/gemm.metal +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/Shaders/gemm.metal @@ -389,9 +389,11 @@ kernel void mm_dot(device const half4* x [[buffer(0)]], // A stage gathers input pixels through tap-major im2col addressing instead of // reading a materialized [M,K] matrix. One staged half4 never crosses a tap // (cinE%4==0 gate on the emit side). p.K is the padded reduction extent -// (multiple of 16): pad k-slices hit W rows that ConstF16 zero-filled, so they -// contribute 0 even though the im2col gather may read a real pixel there. -// Batch-free like the im2col path it replaces. +// (multiple of 16): pad k-slices gather REAL pixels (tap>=kH*kW still lands +// in-bounds) but multiply against W rows that ConstF16 zero-filled, so they +// contribute 0 — correct for finite activations; an inf activation would make +// NaN here (0*inf), which no real graph produces. Batch-free like the im2col +// path it replaces. struct PcGemmIc { uint M, N, K, kv4, flags, mImg, scMod; // K = Kp (padded, W stride) uint outW, inW, inH, sH, sW, pT, pL, kW, cinE; @@ -499,9 +501,11 @@ kernel void mm_ic_sg(device const half4* x [[buffer(0)]], // A stage gathers input pixels through tap-major im2col addressing instead of // reading a materialized [M,K] matrix. One staged half4 never crosses a tap // (cinE%4==0 gate on the emit side). p.K is the padded reduction extent -// (multiple of 16): pad k-slices hit W rows that ConstF16 zero-filled, so they -// contribute 0 even though the im2col gather may read a real pixel there. -// Batch-free like the im2col path it replaces. +// (multiple of 16): pad k-slices gather REAL pixels (tap>=kH*kW still lands +// in-bounds) but multiply against W rows that ConstF16 zero-filled, so they +// contribute 0 — correct for finite activations; an inf activation would make +// NaN here (0*inf), which no real graph produces. Batch-free like the im2col +// path it replaces. struct PcGemmIc32 { uint M, N, K, kv4, flags, mImg, scMod; // K = Kp (padded, W stride) uint outW, inW, inH, sH, sW, pT, pL, kW, cinE; From 6fde8fb4653a07a44a7a033a170da93b9cc237b9 Mon Sep 17 00:00:00 2001 From: sdcb Date: Mon, 28 Sep 2026 21:12:33 +0800 Subject: [PATCH 5/8] Add macOS Metal CI smoke on osx-arm64 --- .github/workflows/test.yml | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 19155f2..cda70dd 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -81,11 +81,37 @@ jobs: if (-not $IsWindows) { chmod +x $exe } & $exe --smoke --count 20 --engine sharp --workers 4 --model tiny --input dataset --case-id '${{ matrix.name }}' --out '${{ matrix.name }}-smoke.json' if ($LASTEXITCODE -ne 0) { throw "smoke failed: $LASTEXITCODE" } + - name: Metal smoke (osx-arm64) + if: matrix.name == 'osx-arm64' + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + $exe = 'artifacts/net10/Sdcb.SimdPaddleOCR.Tests' + $env:DOTNET_ROOT = Join-Path $HOME '.dotnet' + & $exe --smoke --count 20 --engine metal --workers 4 --model tiny --input dataset --case-id 'osx-arm64-metal' --out 'osx-arm64-metal-smoke.json' 2> metal-stderr.txt + if ($LASTEXITCODE -ne 0) { throw "metal smoke failed: $LASTEXITCODE" } + Get-Content metal-stderr.txt + if (-not (Select-String -Path metal-stderr.txt -Pattern '\[metal\] device:' -Quiet)) { + throw 'Metal device never engaged — the smoke silently ran on CPU' + } + if (Select-String -Path metal-stderr.txt -Pattern '\[metal\].*fallback' -Quiet) { + throw 'metal session fell back to CPU' + } + $cpu = Get-Content 'osx-arm64-smoke.json' | ConvertFrom-Json + $mtl = Get-Content 'osx-arm64-metal-smoke.json' | ConvertFrom-Json + for ($i = 0; $i -lt $cpu.rows.Count; $i++) { + if ($cpu.rows[$i].detected -ne $mtl.rows[$i].detected) { + throw "metal detected-line count differs on $($cpu.rows[$i].file): $($cpu.rows[$i].detected) vs $($mtl.rows[$i].detected)" + } + } + Write-Output "metal smoke OK: $($mtl.rows.Count) images, detected-line counts identical to CPU" - uses: actions/upload-artifact@v7 if: always() with: name: smoke-platform-${{ matrix.name }} - path: ${{ matrix.name }}-smoke.json + path: | + ${{ matrix.name }}-smoke.json + ${{ matrix.name }}-metal-smoke.json if-no-files-found: warn smoke-models: From 1bf03502faafdfa72f2656444878f1cfe18ec8fb Mon Sep 17 00:00:00 2001 From: ZHOU Jie Date: Mon, 28 Sep 2026 19:09:10 -0700 Subject: [PATCH 6/8] docs: same-run precision metrics + perf re-measure at review HEAD; fix OcrBackend.Auto comment --- docs/metal-m4.md | 46 ++++++++++++++++------------ src/Sdcb.SimdPaddleOCR/OcrBackend.cs | 2 +- 2 files changed, 27 insertions(+), 21 deletions(-) diff --git a/docs/metal-m4.md b/docs/metal-m4.md index 455da70..dd1591e 100644 --- a/docs/metal-m4.md +++ b/docs/metal-m4.md @@ -1,16 +1,16 @@ # SimdPaddleOCR — Apple Silicon Metal 后端实测（Devin VM / M4 paravirt） -实测机：macOS VM（Apple M4 paravirt GPU，8 vCPU），SDK .NET 10，`test/Sdcb.SimdPaddleOCR.Tests` bench（`--workers 4 --benchmark-kind simd --engine sharp|metal`，warmup=1，n=99，dataset/ 100 张固定图、张张不同尺寸）。commit：`metal-backend` 分支（PR #21），性能数据采集于 `680bf38`（其后的 review 修订不动热路径）。 +实测机：macOS VM（Apple M4 paravirt GPU，8 vCPU），SDK .NET 10，`test/Sdcb.SimdPaddleOCR.Tests` bench（`--workers 4 --benchmark-kind simd --engine sharp|metal`，warmup=1，n=99，dataset/ 100 张固定图、张张不同尺寸）。commit：`metal-backend` 分支（PR #21），下表为 review 修订（concurrent encoder、device 级 MSL lib/PSO cache）之后在同一轮重测的数据。 -**结论：三档 Metal 全部净胜 CPU（tiny 1.9× / small 4.7× / medium 5.2×），检测行数与 CPU 完全一致，文本差异仅 fp16 级噪声（3 图各 1 字符）。受 paravirt 虚拟化所限绝对数偏保守——实测 MMA≈fp32≈3TFLOPS、copy ~84GB/s、dispatch ~45µs，均低于真机 M4，真机数字预期更好。** +**结论：三档 Metal 全部净胜 CPU（tiny 1.75× / small 4.00× / medium 6.47×），检测行数与 CPU 完全一致，文本差异为 fp16 级噪声（vs CPU 逐图差：tiny 11 图、small 4 图、medium 3 图，多数为单字符空格增减；medium 的 exact_lines/CER 反而略优于 CPU）。受 paravirt 虚拟化所限绝对数偏保守——实测 MMA≈fp32≈3TFLOPS、copy ~84GB/s、dispatch ~45µs，均低于真机 M4，真机数字预期更好。** ## 端到端（4 workers，median ms/图，越低越好） | 模型 | CPU(sharp) | Metal | 比值 | CPU img/s | GPU img/s | |---|---:|---:|---:|---:|---:| -| tiny | 73.9 | **38.9** | 1.90× | 13.5 | 25.7 | -| small | 269.7 | **57.1** | 4.72× | 3.7 | 17.5 | -| medium | 935.5 | **178.4** | 5.24× | 1.07 | 5.6 | +| tiny | 68.8 | **39.3** | 1.75× | 14.5 | 25.4 | +| small | 218.4 | **54.6** | 4.00× | 4.6 | 18.3 | +| medium | 693.0 | **107.0** | 6.47× | 1.44 | 9.35 | 注：与 Vulkan/B580 表口径相同（dataset 变 shape 是最不利 GPU 的场景）。Metal 走 `OcrBackend.Metal`（Auto 在 macOS+Metal 可用时自动选）。 @@ -18,25 +18,31 @@ | 模型 | 阶段 | CPU | Metal | 倍率 | |---|---|---:|---:|---:| -| tiny | det_graph | 35.5 | 12.7 | 2.8× | -| tiny | cls_graph | 11.3 | 3.25 | 3.5× | -| tiny | rec_graph | 99.7 | 15.9 | 6.3× | -| small | det_graph | 104.7 | 14.9 | 7.0× | -| small | cls_graph | 12.2 | 2.45 | 5.0× | -| small | rec_graph | 585.5 | 33.3 | 17.6× | -| medium | det_graph | 315.1 | 59.6 | 5.3× | -| medium | cls_graph | 11.9 | 2.61 | 4.6× | -| medium | rec_graph | 2322.2 | 112.6 | 20.6× | - -精度：metal/sharp 检出 lines 完全一致（1016/1009/1013）；逐图 hash 3 张差（fp16 阈值临界像素，与 Vulkan 同型已知限制）；文本 3 图各差 1 字符（CER 口径：对 dataset 真值标注的字符错误率——Metal 0.67% vs CPU 0.14%，在 fp16 噪声范围内）。并发正确性：`--conc`（4 线程 × 72 runs/模型）0 mismatch；det medium 228 节点逐节点对拍（本地 probe 工具，未提交）× 8 跑 bitwise 一致。 +| tiny | det_graph | 44.0 | 13.0 | 3.4× | +| tiny | cls_graph | 12.2 | 4.3 | 2.8× | +| tiny | rec_graph | 101.7 | 18.6 | 5.5× | +| small | det_graph | 124.6 | 13.6 | 9.2× | +| small | cls_graph | 12.0 | 3.2 | 3.8× | +| small | rec_graph | 495.4 | 31.4 | 15.8× | +| medium | det_graph | 317.4 | 33.1 | 9.6× | +| medium | cls_graph | 12.4 | 3.2 | 3.9× | +| medium | rec_graph | 1877.8 | 68.9 | 27.3× | + +精度（与性能表同一轮 JSON 跑出的三个口径）： + +- **检出**：metal/sharp 逐图 `detected` 完全一致（1016/1009/1013，共 300 runs 0 差）。 +- **对真值**（dataset 标注）：exact_lines — tiny 744/1024 = 744/1024、small 811/1024 = 811/1024、medium **897 vs 894（Metal 更高）**；CER — tiny 1.640% vs 1.614%、small 1.392% vs 1.396%、medium 0.665% vs 0.677%。 +- **对 CPU 逐图文本差**（同一轮的两个 JSON 直比）：tiny 11 图 / 49 字符（10 张为单字符空格增删，最差 img-078 差 38 字符）、small 4 图 / 4 字符、medium 3 图 / 3 字符；逐图 hash 差与文本差同集；CLS 旋转角 1 图翻转（tiny，含在上述集合内）。属 fp16 阈值临界噪声，与 Vulkan 同型已知限制。 + +并发正确性：`--conc`（4 线程 × 72 runs/模型）0 mismatch；det medium 228 节点逐节点对拍（本地 probe 工具，未提交）× 8 跑 bitwise 一致。 ## 内存与分配 | 模型 | CPU peak WS | Metal peak WS | |---|---:|---:| -| tiny | 691 MB | 574 MB | -| small | 857 MB | 647 MB | -| medium | 1691 MB | 948 MB | +| tiny | 688 MB | 578 MB | +| small | 955 MB | 649 MB | +| medium | 1709 MB | 942 MB | GPU 侧 arena 是共享 grow-only device buffer + 托管 schedule（MetalSchedule 纯数据 LRU=512，逐 shape 录制代价为零），与 Vulkan 的 plan/arena 形态同构。托管堆分配 ~1.4–2.5MB/图。 @@ -70,4 +76,4 @@ GPU 侧 arena 是共享 grow-only device buffer + 托管 schedule（MetalSchedul - DET bs>1 resize 不支持（同 Vulkan）。 - MSL 通用限制：device `memory_order` 只有 relaxed（非 paravirt 独有）→ 跨 WG 定序走两阶段 split + `memoryBarrierWithScope:Buffers`；paravirt 另限 `char16/uchar16` 不可用。 - 序列化逐 dispatch profile（`SIMD_OCR_GPU_PROF`）有 ~0.4–0.5ms/dispatch 地板，绝对值偏大；真实比例看 `SIMD_OCR_GPU_TIME` 的墙钟分解。 -- 无 CI Metal 依赖：Metal 路径只在 macOS 上被选中，其他 OS 的测试和 bench 不会走到它。 +- CI：osx-arm64 smoke（macos-26 runner，paravirt GPU）以 `--engine metal` 跑 tiny 20 张，断言 stderr 含 `[metal] device:` breadcrumb、无 `fallback`、逐图 `detected` 与 CPU 一致、逐图 `texts` 与 CPU 的编辑距离 ≤3/图且合计 ≤8（fp16 噪声允许量）；其他 OS 的测试不会走到 Metal 路径。 diff --git a/src/Sdcb.SimdPaddleOCR/OcrBackend.cs b/src/Sdcb.SimdPaddleOCR/OcrBackend.cs index 9c2dcd1..d805b70 100644 --- a/src/Sdcb.SimdPaddleOCR/OcrBackend.cs +++ b/src/Sdcb.SimdPaddleOCR/OcrBackend.cs @@ -3,7 +3,7 @@ namespace Sdcb.SimdPaddleOCR; /// Compute backend for the OCR graph stages. public enum OcrBackend { - /// Prefer Vulkan when a usable device exists, else CPU. + /// Prefer Metal on macOS arm64 and Vulkan elsewhere when a usable device exists, else CPU. Auto, /// Always run the CPU interpreter. Cpu, From 860166bbeef125fd6447a3fb5e8504694b982d9f Mon Sep 17 00:00:00 2001 From: ZHOU Jie Date: Mon, 28 Sep 2026 19:43:19 -0700 Subject: [PATCH 7/8] docs: describe actual CI smoke assertions (no texts compare) --- docs/metal-m4.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/metal-m4.md b/docs/metal-m4.md index dd1591e..0cbe3da 100644 --- a/docs/metal-m4.md +++ b/docs/metal-m4.md @@ -76,4 +76,4 @@ GPU 侧 arena 是共享 grow-only device buffer + 托管 schedule（MetalSchedul - DET bs>1 resize 不支持（同 Vulkan）。 - MSL 通用限制：device `memory_order` 只有 relaxed（非 paravirt 独有）→ 跨 WG 定序走两阶段 split + `memoryBarrierWithScope:Buffers`；paravirt 另限 `char16/uchar16` 不可用。 - 序列化逐 dispatch profile（`SIMD_OCR_GPU_PROF`）有 ~0.4–0.5ms/dispatch 地板，绝对值偏大；真实比例看 `SIMD_OCR_GPU_TIME` 的墙钟分解。 -- CI：osx-arm64 smoke（macos-26 runner，paravirt GPU）以 `--engine metal` 跑 tiny 20 张，断言 stderr 含 `[metal] device:` breadcrumb、无 `fallback`、逐图 `detected` 与 CPU 一致、逐图 `texts` 与 CPU 的编辑距离 ≤3/图且合计 ≤8（fp16 噪声允许量）；其他 OS 的测试不会走到 Metal 路径。 +- CI：osx-arm64 smoke（macos-26 runner，paravirt GPU）以 `--engine metal` 跑 tiny 20 张，断言 stderr 含 `[metal] device:`、无 `fallback`、逐图 `detected` 与 CPU 一致；不比较文本，文本漂移靠本文的逐图对比数据把关（已接受的限制：CI 发现不了「框数对、文本错」类回归）；其他 OS 的测试不会走到 Metal 路径。 From 2a5ec513db8bd7ffc47ecd3f45df2a3344a1fc18 Mon Sep 17 00:00:00 2001 From: ZHOU Jie Date: Mon, 28 Sep 2026 19:45:27 -0700 Subject: [PATCH 8/8] Drop orphaned Streamed-run callback doc comments after delegate move to CtcUnitsReady --- src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalRunner.cs | 4 ---- src/Sdcb.SimdPaddleOCR/Backends/Vulkan/GpuRunner.cs | 4 ---- 2 files changed, 8 deletions(-) diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalRunner.cs b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalRunner.cs index 064c8b3..845964d 100644 --- a/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalRunner.cs +++ b/src/Sdcb.SimdPaddleOCR/Backends/Metal/MetalRunner.cs @@ -37,10 +37,6 @@ internal sealed class MetalSchedule public long LastTick; } -/// Streamed-run callback: units [first, first+count) of result -/// (unit i at offsets[i]) are complete. Return false to stop receiving -/// further batches. - /// /// Per-session Metal runtime over a shared : /// owns the fp16 arena, fp32 input/output buffers and the partials scratch — diff --git a/src/Sdcb.SimdPaddleOCR/Backends/Vulkan/GpuRunner.cs b/src/Sdcb.SimdPaddleOCR/Backends/Vulkan/GpuRunner.cs index 56ecf6d..e893649 100644 --- a/src/Sdcb.SimdPaddleOCR/Backends/Vulkan/GpuRunner.cs +++ b/src/Sdcb.SimdPaddleOCR/Backends/Vulkan/GpuRunner.cs @@ -36,10 +36,6 @@ internal sealed class GpuSchedule public long LastTick; } -/// Streamed-run callback: units [first, first+count) of result -/// (unit i at offsets[i]) are complete. Return false to stop receiving -/// further batches. - /// /// Per-session GPU runtime over a shared : owns /// the fp16 arena, fp32 input/output buffers, partials scratch, a private