199 lines
11 KiB
Swift
199 lines
11 KiB
Swift
// family-probe: the step cost of every reserve candidate family of spec 1.13.2 on Apple silicon, standalone (no pack).
|
|
// Counter ASIC 3.0 item 6 (docs/plans/counter-asic-3-reserve.md), 6 October 2026. Same method as dot4-probe.swift
|
|
// (docs/analysis/int8-matrix-family.md section 4): a dependent chain of one op per step per lane, 1,048,576 lanes x
|
|
// 4,096 steps, best of N, GPU start-to-end time, bit-exact against a CPU reference on two whole 32-lane SIMD groups.
|
|
//
|
|
// Every chain has the dot4 probe's glue: acc = OP(acc, x, y); x = x * K + acc; y = rotate(y, 7) ^ (acc + s). The
|
|
// reference rows are the live families: alu (the add-xor-rotate chain of the 9070 XT bench-log entry, 5 ops per step
|
|
// counted, no acc), rotr (dst = rotr(src, src2 AND 31), the live rotr_var text), shflx (dst = dst XOR src of lane
|
|
// (lane XOR 8), the live shfl). The candidate rows are the seven families of 1.13.2:
|
|
// shl acc = y << (x AND 31) variable left shift
|
|
// shr acc = y >> (x AND 31) variable logical right shift
|
|
// bfe acc = extract_bits(y, 7, 13) bit-field extract, immediate offset and width (bfec = the C form)
|
|
// andn acc = y AND NOT x
|
|
// perm acc = bytes (y.b1, y.b3, y.b0, y.b2) byte permute by an immediate selector
|
|
// popc acc = acc + popcount(x) popcount folded by add
|
|
// clz acc = acc + clz(x) count-leading-zeros folded by add (clz(0) = 32)
|
|
// sel acc = bit 5 of y ? x : acc three-register select
|
|
// shfla acc = acc XOR x of lane ((lane + 3) mod 32) the second shuffle form
|
|
// and for comparison the dp4a-class rows of the dot4 probe (dot4u unsigned, dot4s signed, both emulated on Apple).
|
|
// The shifted and extracted forms take y (fresh every step) as the shifted register so the chain never drains to 0;
|
|
// the op count per step is the family's one op plus the same glue in every row. mm8 as a chain (Metal 4 matmul2d) is
|
|
// owed: this toolchain (Swift 5.8) has no Metal 4 tensor API.
|
|
//
|
|
// Build: swiftc -O -o family-probe family-probe.swift -framework Metal
|
|
// Run: ./family-probe [--lanes N] [--steps N] [--reps N]
|
|
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; }
|
|
inline uint rotr_var(uint x, uint n) { n &= 31u; return (x >> n) | (x << ((32u - n) & 31u)); }
|
|
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;
|
|
}
|
|
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;
|
|
}
|
|
|
|
#define CHAIN(NAME, OP) \\
|
|
kernel void NAME(constant uint& steps [[buffer(0)]], constant uint& seed [[buffer(1)]], device uint* out [[buffer(2)]], \\
|
|
uint g [[thread_position_in_grid]], ushort lane [[thread_index_in_simdgroup]]) { \\
|
|
uint x = pm_mix(g ^ seed), y = x ^ 0x5bd1e995u; uint acc = pm_mix(x); \\
|
|
for (uint s = 0u; s < steps; ++s) { OP; x = x * 0x9E3779B1u + acc; y = rotate(y, 7u) ^ (acc + s); } \\
|
|
out[g] = acc ^ x ^ y; \\
|
|
}
|
|
|
|
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;
|
|
}
|
|
CHAIN(probe_rotr, acc = rotr_var(y, x))
|
|
CHAIN(probe_shflx, acc = acc ^ simd_shuffle_xor(x, (ushort)8))
|
|
CHAIN(probe_shl, acc = y << (x & 31u))
|
|
CHAIN(probe_shr, acc = y >> (x & 31u))
|
|
CHAIN(probe_bfe, acc = extract_bits(y, 7u, 13u))
|
|
CHAIN(probe_bfec, acc = (y >> 7u) & 0x1fffu)
|
|
CHAIN(probe_andn, acc = y & ~x)
|
|
CHAIN(probe_perm, uchar4 b_ = as_type<uchar4>(y); acc = as_type<uint>(uchar4(b_.y, b_.w, b_.x, b_.z)))
|
|
CHAIN(probe_popc, acc = acc + popcount(x))
|
|
CHAIN(probe_clz, acc = acc + clz(x))
|
|
CHAIN(probe_sel, acc = select(acc, x, ((y >> 5u) & 1u) != 0u))
|
|
CHAIN(probe_shfla, acc = acc ^ simd_shuffle(x, (ushort)((lane + 3u) & 31u)))
|
|
CHAIN(probe_dot4u, acc = dot4_u(x, y, acc))
|
|
CHAIN(probe_dot4s, acc = uint(dot4_s(x, y, int(acc))))
|
|
"""
|
|
|
|
// ---- CPU reference: one whole 32-lane SIMD group (lanes g0 .. g0+31, lane = g AND 31), 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 rotl(_ v: UInt32, _ n: UInt32) -> UInt32 { (v << n) | (v >> (32 - n)) }
|
|
func rotrVar(_ x: UInt32, _ nIn: UInt32) -> UInt32 { let n = nIn & 31; return (x >> n) | (x << ((32 - n) & 31)) }
|
|
func clz32(_ x: UInt32) -> UInt32 { x == 0 ? 32 : UInt32(x.leadingZeroBitCount) }
|
|
func perm(_ y: UInt32) -> UInt32 {
|
|
((y >> 8) & 0xff) | (((y >> 24) & 0xff) << 8) | ((y & 0xff) << 16) | (((y >> 16) & 0xff) << 24)
|
|
}
|
|
func dot4sRef(_ a: UInt32, _ b: UInt32, _ acc: UInt32) -> UInt32 {
|
|
var r = Int32(bitPattern: 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 UInt32(bitPattern: 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
|
|
}
|
|
// the 32 outputs of the SIMD group whose first lane is g0
|
|
func warpRef(kernel: String, g0: UInt32, seed: UInt32, steps: UInt32) -> [UInt32] {
|
|
var x = [UInt32](repeating: 0, count: 32), y = [UInt32](repeating: 0, count: 32), acc = [UInt32](repeating: 0, count: 32)
|
|
for l in 0..<32 { x[l] = pmMix((g0 + UInt32(l)) ^ seed); y[l] = x[l] ^ 0x5bd1e995; acc[l] = pmMix(x[l]) }
|
|
if kernel == "probe_alu" {
|
|
for s in 0..<steps { for l in 0..<32 { x[l] = x[l] &* 0x9E3779B1 &+ rotl(y[l], 7); y[l] = (y[l] ^ x[l]) &+ s } }
|
|
return (0..<32).map { x[$0] ^ y[$0] }
|
|
}
|
|
for s in 0..<steps {
|
|
let xs = x // shuffles read the other lanes' x of this step
|
|
for l in 0..<32 {
|
|
let xv = x[l], yv = y[l]
|
|
switch kernel {
|
|
case "probe_rotr": acc[l] = rotrVar(yv, xv)
|
|
case "probe_shflx": acc[l] = acc[l] ^ xs[l ^ 8]
|
|
case "probe_shl": acc[l] = yv << (xv & 31)
|
|
case "probe_shr": acc[l] = yv >> (xv & 31)
|
|
case "probe_bfe", "probe_bfec": acc[l] = (yv >> 7) & 0x1fff
|
|
case "probe_andn": acc[l] = yv & ~xv
|
|
case "probe_perm": acc[l] = perm(yv)
|
|
case "probe_popc": acc[l] = acc[l] &+ UInt32(xv.nonzeroBitCount)
|
|
case "probe_clz": acc[l] = acc[l] &+ clz32(xv)
|
|
case "probe_sel": acc[l] = ((yv >> 5) & 1) != 0 ? xv : acc[l]
|
|
case "probe_shfla": acc[l] = acc[l] ^ xs[(l + 3) & 31]
|
|
case "probe_dot4u": acc[l] = dot4uRef(xv, yv, acc[l])
|
|
case "probe_dot4s": acc[l] = dot4sRef(xv, yv, acc[l])
|
|
default: fatalError("no reference for \(kernel)")
|
|
}
|
|
x[l] = xv &* 0x9E3779B1 &+ acc[l]
|
|
y[l] = rotl(yv, 7) ^ (acc[l] &+ s)
|
|
}
|
|
}
|
|
return (0..<32).map { acc[$0] ^ x[$0] ^ y[$0] }
|
|
}
|
|
|
|
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)!
|
|
let names = ["probe_alu", "probe_rotr", "probe_shflx", "probe_shl", "probe_shr", "probe_bfe", "probe_bfec", "probe_andn", "probe_perm",
|
|
"probe_popc", "probe_clz", "probe_sel", "probe_shfla", "probe_dot4u", "probe_dot4s"]
|
|
print("family-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 | ns per step | ratio to alu | warps 0 and last ok |")
|
|
print("|---|---|---|---|---|---|---|---|")
|
|
var aluBest = 0.0
|
|
for name in names {
|
|
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 }
|
|
let p = outBuf.contents().bindMemory(to: UInt32.self, capacity: lanes)
|
|
for g0 in [UInt32(0), UInt32(lanes - 32)] {
|
|
let want = warpRef(kernel: name, g0: g0, seed: seed, steps: steps)
|
|
for l in 0..<32 where p[Int(g0) + l] != want[l] {
|
|
okAll = false
|
|
print("MISMATCH \(name) lane \(g0 + UInt32(l)): gpu \(String(p[Int(g0) + l], radix: 16)) cpu \(String(want[l], radix: 16))")
|
|
break
|
|
}
|
|
}
|
|
}
|
|
if name == "probe_alu" { aluBest = best }
|
|
let stepsPerS = Double(lanes) * Double(steps) / (best / 1000.0)
|
|
let ratio = aluBest > 0 ? best / aluBest : 0
|
|
print(String(format: "| %@ | %d | %u | %.3f | %.2f | %.3f | %.2f | %@ |", name, lanes, steps, best, stepsPerS / 1e9, best * 1e6 / Double(steps), ratio, okAll ? "yes" : "NO"))
|
|
print(String(format: "RESULT FAMILY vendor=apple device=\"%@\" kernel=%@ lanes=%d steps=%u best_ms=%.3f gsteps_per_s=%.2f ns_per_step=%.3f ratio_alu=%.3f ok=%d", dev.name, name, lanes, steps, best, stepsPerS / 1e9, best * 1e6 / Double(steps), ratio, okAll ? 1 : 0))
|
|
}
|
|
print("family-probe: done")
|