from_blob for cuda (#9223)

* from_blob for cuda

* maybe docs?

* minor docs

* example

* waiting 9224

---------

Co-authored-by: George Hotz <[email protected]>

docs/runtime.md

@@ -13,3 +13,48 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | [CPU (C Code)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang compiler | `clang` compiler in system `PATH` |
 | [LLVM (LLVM IR)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_llvm.py) | Runs on CPU using the LLVM compiler infrastructure | llvm libraries installed and findable |
 | [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | Dawn library installed and findable. Download binaries [here](https://github.com/wpmed92/pydawn/releases/tag/v0.1.6). |
+
+## Interoperability
+
+tinygrad provides interoperability with OpenCL and PyTorch, allowing efficient tensor data sharing between frameworks through the `Tensor.from_blob` API. This enables zero-copy operations by working directly with external memory pointers.
+
+**Important**: When using external memory pointers with tinygrad tensors, you must ensure these pointers remain valid throughout the entire lifetime of the tinygrad tensor to prevent memory corruption.
+
+### `CUDA` PyTorch Interoperability
+
+You can seamlessly work with CUDA tensors between PyTorch and tinygrad without data copying:
+```python
+from tinygrad.dtype import _from_torch_dtype
+tensor1 = torch.tensor([1.0, 2.0, 3.0], device=torch.device("cuda"))
+tiny_tensor1 = Tensor.from_blob(tensor1.data_ptr(), tensor1.shape, dtype=_from_torch_dtype(tensor1.dtype), device='CUDA')
+```
+
+### `QCOM` OpenCL Interoperability
+
+tinygrad supports OpenCL interoperability on `QCOM` backend.
+
+Buffer interop allows direct access to OpenCL memory buffers:
+```python
+# create raw opencl buffer.
+cl_buf = cl.clCreateBuffer(cl_context, cl.CL_MEM_READ_WRITE, 0x100, None, status := ctypes.c_int32())
+
+# extract pointers
+cl_buf_desc_ptr = to_mv(ctypes.addressof(cl_buf), 8).cast('Q')[0]
+rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw gpu pointer.
+
+# create tiny tensor
+tiny = Tensor.from_blob(rawbuf_ptr, (8, 8), dtype=dtypes.int, device='QCOM')
+```
+
+And the same for the images:
+```python
+# create cl image.
+cl_img = cl.clCreateImage2D(cl_context, cl.CL_MEM_READ_WRITE, cl.cl_image_format(cl.CL_RGBA, cl.CL_FLOAT), w, h, 0, None, status := ctypes.c_int32())
+
+# extract pointers
+cl_buf_desc_ptr = to_mv(ctypes.addressof(cl_img), 8).cast('Q')[0]
+rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw gpu pointer.
+
+# create tiny tensor
+tiny = Tensor.from_blob(rawbuf_ptr, (h*w*4,), dtype=dtypes.imagef((h,w)), device='QCOM')
+```

examples/torch_cuda_kernel.py

@@ -12,28 +12,23 @@ except ImportError:
 from tinygrad import Tensor, TinyJit
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import get_single_element, Context
+from tinygrad.dtype import _from_torch_dtype
 
-@functools.lru_cache(None)
-def generate_tinygrad_cuda(shape:tuple[int, ...], beam=0) -> CompiledRunner:
-  print("generating code for", shape)
-  tg_data = Tensor.zeros(*shape, device="CUDA").contiguous().realize()
-  with Context(BEAM=beam):
-    @TinyJit
-    def f(x): return x[:, :, 0] * 0.2989 + x[:, :, 1] * 0.5870 + x[:, :, 2] * 0.1140
-    for _ in range(3): f(tg_data)
-    fxn = get_single_element(f.captured.jit_cache).prg
-  print("generated with", fxn.p.global_size, fxn.p.local_size)
-  return fxn
+@TinyJit
+def f(tg_out, tg_data): return tg_out.assign(tg_data[:, :, 0] * 0.2989 + tg_data[:, :, 1] * 0.5870 + tg_data[:, :, 2] * 0.1140).realize()
 
 def custom_kernel(data: torch.Tensor) -> torch.Tensor:
   assert data.dtype == torch.float32
+  tg_data = Tensor.from_blob(data.data_ptr(), data.shape, dtype=_from_torch_dtype(data.dtype), device='CUDA')
+
   out = torch.empty((data.shape[0], data.shape[1]), dtype=data.dtype, device=data.device)
-  fxn = generate_tinygrad_cuda(tuple(data.shape), beam=2)
-  fxn._prg(out.data_ptr(), data.data_ptr(), global_size=fxn.p.global_size, local_size=fxn.p.local_size)
+  tg_out = Tensor.from_blob(out.data_ptr(), out.shape, dtype=_from_torch_dtype(out.dtype), device='CUDA')
+
+  with Context(BEAM=2): f(tg_out, tg_data)
   return out
 
 if __name__ == "__main__":
   for i in range(3):
-    out = custom_kernel(inp:=torch.rand(16, 16, 3))
+    out = custom_kernel(inp:=torch.rand(16, 16, 3, device=torch.device("cuda")))
     torch.cuda.synchronize()
     assert torch.allclose(out, inp[:, :, 0] * 0.2989 + inp[:, :, 1] * 0.5870 + inp[:, :, 2] * 0.1140)

tinygrad/runtime/ops_cuda.py

@@ -68,6 +68,7 @@ class CUDAAllocator(LRUAllocator):
     super().__init__()
   def _alloc(self, size, options:BufferSpec):
     check(cuda.cuCtxSetCurrent(self.dev.context))
+    if options.external_ptr: return cuda.CUdeviceptr_v2(options.external_ptr)
     if options.host: return init_c_var(ctypes.c_void_p(), lambda x: check(cuda.cuMemHostAlloc(ctypes.byref(x), size, 0x01)))
     return init_c_var(cuda.CUdeviceptr(), lambda x: check(cuda.cuMemAlloc_v2(ctypes.byref(x), size)))
   def _free(self, opaque, options:BufferSpec):