ayo
7c5d3c9157
Add native K/M/RR inference engines, HIP kernels, replay and validation harnesses, and Windows DLSS setup and runtime support. Include GPU profiling tools, build and packaging scripts, platform integration patches, and implementation documentation.
56 lines
2.1 KiB
Python
56 lines
2.1 KiB
Python
"""Shared numeric helpers for the Ray Reconstruction port: e4m3 codes, f16, and ZLUDA's m16n8k32 accumulation."""
|
|
import numpy as np
|
|
|
|
_E4M3 = np.zeros(256)
|
|
for _c in range(256):
|
|
_s, _e, _m = _c >> 7, (_c >> 3) & 15, _c & 7
|
|
_v = np.nan if (_e == 15 and _m == 7) else (_m * 2.0 ** -9 if _e == 0 else (1 + _m / 8) * 2.0 ** (_e - 7))
|
|
_E4M3[_c] = -_v if _s else _v
|
|
|
|
|
|
def e4m3(codes):
|
|
return _E4M3[np.asarray(codes, np.uint8)]
|
|
|
|
|
|
def f16(bits):
|
|
return np.asarray(bits, np.uint16).view(np.float16).astype(np.float64)
|
|
|
|
|
|
def f16bits(x):
|
|
return np.asarray(x, np.float64).astype(np.float16).view(np.uint16)
|
|
|
|
|
|
def mma_k32(a_codes, b_codes, c_f16bits, ksteps=None):
|
|
"""D = A.B + C with ZLUDA's lowering: per 32-deep k-step, f32 sum of the exact products and C, one f16 rounding;
|
|
the next k-step takes that f16 value as C. (f64 stands in for f32: the products and short sums are exact.)"""
|
|
a, b = e4m3(a_codes), e4m3(b_codes)
|
|
acc = f16(c_f16bits)
|
|
for k in range(0, a.shape[1], 32):
|
|
acc = (a[:, k:k + 32] @ b[k:k + 32, :] + acc).astype(np.float32).astype(np.float16).astype(np.float64)
|
|
return f16bits(acc)
|
|
|
|
# cvt.rn.satfinite.e4m3 of every f16 bit pattern. Generate this deterministic
|
|
# format table locally; never ship an extracted or generated binary table.
|
|
def _f16_to_e4m3_table():
|
|
h = np.arange(65536, dtype=np.uint32)
|
|
mag = h & 0x7fff
|
|
normal = (mag - 0x2000 + 0x3f + ((mag >> 7) & 1)) >> 7
|
|
exponent = mag >> 10
|
|
significand = (mag & 0x3ff) | np.where(exponent != 0, 0x400, 0)
|
|
clamped_exponent = np.clip(exponent, 1, 8)
|
|
shift = 16 - clamped_exponent
|
|
subnormal = (significand + (np.left_shift(1, shift - 1)) - 1 + ((significand >> shift) & 1)) >> shift
|
|
code = np.where(mag >= 0x2400, normal, subnormal)
|
|
code = np.where(mag > 0x5f00, 0x7e, code)
|
|
result = ((h >> 8) & 0x80) | code
|
|
result = np.where(mag > 0x7c00, 0x7f, result) # All NaNs map to the canonical positive NaN.
|
|
return result.astype(np.uint8)
|
|
|
|
|
|
_Q8 = _f16_to_e4m3_table()
|
|
|
|
|
|
def q8(x_f16):
|
|
"""f16 values (float array) -> e4m3 codes"""
|
|
return _Q8[np.asarray(x_f16, np.float64).astype(np.float16).view(np.uint16)]
|