From a7264dcb2bc9d2a5bcf598aeb9a65e304945799a Mon Sep 17 00:00:00 2001 From: Ahmed Harmouche Date: Sat, 16 Dec 2023 17:49:37 +0100 Subject: [PATCH] inf hack that works on chrome and wgpu (#2712) Co-authored-by: George Hotz <72895+geohot@users.noreply.github.com> --- tinygrad/renderer/cstyle.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 6bcb979f85..4caa820c3b 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -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(bits); }\n" + prg = "fn nan() -> f32 { let bits = 0xffffffffu; return bitcast(bits); }\nfn inf(a: f32) -> f32 { return a/0.0; }\n" prg += "\n".join(prekernel+[f"@group(0) @binding({next(bind_it)}) {'var' if isinstance(dtype, PtrDType) else 'var'} {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, @builtin(local_invocation_id) lindex: vec3) {{\n" + "\n".join(kernel) + "\n}" # noqa: E501 return prg