165 lines
7.6 KiB
Swift
165 lines
7.6 KiB
Swift
// dot4-probe: dp4a-class throughput on Apple silicon, standalone (no pack, no lottery kernel).
|
|
// Counter ASIC 2.0 layer 7 (docs/analysis/int8-matrix-family.md), 5 October 2026.
|
|
//
|
|
// Metal has no dp4a intrinsic and no integer simdgroup_matrix (MSL 4.1 section 2.4 lists half, bfloat and float only;
|
|
// the Metal 4 tensor op matmul2d does carry char x char -> int, MSL 4.1 table 7.3, measured separately when it is).
|
|
// So the per-lane dot4 here is the scalar emulation a conforming Apple miner would run: four sign-extended bytes of
|
|
// each operand multiplied and summed into a wrapping int32 accumulator, exactly the PTX dp4a semantics
|
|
// (PTX ISA 9.4 section 9.7.1.24: d = c; d += Va[i] * Vb[i] for i in 0..3, bytes sign- or zero-extended).
|
|
//
|
|
// Two kernels, same shape as the OpenCL --memprobe ALU chain (proto-opencl/host.c, probe_alu: 1,048,576 lanes x 4,096
|
|
// steps, best of 3):
|
|
// alu x = x * K + rotate(y, 7); y = (y ^ x) + s the card's integer baseline, 5 ops per step counted
|
|
// dot4 acc = dot4(x, y, acc); x = x * K + acc; y = rotate(y, 7) ^ acc one dependent dot4 per step per lane
|
|
// Rates: G steps/s per lane-step, so G dot4/s for the second kernel. Timing is the command buffer's GPU start to end.
|
|
//
|
|
// Build: swiftc -O -o dot4-probe dot4-probe.swift -framework Metal
|
|
// Run: ./dot4-probe [--lanes N] [--steps N] [--reps N] [--signed|--unsigned]
|
|
import Foundation
|
|
import Metal
|
|
|
|
let source = """
|
|
#include <metal_stdlib>
|
|
using namespace metal;
|
|
|
|
inline uint pm_mix(uint x) { x ^= x >> 16; x *= 0x7feb352du; x ^= x >> 15; x *= 0x846ca68bu; x ^= x >> 16; return x; }
|
|
|
|
// dp4a, signed bytes, wrapping int32 accumulate: the exact PTX dp4a.s32.s32 semantics.
|
|
inline int dot4_s(uint a, uint b, int acc) {
|
|
int4 va = int4(as_type<char4>(a));
|
|
int4 vb = int4(as_type<char4>(b));
|
|
return acc + va.x * vb.x + va.y * vb.y + va.z * vb.z + va.w * vb.w;
|
|
}
|
|
// dp4a, unsigned bytes, wrapping uint32 accumulate: dp4a.u32.u32.
|
|
inline uint dot4_u(uint a, uint b, uint acc) {
|
|
uint4 va = uint4(as_type<uchar4>(a));
|
|
uint4 vb = uint4(as_type<uchar4>(b));
|
|
return acc + va.x * vb.x + va.y * vb.y + va.z * vb.z + va.w * vb.w;
|
|
}
|
|
|
|
kernel void probe_alu(constant uint& steps [[buffer(0)]], constant uint& seed [[buffer(1)]],
|
|
device uint* out [[buffer(2)]], uint g [[thread_position_in_grid]]) {
|
|
uint x = pm_mix(g ^ seed), y = x ^ 0x5bd1e995u;
|
|
for (uint s = 0u; s < steps; ++s) { x = x * 0x9E3779B1u + rotate(y, 7u); y = (y ^ x) + s; }
|
|
out[g] = x ^ y;
|
|
}
|
|
|
|
kernel void probe_dot4s(constant uint& steps [[buffer(0)]], constant uint& seed [[buffer(1)]],
|
|
device uint* out [[buffer(2)]], uint g [[thread_position_in_grid]]) {
|
|
uint x = pm_mix(g ^ seed), y = x ^ 0x5bd1e995u;
|
|
int acc = int(pm_mix(x));
|
|
for (uint s = 0u; s < steps; ++s) {
|
|
acc = dot4_s(x, y, acc);
|
|
x = x * 0x9E3779B1u + uint(acc);
|
|
y = rotate(y, 7u) ^ (uint(acc) + s);
|
|
}
|
|
out[g] = uint(acc) ^ x ^ y;
|
|
}
|
|
|
|
kernel void probe_dot4u(constant uint& steps [[buffer(0)]], constant uint& seed [[buffer(1)]],
|
|
device uint* out [[buffer(2)]], uint g [[thread_position_in_grid]]) {
|
|
uint x = pm_mix(g ^ seed), y = x ^ 0x5bd1e995u;
|
|
uint acc = pm_mix(x);
|
|
for (uint s = 0u; s < steps; ++s) {
|
|
acc = dot4_u(x, y, acc);
|
|
x = x * 0x9E3779B1u + acc;
|
|
y = rotate(y, 7u) ^ (acc + s);
|
|
}
|
|
out[g] = acc ^ x ^ y;
|
|
}
|
|
"""
|
|
|
|
// CPU reference of the dot4 chain for one lane, to check the kernel is the arithmetic it claims (bit-exact).
|
|
func pmMix(_ v: UInt32) -> UInt32 {
|
|
var x = v
|
|
x ^= x >> 16; x = x &* 0x7feb352d; x ^= x >> 15; x = x &* 0x846ca68b; x ^= x >> 16
|
|
return x
|
|
}
|
|
func dot4sRef(_ a: UInt32, _ b: UInt32, _ acc: Int32) -> Int32 {
|
|
var r = acc
|
|
for i in 0..<4 {
|
|
let ba = Int32(Int8(truncatingIfNeeded: a >> (8 * UInt32(i))))
|
|
let bb = Int32(Int8(truncatingIfNeeded: b >> (8 * UInt32(i))))
|
|
r = r &+ ba &* bb
|
|
}
|
|
return r
|
|
}
|
|
func dot4uRef(_ a: UInt32, _ b: UInt32, _ acc: UInt32) -> UInt32 {
|
|
var r = acc
|
|
for i in 0..<4 {
|
|
let ba = UInt32(UInt8(truncatingIfNeeded: a >> (8 * UInt32(i))))
|
|
let bb = UInt32(UInt8(truncatingIfNeeded: b >> (8 * UInt32(i))))
|
|
r = r &+ ba &* bb
|
|
}
|
|
return r
|
|
}
|
|
func rotl(_ v: UInt32, _ n: UInt32) -> UInt32 { (v << n) | (v >> (32 - n)) }
|
|
func laneRef(kernel: String, g: UInt32, seed: UInt32, steps: UInt32) -> UInt32 {
|
|
var x = pmMix(g ^ seed), y = x ^ 0x5bd1e995
|
|
switch kernel {
|
|
case "probe_alu":
|
|
for s in 0..<steps { x = x &* 0x9E3779B1 &+ rotl(y, 7); y = (y ^ x) &+ s }
|
|
return x ^ y
|
|
case "probe_dot4s":
|
|
var acc = Int32(bitPattern: pmMix(x))
|
|
for s in 0..<steps { acc = dot4sRef(x, y, acc); x = x &* 0x9E3779B1 &+ UInt32(bitPattern: acc); y = rotl(y, 7) ^ (UInt32(bitPattern: acc) &+ s) }
|
|
return UInt32(bitPattern: acc) ^ x ^ y
|
|
default:
|
|
var acc = pmMix(x)
|
|
for s in 0..<steps { acc = dot4uRef(x, y, acc); x = x &* 0x9E3779B1 &+ acc; y = rotl(y, 7) ^ (acc &+ s) }
|
|
return acc ^ x ^ y
|
|
}
|
|
}
|
|
|
|
var lanes = 1 << 20, steps: UInt32 = 4096, reps = 3
|
|
var args = Array(CommandLine.arguments.dropFirst())
|
|
while !args.isEmpty {
|
|
let a = args.removeFirst()
|
|
switch a {
|
|
case "--lanes": lanes = Int(args.removeFirst())!
|
|
case "--steps": steps = UInt32(args.removeFirst())!
|
|
case "--reps": reps = Int(args.removeFirst())!
|
|
default: print("unknown argument \(a)"); exit(2)
|
|
}
|
|
}
|
|
|
|
guard let dev = MTLCreateSystemDefaultDevice() else { print("no Metal device"); exit(1) }
|
|
let lib: MTLLibrary
|
|
do { lib = try dev.makeLibrary(source: source, options: nil) } catch { print("compile failed: \(error)"); exit(1) }
|
|
let queue = dev.makeCommandQueue()!
|
|
let outBuf = dev.makeBuffer(length: lanes * 4, options: .storageModeShared)!
|
|
print("dot4-probe on \(dev.name), lanes \(lanes), steps \(steps), best of \(reps), GPU start-to-end time")
|
|
print("| kernel | lanes | steps | best ms | G steps/s (= G dot4/s for dot4 rows) | ns per dependent step | lane 0 ok |")
|
|
print("|---|---|---|---|---|---|---|")
|
|
for name in ["probe_alu", "probe_dot4s", "probe_dot4u"] {
|
|
let fn = lib.makeFunction(name: name)!
|
|
let pso = try! dev.makeComputePipelineState(function: fn)
|
|
let tg = min(256, pso.maxTotalThreadsPerThreadgroup)
|
|
var best = Double.infinity
|
|
var okAll = true
|
|
for r in 0..<reps {
|
|
var st = steps
|
|
var seed = UInt32(0x2468ace) &+ UInt32(r) &* 0x9E3779B9
|
|
let cb = queue.makeCommandBuffer()!
|
|
let enc = cb.makeComputeCommandEncoder()!
|
|
enc.setComputePipelineState(pso)
|
|
enc.setBytes(&st, length: 4, index: 0)
|
|
enc.setBytes(&seed, length: 4, index: 1)
|
|
enc.setBuffer(outBuf, offset: 0, index: 2)
|
|
enc.dispatchThreads(MTLSize(width: lanes, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: tg, height: 1, depth: 1))
|
|
enc.endEncoding()
|
|
cb.commit()
|
|
cb.waitUntilCompleted()
|
|
let ms = (cb.gpuEndTime - cb.gpuStartTime) * 1000.0
|
|
if ms < best { best = ms }
|
|
// bit-exactness of the kernel against the CPU reference on two lanes
|
|
let p = outBuf.contents().bindMemory(to: UInt32.self, capacity: lanes)
|
|
for g in [UInt32(0), UInt32(lanes - 1)] {
|
|
let want = laneRef(kernel: name, g: g, seed: seed, steps: steps)
|
|
if p[Int(g)] != want { okAll = false; print("MISMATCH \(name) lane \(g): gpu \(String(p[Int(g)], radix: 16)) cpu \(String(want, radix: 16))") }
|
|
}
|
|
}
|
|
let stepsPerS = Double(lanes) * Double(steps) / (best / 1000.0)
|
|
print(String(format: "| %@ | %d | %u | %.3f | %.2f | %.3f | %@ |", name, lanes, steps, best, stepsPerS / 1e9, best * 1e6 / Double(steps), okAll ? "yes" : "NO"))
|
|
}
|
|
print("dot4-probe: done")
|