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

dynamic tc function render (#3387)

Browse Source 
hip cant be done right now
This commit is contained in:
qazal
2024-02-15 11:19:46 +01:00
committed by GitHub
parent 3c4c8c9e16
commit 9d4d63fcfc
1 changed files with 9 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/renderer/cstyle.py
+9 -4
View File
@@ -191,10 +191,7 @@ class OpenCLLanguage(CStyleLanguage):
OpenCLRenderer = functools.partial(uops_to_cstyle, OpenCLLanguage())
class MetalLanguage(CStyleLanguage):
kernel_prefix = """#include <metal_stdlib>\nusing namespace metal;\ntemplate<typename T, typename S, typename U> U __metal_wmma(T m, T n, U o) {
S a,b,c; a.thread_elements()[0] = m.x; a.thread_elements()[1] = m.y; b.thread_elements()[0] = n.x; b.thread_elements()[1] = n.y;
c.thread_elements()[0] = o.x; c.thread_elements()[1] = o.y; simdgroup_multiply_accumulate(c, a, b, c);
return U(c.thread_elements()[0], c.thread_elements()[1]);\n}\nkernel """
kernel_prefix = "kernel "
buffer_prefix = "device "
smem_prefix = "threadgroup "
arg_int_prefix = "constant int&"
@@ -205,6 +202,14 @@ class MetalLanguage(CStyleLanguage):
extra_args = ['uint3 gid [[threadgroup_position_in_grid]]', 'uint3 lid [[thread_position_in_threadgroup]]']
def render_cast(self, x: List[str], var_dtype: DType, bitcast=False) -> str:
return f"as_type<{var_dtype.name}>({x[0]})" if bitcast else super().render_cast(x, var_dtype)
def render_kernel(self, function_name, kernel, bufs, local_size, uops, prefix=None):
prefix = ["#include <metal_stdlib>","using namespace metal;"]
if any(uop.uop == UOps.WMMA for uop in uops): prefix.append("""template<typename T, typename S, typename U> U __metal_wmma(T m, T n, U o) {
S a,b,c; a.thread_elements()[0] = m.x; a.thread_elements()[1] = m.y; b.thread_elements()[0] = n.x; b.thread_elements()[1] = n.y;
c.thread_elements()[0] = o.x; c.thread_elements()[1] = o.y; simdgroup_multiply_accumulate(c, a, b, c);
return U(c.thread_elements()[0], c.thread_elements()[1]);\n}""")
return super().render_kernel(function_name, kernel, bufs, local_size, uops, prefix)
MetalRenderer = functools.partial(uops_to_cstyle, MetalLanguage())
code_for_op_half = {