deepcrayon
/
tinygrab
1
0
Fork 0
Code

moved global_max and local_max to LinearizerOptions also added assert for max bufs (#1446)

pull/1449/head
Diogo 2023-08-05 21:23:18 -04:00 committed by GitHub
parent 7b8d06c9f1
commit e2af95c2f8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tinygrad/renderer/wgsl.py
View File

@ -11,8 +11,6 @@ class WGSLLanguage(CStyleLanguage):
gid = [f"i32(gindex.{'xyz'[x]})" for x in range(3)]
lid = [f"i32(lindex.{'xyz'[x]})" for x in range(3)]
size_prefix = "let"
global_max = [65535, 65535, 65535]
local_max = [256, 256, 64]
barrier="workgroupBarrier();"
generic_var_prefix = "var "
external_local_bufs = True

3
tinygrad/runtime/ops_webgpu.py
View File

@ -14,6 +14,7 @@ device = get_default_device()
class WebGPUProgram:
def __init__(self, name: str, prg: str): self.name,self.prg = name,device.create_shader_module(code=prg)
def __call__(self, global_size, local_size, *bufs, wait=False):
assert len(bufs) <= 8, "WEBGPU only supports 8 buffers"
binding_layouts = [{"binding": i, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}} for i in range(len(bufs))]
bindings = [{"binding": i, "resource": {"buffer": x._buf, "offset": 0, "size": x._buf.size}} for i, x in enumerate(bufs)]
bind_group_layout = device.create_bind_group_layout(entries=binding_layouts)
@ -36,4 +37,4 @@ class RawWebGPUBuffer(RawBufferCopyIn):
def toCPU(self) -> np.ndarray: return np.frombuffer(device.queue.read_buffer(self._buf, 0), dtype=np.dtype(self.dtype.np, metadata={"backing": self})) # type: ignore
renderer = functools.partial(uops_to_cstyle, WGSLLanguage())
WebGpuBuffer = Compiled(RawWebGPUBuffer, LinearizerOptions(supports_float4=False), renderer, WebGPUProgram)
WebGpuBuffer = Compiled(RawWebGPUBuffer, LinearizerOptions(supports_float4=False, local_max=[256, 256, 64], global_max=[65535, 65535, 65535]), renderer, WebGPUProgram)