open-standards/wai/wai-sink-swift/Tests/WaiSinkTests/IntDecodeTests.swift

// The integer decode a native sink runs in its own process (native-sink §13), through the same C
// interface every native integration links: each staged vector of each integer capability decodes
// with the device's GPU to the bytes of the decode on the CPU, and each answer's dispatch report
// names the unit that ran it. On a host whose GPU does not pass the start-up check the GPU half
// prints SKIPPED, and fails when WAI_SINK_REQUIRE_GPU is set.
import CryptoKit
import XCTest
import WaiSink

final class IntDecodeTests: XCTestCase {
    struct Case {
        let capability: String
        let name: String
        let payload: Data?
        let envelope: Data?
        let set: [String: Data]
    }

    func dir() -> URL {
        #if SWIFT_PACKAGE
        return Bundle.module.resourceURL!.appendingPathComponent("Resources/int")
        #else
        return Bundle(for: TestBundleMarker.self).resourceURL!.appendingPathComponent("int")
        #endif
    }

    func cases() throws -> [Case] {
        let dir = self.dir()
        let index = try JSONSerialization.jsonObject(with: Data(contentsOf: dir.appendingPathComponent("cases.json"))) as! [[String: Any]]
        return try index.map { c in
            let file = { (n: String) in try Data(contentsOf: dir.appendingPathComponent(n)) }
            var set: [String: Data] = [:]
            for (k, v) in c["set"] as! [String: String] { set[k] = try file(v) }
            return Case(capability: c["capability"] as! String, name: c["name"] as! String,
                        payload: try (c["payload"] as? String).map(file), envelope: try (c["envelope"] as? String).map(file), set: set)
        }
    }

    func decode(_ d: WaiIntDecoder, _ c: Case) throws -> WaiIntDecoded {
        if let e = c.envelope { return try d.decode(envelope: e, set: c.set) }
        return try d.decode(capability: c.capability, payload: c.payload!, set: c.set)
    }

    func testEveryCapabilityDecodesOnTheGPUToTheCPUsBytes() throws {
        let cases = try cases()
        XCTAssertGreaterThanOrEqual(cases.count, 8)
        let cpu = try WaiIntDecoder(gpu: .off)
        XCTAssertEqual(try cpu.gpu()["state"] as? String, "off")
        let dev = try WaiIntDecoder(gpu: .auto, unit: .gpuFirst)
        let gpu = try dev.gpu()
        let ready = gpu["state"] as? String == "ready"
        var onGPU = 0
        for c in cases {
            let want = try decode(cpu, c)
            XCTAssertNil(want.refused, "\(c.name): refused on the CPU")
            XCTAssertFalse(want.bytes.isEmpty, c.name)
            XCTAssertEqual(want.unit, "cpu", c.name)
            let got = try decode(dev, c)
            XCTAssertEqual(got.bytes, want.bytes, "\(c.name): the GPU's bytes are not the CPU's")
            if c.capability == "wai.video.int_motion" {
                // No network: neural dispatch does not apply, and it decodes on the CPU.
                XCTAssertEqual(got.unit, "cpu", c.name)
                XCTAssertTrue(got.answer["dispatch"] is NSNull, c.name)
                continue
            }
            let ran = got.dispatch["ran"] as? [String: Any]
            XCTAssertEqual(ran?["tier"] as? String, "decode-equivalence", c.name)
            if ready {
                XCTAssertEqual(got.unit, "gpu", c.name)
                XCTAssertEqual(ran?["backend"] as? String, "wgsl", c.name)
                onGPU += 1
            } else {
                XCTAssertEqual(got.unit, "cpu", c.name)
                let refused = (got.dispatch["refused"] as? [[String: Any]])?.first
                XCTAssertEqual(refused?["unit"] as? String, "gpu", c.name)
            }
        }
        let adapter = gpu["adapter"] as? [String: Any]
        print("WAI-SINK int decode: \(cases.count) vectors, \(onGPU) on the GPU (state \(gpu["state"] ?? "?"), adapter \(adapter?["name"] ?? "none"), backend \(adapter?["backend"] ?? "none")), every one the CPU's bytes")
        if !ready {
            print("WAI-SINK int decode: GPU SKIPPED: \(gpu["why"] ?? "")")
            XCTAssertNil(ProcessInfo.processInfo.environment["WAI_SINK_REQUIRE_GPU"], "WAI_SINK_REQUIRE_GPU is set and the decoder has no GPU")
        }
    }

    // The integer frame forms (native-sink §13.6): a stream of each, its objects in order, with
    // the GPU and on the CPU; every frame its committed SHA-256, the GPU's bytes the CPU's, each
    // answer's `frame` naming the object's index.
    func testEveryFrameFormDecodesWithTheGPUToTheCPUsFrames() throws {
        let dir = self.dir()
        let index = try JSONSerialization.jsonObject(with: Data(contentsOf: dir.appendingPathComponent("frames.json"))) as! [[String: Any]]
        XCTAssertEqual(index.count, 3)
        let cpu = try WaiIntDecoder(gpu: .off)
        let dev = try WaiIntDecoder(gpu: .auto, unit: .gpuFirst)
        var frames = 0, onGPU = 0
        for c in index {
            let name = c["name"] as! String
            var set: [String: Data] = [:]
            for (k, v) in c["set"] as! [String: String] { set[k] = try Data(contentsOf: dir.appendingPathComponent(v)) }
            let cap = c["capability"] as! String
            let listing = try (c["listing"] as? String).map { try Data(contentsOf: dir.appendingPathComponent($0)) }
            let sc = try WaiIntFrameStream(capability: cap, set: set, listing: listing)
            let sg = try WaiIntFrameStream(capability: cap, set: set, listing: listing)
            for (i, o) in (c["objects"] as! [[String: String]]).enumerated() {
                let env = try Data(contentsOf: dir.appendingPathComponent(o["envelope"]!))
                let want = try cpu.decode(frame: env, stream: sc)
                XCTAssertNil(want.refused, "\(name) \(i): refused on the CPU")
                let hash = SHA256.hash(data: want.bytes).map { String(format: "%02x", $0) }.joined()
                XCTAssertEqual(hash, o["frame_sha256"], "\(name) \(i): the frame's SHA-256")
                XCTAssertEqual(want.frame?["index"] as? Int, i, "\(name) \(i)")
                let got = try dev.decode(frame: env, stream: sg)
                XCTAssertEqual(got.bytes, want.bytes, "\(name) \(i): the GPU's frame is not the CPU's")
                XCTAssertEqual(got.frame?["index"] as? Int, i, "\(name) \(i)")
                if got.unit == "gpu" { onGPU += 1 }
                frames += 1
            }
        }
        print("WAI-SINK int frames: \(frames) frames of 3 frame forms, \(onGPU) on the GPU, every one the CPU's bytes and its committed SHA-256")
    }

    func testAFloatCapabilityHasNoIntegerPath() throws {
        let d = try WaiIntDecoder(gpu: .auto)
        let o = try d.decode(capability: "wai.neural.bmshj2018", payload: Data(count: 32))
        XCTAssertNil(o.unit)
        XCTAssertEqual(o.refused, "dispatch_no_path")
        for r in o.dispatch["refused"] as! [[String: Any]] {
            XCTAssertEqual(r["reason"] as? String, "dispatch_arithmetic_mismatch")
        }
    }
}