mawei/onepilot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update251203 (#233)

Browse Source
This commit is contained in:
carrot
2025-12-03 10:28:27 +09:00
committed by GitHub
parent d6899edd97
commit c5ebcbcb97
347 changed files with 8678 additions and 13489 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad_repo/extra/amdpci/proclogs.py
+1 -1
View File
@@ -37,7 +37,7 @@ def main():
dev = PCIIface(None, 0)
for x, y in dev.dev_impl.__dict__.items():
if isinstance(y, AMRegister):
for inst, addr in y.addr.keys(): reg_names[addr] = f"{x}, xcc={inst}"
for inst, addr in y.addr.items(): reg_names[addr] = f"{x}, xcc={inst}"
with open(sys.argv[1], 'r') as f:
log_content = log_content_them = f.read()
tinygrad_repo/extra/archprobe.py
+1 -1
View File
@@ -1,7 +1,7 @@
# copying the kernels from https://github.com/microsoft/ArchProbe into Python
import numpy as np
import pickle
from tinygrad.runtime.ops_gpu import CLProgram, CLBuffer
from tinygrad.runtime.ops_cl import CLProgram, CLBuffer
from tinygrad import dtypes
from tqdm import trange, tqdm
from matplotlib import pyplot as plt
tinygrad_repo/extra/assembly/assembly_rdna.py
+1 -1
View File
@@ -4,7 +4,7 @@ from tinygrad import dtypes
from tinygrad.codegen.assembly import AssemblyCodegen, Register
from tinygrad.codegen.opt.kernel import Ops
from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH
from tinygrad.runtime.ops_cl import ROCM_LLVM_PATH
# ugh, is this really needed?
from extra.helpers import enable_early_exec
tinygrad_repo/extra/assembly/rocm/rdna3/asm.py
+1 -1
View File
@@ -5,7 +5,7 @@ from tinygrad.helpers import colored
from extra.helpers import enable_early_exec
early_exec = enable_early_exec()
from tinygrad.runtime.ops_gpu import CLProgram, CLBuffer, ROCM_LLVM_PATH
from tinygrad.runtime.ops_cl import CLProgram, CLBuffer, ROCM_LLVM_PATH
ENABLE_NON_ASM = False
tinygrad_repo/extra/backends/clang_graph.py
+4 -4
View File
@@ -10,13 +10,13 @@ from tinygrad.renderer.cstyle import ClangRenderer
render_dtype = ClangRenderer().render_dtype
class ClangGraph(GraphRunner):
def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, int]):
super().__init__(jit_cache, input_rawbuffers, var_vals)
if not all(isinstance(ji.prg, CompiledRunner) for ji in jit_cache): raise GraphException
prgs = '\n'.join(dedup([cast(CompiledRunner, ji.prg).p.src for ji in jit_cache]))
args = [f"{render_dtype(x.dtype)}* arg{i}" for i,x in enumerate(input_rawbuffers)]
args += sorted([f"int {v.expr}" for v in var_vals])
args += sorted([f"int {v}" for v in var_vals])
code = ["void batched("+','.join(args)+") {"]
for ji in jit_cache:
args = []
@@ -34,6 +34,6 @@ class ClangGraph(GraphRunner):
assert compiler is not None
self._prg = ClangProgram("batched", compiler.compile(prgs+"\n"+"\n".join(code))) # no point in caching the pointers
def __call__(self, rawbufs: List[Buffer], var_vals: Dict[Variable, int], wait=False):
def __call__(self, rawbufs: List[Buffer], var_vals: Dict[str, int], wait=False):
return cpu_time_execution(
lambda: self._prg(*[x._buf for x in rawbufs], *[x[1] for x in sorted(var_vals.items(), key=lambda x: x[0].expr)]), enable=wait)
lambda: self._prg(*[x._buf for x in rawbufs], *[x[1] for x in sorted(var_vals.items(), key=lambda x: x[0])]), enable=wait)
tinygrad_repo/extra/backends/hsa_graph.py
+4 -4
View File
@@ -26,7 +26,7 @@ class VirtAQLQueue(AQLQueue):
self.available_packet_slots -= 1
class HSAGraph(MultiGraphRunner):
def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, int]):
super().__init__(jit_cache, input_rawbuffers, var_vals)
# Check all jit items are compatible.
@@ -53,7 +53,7 @@ class HSAGraph(MultiGraphRunner):
self.ji_kargs_structs[j] = ji.prg._prg.args_struct_t.from_address(kernargs_ptrs[ji.prg.dev])
kernargs_ptrs[ji.prg.dev] += round_up(ctypes.sizeof(ji.prg._prg.args_struct_t), 16)
for i in range(len(ji.bufs)): self.ji_kargs_structs[j].__setattr__(f'f{i}', cast(Buffer, ji.bufs[i])._buf)
for i in range(len(ji.prg.p.vars)): self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[ji.prg.p.vars[i]])
for i in range(len(ji.prg.p.vars)): self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[ji.prg.p.vars[i].expr])
# Build queues.
self.virt_aql_queues: Dict[Compiled, VirtAQLQueue] = {dev:VirtAQLQueue(dev, 2*len(self.jit_cache)+16) for dev in self.devices}
@@ -106,7 +106,7 @@ class HSAGraph(MultiGraphRunner):
for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 0)
hsa.hsa_signal_silent_store_relaxed(self.finish_signal, 0)
def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int], wait=False) -> Optional[float]:
def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[str, int], wait=False) -> Optional[float]:
# Wait and restore signals
hsa.hsa_signal_wait_scacquire(self.finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 1)
@@ -123,7 +123,7 @@ class HSAGraph(MultiGraphRunner):
# Update var_vals
for j in self.jc_idx_with_updatable_var_vals:
for i,v in enumerate(cast(CompiledRunner, self.jit_cache[j].prg).p.vars):
self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v])
self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v.expr])
# Update launch dims
for j in self.jc_idx_with_updatable_launch_dims:
tinygrad_repo/extra/backends/rdna.py
+4 -4
View File
@@ -29,10 +29,10 @@ def uops_to_rdna(function_name:str, uops:UOpGraph) -> str:
r: Dict[UOp, str] = {}
for u in uops:
if u.uop == UOps.SPECIAL:
if u.arg[1].startswith("lidx"):
r[u] = f'v{u.arg[0]}'
elif u.arg[1].startswith("gidx"):
r[u] = f's{2+u.arg[0]}'
if u.arg.startswith("lidx"):
r[u] = f'v{u.src[0].arg}'
elif u.arg.startswith("gidx"):
r[u] = f's{2+u.src[0].arg}'
else:
raise NotImplementedError
elif u.uop == UOps.CONST:
tinygrad_repo/extra/export_model.py
+6 -3
View File
@@ -10,7 +10,7 @@ from tinygrad.uop.ops import Ops
import json
from collections import OrderedDict
EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "GPU"]
EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]
def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
@@ -67,11 +67,12 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
forward_args = ",".join(f"{dtype}{'*' if name not in symbolic_vars.values() else ''} {name}" for name,dtype,_ in (outputs+inputs if wasm else inputs+outputs))
if not wasm:
thread_id = 0 # NOTE: export does not support threading, thread_id is always 0
for name,cl in bufs_to_save.items():
weight = ''.join(["\\x%02X"%x for x in bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)}, {thread_id});" for (name, args, _global_size, _local_size) in statements] + ["}"]
return '\n'.join(headers + cprog)
else:
if bufs_to_save:
@@ -239,7 +240,9 @@ export default {model_name};
def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False):
assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"
with Context(JIT=2): run,special_names = jit_model(model, *inputs)
# NOTE: CPU_COUNT=1, since export does not support threading
with Context(JIT=2, CPU_COUNT=1): run,special_names = jit_model(model, *inputs)
functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
state = get_state_dict(model)
weight_names = {id(x.uop.base.realized): name for name, x in state.items()}
tinygrad_repo/extra/gemm/amd_uop_matmul.py
+34 -34
View File
@@ -65,7 +65,7 @@ def top_spec_kernel3():
c = a@b
sink = c.schedule()[-1].ast
L = 16
sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(dtypes.int, N//BM, 0), 2:UOp.range(dtypes.int, N//BN, 1)})
sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(N//BM, 0), 2:UOp.range(N//BN, 1)})
sink = graph_rewrite(sink, view_left+pm)
axis_types = (AxisType.GLOBAL, AxisType.LOCAL, AxisType.GLOBAL, AxisType.LOCAL, AxisType.REDUCE)
return sink.replace(arg=KernelInfo(name="top_"+to_colored(sink.full_shape, axis_types), axis_types=axis_types))
@@ -186,7 +186,7 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
c_regs = UOp(Ops.DEFINE_REG, dtypes.float.ptr(TM * nbIterWaveM * TN * nbIterWaveN), arg=2)
i = UOp.range(dtypes.int, c_regs.dtype.size, 16)
i = UOp.range(c_regs.dtype.size, 16)
init_store = c_regs[i].store(UOp.const(dtypes.float, 0.0), i)
if kernel4:
@@ -197,53 +197,53 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
kId = 0
# load from globals into locals
i = UOp.range(dtypes.int, nbReadsB, 0)
i = UOp.range(nbReadsB, 0)
index_x = BN * blockIdx_x + rBIdx
index_y = rBIdy + i * strideReadB + kId
Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)
i = UOp.range(dtypes.int, nbReadsA, 1)
i = UOp.range(nbReadsA, 1)
index_x = rAIdx + kId
index_y = BM * blockIdx_y + rAIdy + i * strideReadA
As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
# iterate over the middle chunk
kId_range = UOp.range(dtypes.int, N//BK-1, 2)
kId_range = UOp.range(N//BK-1, 2)
kId = kId_range*BK
barrier = UOp.barrier(As_store, Bs_store)
# load from globals into registers (next round)
i = UOp.range(dtypes.int, nbReadsB, 3)
i = UOp.range(nbReadsB, 3)
index_x = BN * blockIdx_x + rBIdx
index_y = rBIdy + i * strideReadB + kId + BK
regB_store = regB[i].store(b[N * index_y + index_x].load(), i)
i = UOp.range(dtypes.int, nbReadsA, 4)
i = UOp.range(nbReadsA, 4)
index_x = rAIdx + kId + BK
index_y = BM * blockIdx_y + rAIdy + i * strideReadA
regA_store = regA[i].store(a[N * index_y + index_x].load(), i)
def inner_loop(first_range, inp_dep=()):
# inner unroll
k = UOp.range(dtypes.int, BK, first_range+0)
k = UOp.range(BK, first_range+0)
# load from locals into registers
iterWave = UOp.range(dtypes.int, nbIterWaveN, first_range+1)
i = UOp.range(dtypes.int, TN, first_range+2)
iterWave = UOp.range(nbIterWaveN, first_range+1)
i = UOp.range(TN, first_range+2)
index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(*inp_dep), iterWave, i)
iterWave = UOp.range(dtypes.int, nbIterWaveM, first_range+3)
i = UOp.range(dtypes.int, TM, first_range+4)
iterWave = UOp.range(nbIterWaveM, first_range+3)
i = UOp.range(TM, first_range+4)
index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(*inp_dep), iterWave, i)
# do the GEMM math
iterWaveM = UOp.range(dtypes.int, nbIterWaveM, first_range+5)
yt = UOp.range(dtypes.int, TM, first_range+6)
iterWaveN = UOp.range(dtypes.int, nbIterWaveN, first_range+7)
xt = UOp.range(dtypes.int, TN, first_range+8)
iterWaveM = UOp.range(nbIterWaveM, first_range+5)
yt = UOp.range(TM, first_range+6)
iterWaveN = UOp.range(nbIterWaveN, first_range+7)
xt = UOp.range(TN, first_range+8)
x = iterWaveN * TN + xt
y = iterWaveM * TM + yt
c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
@@ -256,12 +256,12 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
sink = inner_loop(5, (barrier, regB_store, regA_store)).barrier()
# load from registers into locals
i = UOp.range(dtypes.int, nbReadsB, 14)
i = UOp.range(nbReadsB, 14)
index_x = BN * blockIdx_x + rBIdx
index_y = rBIdy + i * strideReadB + kId + BK
Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(regB[i].load(sink), i, kId_range)
i = UOp.range(dtypes.int, nbReadsA, 15)
i = UOp.range(nbReadsA, 15)
index_x = rAIdx + kId + BK
index_y = BM * blockIdx_y + rAIdy + i * strideReadA
As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(regA[i].load(sink), i, kId_range)
@@ -269,40 +269,40 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
# final iteration without the copy
sink = inner_loop(16, (UOp.barrier(Bs_store, As_store),))
else:
kId_range = UOp.range(dtypes.int, N//BK, 0)
kId_range = UOp.range(N//BK, 0)
kId = kId_range*BK
# load from globals into locals
i = UOp.range(dtypes.int, nbReadsB, 1)
i = UOp.range(nbReadsB, 1)
index_x = BN * blockIdx_x + rBIdx
index_y = rBIdy + i * strideReadB + kId
Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)
i = UOp.range(dtypes.int, nbReadsA, 2)
i = UOp.range(nbReadsA, 2)
index_x = rAIdx + kId
index_y = BM * blockIdx_y + rAIdy + i * strideReadA
As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
barrier = UOp.barrier(As_store, Bs_store)
k = UOp.range(dtypes.int, BK, 3)
k = UOp.range(BK, 3)
# load from locals into registers
iterWave = UOp.range(dtypes.int, nbIterWaveN, 4)
i = UOp.range(dtypes.int, TN, 5)
iterWave = UOp.range(nbIterWaveN, 4)
i = UOp.range(TN, 5)
index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(barrier), iterWave, i)
iterWave = UOp.range(dtypes.int, nbIterWaveM, 6)
i = UOp.range(dtypes.int, TM, 7)
iterWave = UOp.range(nbIterWaveM, 6)
i = UOp.range(TM, 7)
index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(barrier), iterWave, i)
# do the GEMM math
iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 8)
yt = UOp.range(dtypes.int, TM, 9)
iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 10)
xt = UOp.range(dtypes.int, TN, 12)
iterWaveM = UOp.range(nbIterWaveM, 8)
yt = UOp.range(TM, 9)
iterWaveN = UOp.range(nbIterWaveN, 10)
xt = UOp.range(TN, 12)
x = iterWaveN * TN + xt
y = iterWaveM * TM + yt
c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
@@ -310,10 +310,10 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
iterWaveM, iterWaveN, yt, xt, k, kId_range)
# store c_regs into c
iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 1000)
yt = UOp.range(dtypes.int, TM, 1001)
iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 1002)
xt = UOp.range(dtypes.int, TN, 1003)
iterWaveM = UOp.range(nbIterWaveM, 1000)
yt = UOp.range(TM, 1001)
iterWaveN = UOp.range(nbIterWaveN, 1002)
xt = UOp.range(TN, 1003)
xOut = blockIdx_x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave
yOut = blockIdx_y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave
indexC = N * (yOut + yt) + xOut + xt
tinygrad_repo/extra/gemm/intel_xmx.py
+5 -5
View File
@@ -1,6 +1,6 @@
#!/usr/bin/env python3
import numpy as np
from tinygrad.runtime.ops_gpu import CLProgram, CLCompiler
from tinygrad.runtime.ops_cl import CLProgram, CLCompiler
from tinygrad import Device, dtypes
from tinygrad.device import Buffer
from hexdump import hexdump
@@ -11,7 +11,7 @@ from hexdump import hexdump
# https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_split_matrix_multiply_accumulate.html
# https://hc34.hotchips.org/assets/program/conference/day1/GPU%20HPC/Intel_s%20Ponte%20Vecchio%20GPU%20-%20Architecture%20Systems%20and%20Software%20FINAL.pdf
device = Device["GPU"]
device = Device["CL"]
# NOTE: only the subgroup type 8 ones work
prog = CLProgram(device, "test", CLCompiler(device, "test").compile(f"""
@@ -26,9 +26,9 @@ __kernel void test(__global float* data0, const __global int* data1, const __glo
"""))
#with open("/tmp/test.elf", "wb") as f: f.write(prog.lib)
a = Buffer("GPU", 8, dtypes.float32).allocate()
b = Buffer("GPU", 0x10, dtypes.float16).allocate()
c = Buffer("GPU", 8*0x10, dtypes.float16).allocate()
a = Buffer("CL", 8, dtypes.float32).allocate()
b = Buffer("CL", 0x10, dtypes.float16).allocate()
c = Buffer("CL", 8*0x10, dtypes.float16).allocate()
row = np.array([1,2,3,4,5,6,7,8,1,2,3,4,5,6,7,8], np.float16)
mat = np.random.random((8, 0x10)).astype(np.float16)
tinygrad_repo/extra/gemm/max_matmul.py
+2 -2
View File
@@ -56,7 +56,7 @@ def randoms():
def ast_to_cuda_prog(compiler, ast, opts):
k = Kernel(ast)
k.apply_opts(opts)
p = get_program(k.get_optimized_ast(), k.opts)
p = get_program(k.ast, k.opts, k.applied_opts)
return CUDAProgram(device, p.function_name, compiler.compile(p.src))
if __name__ == "__main__":
@@ -75,7 +75,7 @@ if __name__ == "__main__":
if GEMM_VARIATION == "max" and (M%64)==0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.float and DTYPE_ACC == dtypes.float:
print("Using CUDA and triton-generated kernel")
# See nv_triton_gemm.annotated.ptx for PTX code which was generated from `PYTHONPATH=. DEBUG=6 CUDA=1 PTX=1 python3 extra/gemm/triton_nv_matmul.py`
# See nv_triton_gemm.annotated.ptx for PTX code which was generated from `PYTHONPATH=. DEBUG=6 CUDA=1 CUDA_PTX=1 python3 extra/gemm/triton_nv_matmul.py`
# this kernel with M=N=K=4096 does 162TFLOPS, vs torch at 144TFLOPS and BEAM=8 tinygrad at 138TFLOPS. theo max is 165TFLOPS.
# WMMA element size is (M, N, K) = (16, 8, 16)
tinygrad_repo/extra/gemm/simple_matmul.py
+1 -1
View File
@@ -2,7 +2,7 @@ import numpy as np
from tinygrad import dtypes, Tensor
from tinygrad.helpers import getenv, get_single_element
from tinygrad.dtype import _to_np_dtype
from tinygrad.codegen.opt.kernel import OptOps
from tinygrad.codegen.opt import OptOps
from tinygrad.engine.realize import lower_schedule
dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
tinygrad_repo/extra/gemm/tinygrad_nv_matmul.py
+1 -1
View File
@@ -29,7 +29,7 @@ if __name__ == "__main__":
Opt(op=OptOps.LOCAL, axis=0, amt=2),
]
k.apply_opts(opts)
prg = get_program(k.get_optimized_ast(), k.opts)
prg = get_program(k.ast, k.opts, k.applied_opts)
new_src = prg.src
# can mod source here
prg = replace(prg, src=new_src)
tinygrad_repo/extra/gemm/triton_nv_matmul.py
+1 -1
View File
@@ -43,7 +43,7 @@ def matmul_kernel(c_ptr, a_ptr, b_ptr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N:
c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
tl.store(c_ptrs, c)
# CUDA=1 PTX=1 python3 extra/gemm/triton_nv_matmul.py
# CUDA=1 CUDA_PTX=1 python3 extra/gemm/triton_nv_matmul.py
if __name__ == "__main__":
BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 128, 64
M, N, K = 4096, 4096, 4096
tinygrad_repo/extra/mcts_search.py
+1 -1
View File
@@ -88,7 +88,7 @@ def mcts_search(lin:Kernel, rawbufs:List[Buffer], amt:int) -> Kernel:
return ret
rawbufs = _ensure_buffer_alloc(rawbufs)
var_vals = {k:(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
var_vals = {k.expr:(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
dev = Device[lin.opts.device]
root = MCTSNode(lin)
tinygrad_repo/extra/models/inception.py
+4 -2
View File
@@ -270,8 +270,10 @@ class FidInceptionV3:
self.Mixed_7b = inception.Mixed_7b
self.Mixed_7c = inception.Mixed_7c
def load_from_pretrained(self):
state_dict = torch_load(str(fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")))
def load_from_pretrained(self, path=None):
if path is None:
path = fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")
state_dict = torch_load(str(path))
for k,v in state_dict.items():
if k.endswith(".num_batches_tracked"):
state_dict[k] = v.reshape(1)
tinygrad_repo/extra/models/llama.py
+2 -5
View File
@@ -249,8 +249,5 @@ def convert_from_gguf(weights:dict[str, Tensor], n_layers:int):
return sd
def fix_bf16(weights:dict[Any, Tensor]):
if getenv("SUPPORT_BF16", 1):
# TODO: without casting to float16, 70B llama OOM on tinybox.
return {k:v.cast(dtypes.float32).cast(dtypes.float16) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}
# TODO: check if device supports bf16
return {k:v.llvm_bf16_cast(dtypes.half).to(v.device) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}
# TODO: without casting to float16, 70B llama OOM on tinybox.
return {k:v.cast(dtypes.float32).cast(dtypes.float16) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}
tinygrad_repo/extra/nv_gpu_driver/nv_ioctl.py
+1 -1
View File
@@ -272,4 +272,4 @@ def compare_launch_state(states, good_states):
return True, "PASS"
# IOCTL=1 PTX=1 CUDA=1 python3 test/test_ops.py TestOps.test_tiny_add
# IOCTL=1 CUDA=1 CUDA_PTX=1 python3 test/test_ops.py TestOps.test_tiny_add
tinygrad_repo/extra/optimization/generate_dataset.sh
+1 -1
View File
@@ -7,7 +7,7 @@ rm $LOGOPS
test/external/process_replay/reset.py
CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/test_winograd.py test/models/test_real_world.py --durations=20
GPU=1 python3 -m pytest test/test_tiny.py
CL=1 python3 -m pytest test/test_tiny.py
# extract, sort and uniq
extra/optimization/extract_dataset.py
tinygrad_repo/extra/optimization/helpers.py
+2 -2
View File
@@ -1,6 +1,6 @@
# stuff needed to unpack a kernel
from tinygrad import Variable
from tinygrad.codegen.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.dtype import dtypes, PtrDType
from tinygrad.shape.shapetracker import ShapeTracker
@@ -115,7 +115,7 @@ def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_
assert dev.compiler is not None
rawbufs = _ensure_buffer_alloc(rawbufs)
var_vals: dict[Variable, int] = {k:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
p = get_program(lin.get_optimized_ast(), lin.opts)
tms = _time_program(p, dev.compiler.compile(p.src), var_vals, rawbufs,
max_global_size=max_global_size if allow_test_size else None, clear_l2=clear_l2, cnt=cnt, name=to_function_name(lin.name))
tinygrad_repo/extra/optimization/test_beam_search.py
+9 -11
View File
@@ -16,9 +16,9 @@ class TestBeamSearch(unittest.TestCase):
BEAM.value = self.old_beam
def test_variable_ast_beam(self):
with Context(IGNORE_OOB=1):
a = rand(3, 3).reshape((Variable("a", 1, 10).bind(3), 3))
a = (a+1).realize()
vi = Variable("a", 1, 10).bind(3)
a = rand(10, 3)[:vi]
a = (a+1).realize()
def test_big_prime_number(self):
a = rand(367, 367)
@@ -42,18 +42,16 @@ class TestBeamSearch(unittest.TestCase):
def test_variable_big_prime_number(self):
v = Variable("v", 1, 400).bind(367)
a = rand(367, 367)
b = rand(367, 367)
with Context(IGNORE_OOB=1):
c = (a.reshape(367, v) @ b.reshape(v, 367)).realize()
np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), atol=1e-4, rtol=1e-4)
a = rand(367, 400)
b = rand(400, 367)
c = (a[:, :v] @ b[:v, :]).realize()
np.testing.assert_allclose(c.numpy(), a[:, :367].numpy() @ b[:367, :].numpy(), atol=1e-4, rtol=1e-4)
def test_variable_shrink_prime_number(self):
v = Variable("v", 1, 400).bind(367)
a = rand(400, 367)
with Context(IGNORE_OOB=1):
b = (a.shrink(((0,v), None))+1).reshape(367,367).realize()
np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4)
b = (a.shrink(((0,v), None))+1)[:367,:367].realize()
np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4)
def test_no_mutate_rawbuffers(self):
a = rand(3, 3).realize()
tinygrad_repo/extra/qcom_gpu_driver/qcom_opencl_interop.py
+2 -2
View File
@@ -1,6 +1,6 @@
import ctypes, array
from hexdump import hexdump
from tinygrad.runtime.ops_gpu import GPUDevice
from tinygrad.runtime.ops_cl import CLDevice
from tinygrad.helpers import getenv, to_mv, mv_address
from tinygrad.dtype import dtypes
from tinygrad import Tensor, TinyJit
@@ -8,7 +8,7 @@ from tinygrad.runtime.autogen import opencl as cl
if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl # noqa: F401 # pylint: disable=unused-import
# create raw opencl buffer.
gdev = GPUDevice()
gdev = CLDevice()
cl_buf = cl.clCreateBuffer(gdev.context, cl.CL_MEM_READ_WRITE, 0x100, None, status := ctypes.c_int32())
assert status.value == 0
tinygrad_repo/extra/remu/src/thread.rs
+11 -1
View File
@@ -673,6 +673,7 @@ impl<'a> Thread<'a> {
39 => f32::log2(s0),
42 => 1.0 / s0,
43 => 1.0 / s0,
46 => 1.0 / f32::sqrt(s0),
51 => f32::sqrt(s0),
_ => todo_instr!(instruction)?,
}
@@ -1246,7 +1247,7 @@ impl<'a> Thread<'a> {
}
let ret = match op {
257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 531 | 537 | 540 | 551 | 567 | 796 => {
257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 430 | 531 | 537 | 540 | 551 | 567 | 796 => {
let s0 = f32::from_bits(s0).negate(0, neg).absolute(0, abs);
let s1 = f32::from_bits(s1).negate(1, neg).absolute(1, abs);
let s2 = f32::from_bits(s2).negate(2, neg).absolute(2, abs);
@@ -1258,6 +1259,7 @@ impl<'a> Thread<'a> {
272 => f32::max(s0, s1),
299 => f32::mul_add(s0, s1, f32::from_bits(self.vec_reg[vdst])),
426 => s0.recip(),
430 => 1.0 / f32::sqrt(s0),
531 => f32::mul_add(s0, s1, s2),
537 => f32::min(f32::min(s0, s1), s2),
540 => f32::max(f32::max(s0, s1), s2),
@@ -2625,6 +2627,14 @@ mod test_vop1 {
assert_eq!(thread.vec_reg[3], 1071644672);
}
#[test]
fn test_v_rsq_f32() {
let mut thread = _helper_test_thread();
thread.vec_reg[0] = f32::to_bits(4.0);
r(&vec![0x7E005D00, END_PRG], &mut thread);
assert_eq!(f32::from_bits(thread.vec_reg[0]), 0.5);
}
#[test]
fn test_v_frexp_exp_i32_f64() {
[(3573412790272.0, 42), (69.0, 7), (2.0, 2), (f64::NEG_INFINITY, 0)]
tinygrad_repo/extra/replay_pkl.py
+1 -1
View File
@@ -58,7 +58,7 @@ if __name__ == "__main__":
GlobalCounters.kernel_count -= 1
if not getenv("NOOPT"): k.apply_opts(hand_coded_optimizations(k))
p2 = get_program(k.get_optimized_ast(), k.opts)
p2 = get_program(k.ast, k.opts, k.applied_opts)
new_ei = replace(ei, prg=CompiledRunner(p2))
new_ei.run()
new_jit.append(new_ei)
tinygrad_repo/extra/sqtt/README.md
+1 -1
View File
@@ -4,7 +4,7 @@
Only supported on 7900XTX, requires either AM (`rmmod amdgpu`) or disabling power gating on AMD (`ppfeaturemask=0xffff3fff`, don't forget to rebuild initramfs)
SQTT is implemented on top of normal tinygrad PROFILE=1, `PROFILE=1 SQTT=1` to get profile pickle with sqtt data embedded in it.
SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.
`SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.
tinygrad_repo/extra/test_hcopt.py
+40
View File
@@ -0,0 +1,40 @@
import time
from extra.optimization.helpers import load_worlds, ast_str_to_ast
from tinygrad import Device
from tinygrad.codegen.lowerer import pm_lowerer, get_index
from tinygrad.uop.ops import graph_rewrite
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.codegen.opt.postrange import Scheduler
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
from tinygrad.helpers import getenv
if __name__ == "__main__":
renderer = Device.default.renderer
ast_strs = load_worlds()
if (n:=getenv("N", -1)) != -1: ast_strs = ast_strs[n:n+1]
good = 0
for i, ast_str in enumerate(ast_strs):
ast = ast_str_to_ast(ast_str)
st = time.perf_counter()
lin = Kernel(ast, renderer)
opt1 = hand_coded_optimizations(lin)
et_lin = time.perf_counter() - st
lowered = graph_rewrite(ast, pm_lowerer, ctx=get_index(ast), bottom_up=True)
st = time.perf_counter()
sch = Scheduler(lowered, renderer)
sch.convert_loop_to_global()
sch.simplify_merge_adjacent()
opt2 = hand_coded_optimizations(sch)
et_sch = time.perf_counter() - st
if opt1 != opt2:
print(f"******* {i:6d}")
print("Kernel: ", lin.colored_shape(), "->", lin.apply_opts(opt1).colored_shape())
print("Scheduler: ", sch.colored_shape(), "->", sch.apply_opts(opt2).colored_shape())
print(opt1)
print(opt2)
else:
good += 1
print(f"******* {i:6d} MATCH {good/(i+1)*100:.2f}% -- {et_lin/et_sch:4.2f}x speedup")
tinygrad_repo/extra/test_pyrender.py
+20
View File
@@ -0,0 +1,20 @@
from extra.optimization.helpers import load_worlds, ast_str_to_ast
from tinygrad.helpers import tqdm
from tinygrad.uop.ops import pyrender, UOp, Ops
from tinygrad import dtypes
from tinygrad.shape.shapetracker import ShapeTracker, View
inf, nan = float('inf'), float('nan')
if __name__ == "__main__":
ast_strs = load_worlds()
for i, ast_str in enumerate(tqdm(ast_strs)):
good_ast = ast_str_to_ast(ast_str)
code = '\n'.join(pyrender(good_ast))
print("\n***************\n\n"+code)
exec(code)
if str(good_ast) != str(ast):
print(code)
print("MISMATCH")
print(good_ast)
print(ast)
break
tinygrad_repo/extra/thneed.py
+5 -5
View File
@@ -4,13 +4,13 @@ import struct
import json
import traceback
import numpy as np
from tinygrad.runtime.ops_gpu import CLProgram, compile_gpu
from tinygrad.runtime.ops_cl import CLProgram, compile_gpu
from tinygrad.device import Device
from tinygrad.helpers import DEBUG, getenv
from collections import defaultdict
import pyopencl as cl
from tinygrad.runtime.ops_gpu import OSX_TIMING_RATIO
CL = Device["GPU"]
from tinygrad.runtime.ops_cl import OSX_TIMING_RATIO
CL = Device["CL"]
DEBUGCL = getenv("DEBUGCL", 0)
FLOAT16 = getenv("FLOAT16", 0)
@@ -110,7 +110,7 @@ class Thneed:
prgs = {}
for o in jdat['binaries']:
nptr = ptr + o['length']
prgs[o['name']] = CLProgram(Device["GPU"], o['name'], weights[ptr:nptr])
prgs[o['name']] = CLProgram(Device["CL"], o['name'], weights[ptr:nptr])
ptr = nptr
# populate the cl_cache
@@ -267,7 +267,7 @@ class Thneed:
for prg, args in self.cl_cache:
events.append(prg.clprg(CL.queue, *args))
mt = time.monotonic()
Device["GPU"].synchronize()
Device["CL"].synchronize()
et = time.monotonic() - st
print(f"submit in {(mt-st)*1000.0:.2f} ms, total runtime is {et*1000.0:.2f} ms")
tinygrad_repo/extra/to_movement_ops.py
+3 -3
View File
@@ -2,7 +2,6 @@ import itertools
from enum import Enum, auto
from collections import defaultdict
from typing import List, Tuple, DefaultDict
from extra.optimization.helpers import load_worlds, ast_str_to_ast
from tinygrad.helpers import prod, tqdm
from tinygrad.uop.ops import UOp, Ops
from tinygrad.shape.shapetracker import ShapeTracker
@@ -36,7 +35,7 @@ def to_movement_ops(st: ShapeTracker) -> List[Tuple[MovementOps, Tuple]]:
to_apply:List[Tuple[MovementOps, Tuple]] = []
for i, v in enumerate(st.views):
real_shape = tuple(y-x for x,y in v.mask) if v.mask else v.shape
offset = v.offset + sum(st*(s-1) for s,st in zip(real_shape, v.strides) if st<0)
offset = (v.offset or 0) + sum(st*(s-1) for s,st in zip(real_shape, v.strides) if st<0)
real_offset = offset + (sum(x*st for (x,_),st in zip(v.mask, v.strides)) if v.mask else 0)
real_real_shape = [s for s,st in zip(real_shape, v.strides) if st]
strides: List[int] = [abs(st) if isinstance(st,int) else st for st in v.strides if st]
@@ -121,7 +120,7 @@ def st_equivalent(st1: ShapeTracker, st2: ShapeTracker):
if i > 1000:
print("WARNING: did not search all possible combinations")
break
var_vals = {k:v for k,v in zip(vs, ranges)}
var_vals = {k.expr:v for k,v in zip(vs, ranges)}
r1 = sym_infer(idx1, var_vals) if sym_infer(valid1, var_vals) else 0
r2 = sym_infer(idx2, var_vals) if sym_infer(valid2, var_vals) else 0
if r1 != r2: return False
@@ -147,6 +146,7 @@ def test_rebuild_bufferop_st(ast:UOp):
for src in ast.src: test_rebuild_bufferop_st(src)
if __name__ == "__main__":
from extra.optimization.helpers import load_worlds, ast_str_to_ast
ast_strs = load_worlds(False, False, True)[:2000]
for ast_str in tqdm(ast_strs):
test_rebuild_bufferop_st(ast_str_to_ast(ast_str))
tinygrad_repo/extra/torch_backend/backend.py
+17 -11
View File
@@ -177,22 +177,28 @@ def cached_to_movement_ops(shape, st) -> list:
from tinygrad.shape.shapetracker import ShapeTracker, View
from extra.to_movement_ops import to_movement_ops, apply_mop, MovementOps
@wrap_view_op
def _as_strided(tensor:Tensor, size, stride, storage_offset=None):
# multiple as_strided do not compound
base = canonical_base(tensor)
# TODO: this is heavyweight
st = ShapeTracker(base.uop.st.views + (View.create(tuple(size), tuple(stride), storage_offset),))
ret = base
if TORCH_DEBUG >= 1: print("**** as_strided", tensor.shape, size, stride, st)
if prod(size) == 1: return ret.flatten()[storage_offset].reshape(size)
for mo in cached_to_movement_ops(tuple(base.shape), st): ret = apply_mop(ret, mo)
return ret
@torch.library.impl("aten::as_strided", "privateuseone")
def as_strided(tensor:torch.Tensor, size, stride, storage_offset=None):
storage_offset = storage_offset or tensor.storage_offset()
@wrap_view_op
def _as_strided(tensor:Tensor, size, stride, storage_offset=None):
# multiple as_strided do not compound
base = canonical_base(tensor)
# TODO: this is heavyweight
st = ShapeTracker(base.uop.st.views + (View.create(tuple(size), tuple(stride), storage_offset),))
ret = base
if TORCH_DEBUG >= 1: print("**** as_strided", tensor.shape, size, stride, st)
if prod(size) == 1: return ret.flatten()[storage_offset].reshape(size)
for mo in cached_to_movement_ops(tuple(base.shape), st): ret = apply_mop(ret, mo)
return ret
return _as_strided(tensor, size, stride, storage_offset)
@torch.library.impl("aten::_reshape_alias", "privateuseone")
def _reshape_alias(tensor:torch.Tensor, size, stride):
return _as_strided(tensor, size, stride)
@torch.library.impl("aten::empty_strided", "privateuseone")
def empty_strided(size, stride, dtype, layout=None, device=None, pin_memory=False):
if TORCH_DEBUG: print(f"empty_strided {size=} {stride=} {dtype=} {layout=} {device=} {pin_memory=}")