geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

inf hack that works on chrome and wgpu (#2712)

Browse Source 
Co-authored-by: George Hotz <[email protected]>
This commit is contained in:
Ahmed Harmouche
2023-12-16 11:49:37 -05:00
committed by GitHub
co-authored by George Hotz
parent 877c78b4ce
commit a7264dcb2b
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/renderer/cstyle.py
+2 -2
View File
@@ -345,13 +345,13 @@ class WGSLLanguage(CStyleLanguage):
def render_const(self, x:Union[float,int], var_dtype) -> str:
if math.isnan(x): return "nan()"
elif math.isinf(x): return ("-" if x < 0 else "") + "(1.0/0.0)"
elif math.isinf(x): return ("-" if x < 0 else "") + "inf(1.0)"
return f"({super().render_const(x, var_dtype)})"
def render_kernel(self, function_name:str, kernel:List[str], bufs:List[Tuple[str,DType]], local_size:List[int], prekernel:List[str]) -> str:
local_size = local_size[::-1] if local_size else [1]
bind_it = iter(range(len(bufs)))
prg = "fn nan() -> f32 { let bits = 0xffffffffu; return bitcast<f32>(bits); }\n"
prg = "fn nan() -> f32 { let bits = 0xffffffffu; return bitcast<f32>(bits); }\nfn inf(a: f32) -> f32 { return a/0.0; }\n"
prg += "\n".join(prekernel+[f"@group(0) @binding({next(bind_it)}) {'var<storage,read_write>' if isinstance(dtype, PtrDType) else 'var<uniform>'} {name}: {f'array<{self.type_map[dtype]}>' if isinstance(dtype, PtrDType) else 'i32'};" for name,dtype in bufs]) # noqa: E501
prg += f"\n@compute @workgroup_size({','.join([str(x) for x in local_size])}) fn {function_name}(@builtin(workgroup_id) gindex: vec3<u32>, @builtin(local_invocation_id) lindex: vec3<u32>) {{\n" + "\n".join(kernel) + "\n}" # noqa: E501
return prg