github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-07 03:00:26 -04:00
Code Issues Packages Projects Releases Wiki Activity

webgl backend in extra (#3041)

Browse Source 
* WebGL WIP

* 84% of ops passing test

* tests passing 100%

* Cleanup, refactor

* Shave off some lines

* Work on dtypes

* TestOps at 100% again

* Efficient net shaders compile in browser webgl2

* Compile all efficientnet shaders in browser

* Create empty textures for tensor buffers

* Run program. Up next weight loading

* Exported WebGL model working

* Add tests, refactor

* Explicit cast alu for GLSL

* Fix CI tests

* WebGL efficientnet demo

* Compile and run yolov8 in browser

* Fix imports

* Simplify yolo compile

* Fix bool*bool and cast cmplt to float

* More tests

* Do std tests pass on CI?

* Skip std tests on CI

* Remove explicit_cast_alu hack, and solve it in code_for_op

* Move to new dtype-less alloc api

* Remove local size hack: optimize local_size only if device has local

* Remove glsl.py, and move content to cstyle

* dont_use_locals in opts

* Fix dtype tests

* type_map in CStyleLanguage

* Make core changes smaller, cleaner, refactor export_model and demo

* Skip pad_slice

* Simplify: render_const, render_conditional

* solve bool alu for other binops, cleaner ops_webgl

* Fix noopt hack

* Remove some skipIfs

* WebGL image hack

* type_names is a better name

* global_max

* Fix dtype import

* Fix type_names -> type_map

* Fix lint

* Remove webgpu, back to 5k lines (#3040)

* remove webgpu

* max 5000 lines

* revert those to master

* retain that cstyle

---------

Co-authored-by: Ahmed Harmouche <ahmedharmouche92@gmail.com>
This commit is contained in:
George Hotz
2024-01-08 09:29:13 -08:00
committed by GitHub
parent 8cbcd1b342
commit c5a941d466
9 changed files with 596 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
.github/workflows/test.yml vendored
View File

@@ -334,6 +334,38 @@ jobs:
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/rocm/hip/lib
MOCKHIP=1 HIP=1 python -m pytest -s test/test_hip_rdna3.py
# testwebgl:
# name: WebGL Tests
# runs-on: ubuntu-latest
# timeout-minutes: 20
#
# steps:
# - name: Checkout Code
# uses: actions/checkout@v3
# - name: Set up Python 3.11
# uses: actions/setup-python@v4
# with:
# python-version: 3.11
# - name: Cache python packages
# uses: actions/cache@v3
# with:
# path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
# key: webgl-testing-packages-${{ hashFiles('**/setup.py') }}
# - name: Install Dependencies
# run: pip install -e '.[webgl,testing]' --extra-index-url https://download.pytorch.org/whl/cpu
# - name: Cache downloads
# uses: actions/cache@v3
# with:
# path: ~/Library/Caches/tinygrad/downloads/
# key: downloads-cache-webgl-${{ env.DOWNLOAD_CACHE_VERSION }}
# - name: Prepare
# run: |
# sudo apt-get -y install xvfb
# sudo /usr/bin/Xvfb :0 -screen 0 4096x4096x24+32 &
# - name: Run selected webgl tests
# run: WEBGL=1 python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_jit.py
# - name: Build WebGL Efficientnet
# run: WEBGL=1 python -m examples.compile_efficientnet
tests:
strategy:
@@ -444,4 +476,4 @@ jobs:
# - name: Install dependencies
# run: pip install -e '.[testing,arm]' --extra-index-url https://download.pytorch.org/whl/cpu
# - name: Test arm
# run: CI=1 ARM64=1 CLANG=1 python -m pytest -n=auto test/ -k 'not (test_nn.py and (test_conv_transpose2d or test_conv2d))' --ignore=test/models --ignore=test/test_speed_v_torch.py --ignore=test/test_net_speed.py --ignore=test/test_specific_conv.py --ignore=test/unit/test_disk_tensor.py
# run: CI=1 ARM64=1 CLANG=1 python -m pytest -n=auto test/ -k 'not (test_nn.py and (test_conv_transpose2d or test_conv2d))' --ignore=test/models --ignore=test/test_speed_v_torch.py --ignore=test/test_net_speed.py --ignore=test/test_specific_conv.py --ignore=test/unit/test_disk_tensor.py

4
examples/compile_efficientnet.py
View File

@@ -9,12 +9,12 @@ import ast
if __name__ == "__main__":
model = EfficientNet(0)
model.load_from_pretrained()
mode = "clang" if getenv("CLANG", "") != "" else "webgpu" if getenv("WEBGPU", "") != "" else ""
mode = "clang" if getenv("CLANG", "") != "" else "webgpu" if getenv("WEBGPU", "") != "" else "webgl" if getenv("WEBGL", "") != "" else ""
prg, inp_sizes, out_sizes, state = export_model(model, mode, Tensor.randn(1,3,224,224))
dirname = Path(__file__).parent
if getenv("CLANG", "") == "":
safe_save(state, (dirname / "net.safetensors").as_posix())
ext = "js" if getenv("WEBGPU", "") != "" else "json"
ext = "js" if getenv("WEBGPU", "") != "" or getenv("WEBGL", "") != "" else "json"
with open(dirname / f"net.{ext}", "w") as text_file:
text_file.write(prg)
else:

23
examples/webgl/yolov8/compile.py Normal file
View File

@@ -0,0 +1,23 @@
from pathlib import Path
from examples.yolov8 import YOLOv8
from tinygrad.tensor import Tensor
from tinygrad.nn.state import safe_save
from extra.export_model import export_model
from tinygrad.helpers import fetch
from tinygrad.helpers import getenv
from tinygrad.device import Device
from tinygrad.nn.state import safe_load, load_state_dict
if __name__ == "__main__":
Device.DEFAULT = "WEBGL"
yolo_variant = 'n'
yolo_infer = YOLOv8(w=0.25, r=2.0, d=0.33, num_classes=80)
weights_location = Path(__file__).parents[1] / "weights" / f'yolov8{yolo_variant}.safetensors'
fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{yolo_variant}.safetensors', weights_location)
state_dict = safe_load(weights_location)
load_state_dict(yolo_infer, state_dict)
prg, inp_sizes, out_sizes, state = export_model(yolo_infer, Device.DEFAULT.lower(), Tensor.randn(1,3,640,640))
dirname = Path(__file__).parent
safe_save(state, (dirname / "net.safetensors").as_posix())
with open(dirname / f"net.js", "w") as text_file:
text_file.write(prg)

223
examples/webgl/yolov8/index.html Normal file
View File

@@ -0,0 +1,223 @@
<!-- Pre and post-processing functions from: https://github.com/AndreyGermanov/yolov8_onnx_javascript -->
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>YOLOv8 tinygrad WebGL</title>
<script src="./net.js"></script>
<style>
body {
text-align: center;
font-family: Arial, sans-serif;
margin: 0;
padding: 0;
overflow: hidden;
}
.video-container {
position: relative;
width: 100%;
margin: 0 auto;
}
#video, #canvas {
position: absolute;
top: 0;
left: 0;
width: 100%;
height: auto;
}
#canvas {
background: transparent;
}
h1 {
margin-top: 20px;
}
</style>
</head>
<body>
<h1>YOLOv8 tinygrad WebGL</h1>
<div class="video-container">
<video id="video" muted autoplay playsinline></video>
<canvas id="canvas"></canvas>
</div>
<script>
let net = null;
const video = document.getElementById('video');
const canvas = document.getElementById('canvas');
const context = canvas.getContext('2d');
const offscreenCanvas = document.createElement('canvas');
offscreenCanvas.width = 640;
offscreenCanvas.height = 640;
const offscreenContext = offscreenCanvas.getContext('2d');
if (navigator.mediaDevices && navigator.mediaDevices.getUserMedia) {
navigator.mediaDevices.getUserMedia({ audio: false, video: true }).then(function (stream) {
video.srcObject = stream;
video.onloadedmetadata = function() {
canvas.width = video.clientWidth;
canvas.height = video.clientHeight;
}
});
}
async function processFrame() {
offscreenContext.drawImage(video, 0, 0, 640, 640);
const boxes = await detectObjectsOnFrame(offscreenContext);
drawBoxes(offscreenCanvas, boxes);
requestAnimationFrame(processFrame);
}
requestAnimationFrame(processFrame);
function drawBoxes(offscreenCanvas, boxes) {
const canvas = document.querySelector("canvas");
const ctx = canvas.getContext("2d");
ctx.clearRect(0, 0, canvas.width, canvas.height);
ctx.lineWidth = 3;
ctx.font = "20px serif";
const scaleX = canvas.width / 640;
const scaleY = canvas.height / 640;
boxes.forEach(([x1, y1, x2, y2, label]) => {
const classIndex = yolo_classes.indexOf(label);
const color = classColors[classIndex];
const textWidth = ctx.measureText(label).width;
ctx.strokeStyle = color;
ctx.fillStyle = color;
let adjustedX1 = x1 * scaleX;
let adjustedY1 = y1 * scaleY;
let adjustedX2 = x2 * scaleX;
let adjustedY2 = y2 * scaleY;
let boxWidth = adjustedX2 - adjustedX1;
let boxHeight = adjustedY2 - adjustedY1;
ctx.strokeRect(adjustedX1, adjustedY1, boxWidth, boxHeight);
ctx.fillRect(adjustedX1, adjustedY1 - 25, textWidth + 10, 25);
ctx.fillStyle = "#000000";
ctx.fillText(label, adjustedX1, adjustedY1 - 7);
});
}
async function detectObjectsOnFrame(offscreenContext) {
if (!net) net = await loadNet();
let start = performance.now();
const [input,img_width,img_height] = await prepareInput(offscreenContext);
console.log("Preprocess took: " + (performance.now() - start) + " ms");
start = performance.now();
const output = net(new Float32Array(input));
console.log("Inference took: " + (performance.now() - start) + " ms");
start = performance.now();
let out = processOutput(output,img_width,img_height);
console.log("Postprocess took: " + (performance.now() - start) + " ms");
return out;
}
async function prepareInput(offscreenContext) {
return new Promise(resolve => {
const [img_width,img_height] = [640, 640]
const imgData = offscreenContext.getImageData(0,0,640,640);
const pixels = imgData.data;
const red = [], green = [], blue = [];
for (let index=0; index<pixels.length; index+=4) {
red.push(pixels[index]/255.0);
green.push(pixels[index+1]/255.0);
blue.push(pixels[index+2]/255.0);
}
const input = [...red, ...green, ...blue];
resolve([input, img_width, img_height])
})
}
const loadNet = async () => {
try {
const safetensor = await (new Uint8Array(await (await fetch("./net.safetensors")).arrayBuffer()));
const gl = document.createElement("canvas").getContext("webgl2");
return setupNet(gl, safetensor);
} catch (e) {
console.log(e);
return null;
}
}
function processOutput(output, img_width, img_height) {
let boxes = [];
for (let index=0;index<8400;index++) {
const [class_id,prob] = [...Array(80).keys()]
.map(col => [col, output[8400*(col+4)+index]])
.reduce((accum, item) => item[1]>accum[1] ? item : accum,[0,0]);
if (prob < 0.25) {
continue;
}
const label = yolo_classes[class_id];
const xc = output[index];
const yc = output[8400+index];
const w = output[2*8400+index];
const h = output[3*8400+index];
const x1 = (xc-w/2)/640*img_width;
const y1 = (yc-h/2)/640*img_height;
const x2 = (xc+w/2)/640*img_width;
const y2 = (yc+h/2)/640*img_height;
boxes.push([x1,y1,x2,y2,label,prob]);
}
boxes = boxes.sort((box1,box2) => box2[5]-box1[5])
const result = [];
while (boxes.length>0) {
result.push(boxes[0]);
boxes = boxes.filter(box => iou(boxes[0],box)<0.7);
}
return result;
}
function iou(box1,box2) {
return intersection(box1,box2)/union(box1,box2);
}
function union(box1,box2) {
const [box1_x1,box1_y1,box1_x2,box1_y2] = box1;
const [box2_x1,box2_y1,box2_x2,box2_y2] = box2;
const box1_area = (box1_x2-box1_x1)*(box1_y2-box1_y1)
const box2_area = (box2_x2-box2_x1)*(box2_y2-box2_y1)
return box1_area + box2_area - intersection(box1,box2)
}
function intersection(box1,box2) {
const [box1_x1,box1_y1,box1_x2,box1_y2] = box1;
const [box2_x1,box2_y1,box2_x2,box2_y2] = box2;
const x1 = Math.max(box1_x1,box2_x1);
const y1 = Math.max(box1_y1,box2_y1);
const x2 = Math.min(box1_x2,box2_x2);
const y2 = Math.min(box1_y2,box2_y2);
return (x2-x1)*(y2-y1)
}
const yolo_classes = [
'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase',
'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant',
'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
];
function generateColors(numColors) {
const colors = [];
for (let i = 0; i < 360; i += 360 / numColors) {
colors.push(`hsl(${i}, 100%, 50%)`);
}
return colors;
}
const classColors = generateColors(yolo_classes.length);
</script>
</body>
</html>

70
extra/backends/cstyle.py Normal file
View File

@@ -0,0 +1,70 @@
# TODO: how much of this can be merged with above?
class WGSLLanguage(CStyleLanguage):
code_for_workitem = {"g": lambda x: f"i32(gindex.{'xyz'[x]})", "l": lambda x: f"i32(lindex.{'xyz'[x]})"}
size_prefix = "let"
barrier="workgroupBarrier();"
generic_var_prefix = "var "
external_local_bufs = True
code_for_op = { **CStyleLanguage().code_for_op,
BinaryOps.CMPLT: lambda x,y,dtype: f"f32({x}<{y})", BinaryOps.CMPEQ: lambda x,y,dtype: f"f32({x}=={y})",
TernaryOps.MULACC: lambda x,y,z,dtype: f"fma({x},{y},{z})", TernaryOps.WHERE: lambda a,b,c,dtype: f"select({c},{b},bool({a}))" }
# HACK: write bool as f32
type_map = {dtypes.float: "f32", dtypes.half: "f16", dtypes.int32: "i32", dtypes.uint32: "u32", dtypes.bool: "f32"}
def render_local(self, name: str, dtype:DType, size: int): return f"var<workgroup> {name}: array<{self.type_map[dtype]},{size}>;"
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 "") + "inf(1.0)"
return f"({super().render_const(x, var_dtype)})"
def render_if(self, cond: str): return f"if (bool({cond})) {{"
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); }\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
def render_cast(self, x:List[str], var_dtype:DType, bitcast=False) -> str:
if self.type_map[var_dtype]: return f"bitcast<{self.type_map[var_dtype]}>({x[0]})" if bitcast else f"{self.type_map[var_dtype]}({x[0]})"
raise NotImplementedError(f"no cast for {var_dtype}")
WGSLRenderer = functools.partial(uops_to_cstyle, WGSLLanguage())
class GLSLLanguage(CStyleLanguage):
type_map = {dtypes.float: "float", dtypes.half: "float", dtypes.int32: "int", dtypes.uint32: "uint", dtypes.bool: "bool"}
sampler_prefix = {dtypes.float64: "d", dtypes.float: "", dtypes.half: "", dtypes.int32: "i", dtypes.uint32: "u", dtypes.bool: "i"}
fragment_center_offset = 0.5
code_for_workitem = {"i": lambda x, offset=fragment_center_offset:f"int(gl_FragCoord.y-{offset}) * width + int(gl_FragCoord.x-{offset})"}
code_for_op = {**CStyleLanguage().code_for_op, **{op: lambda a,b,dtype,charforop=charforop: f"bool(int({a}){charforop}int({b}))" \
if dtype == dtypes.bool else f"({a}{charforop}{b})" for op,charforop in [(BinaryOps.MUL,"*"),(BinaryOps.ADD,"+"),(BinaryOps.DIV,"/")]},
BinaryOps.CMPLT: lambda a,b,dtype: f"(float({a})<float({b}))" if dtype == dtypes.bool else f"({a}<{b})",
BinaryOps.MOD: lambda a,b,dtype: f"(int({a})%int({b}))", TernaryOps.WHERE: lambda a,b,c,dtype: f"(float({a})!=0.0?{b}:{c})"}
def render_const(self, x:Union[float,int], var_dtype) -> str:
if math.isnan(x): return "(0.0 / 0.0)"
elif math.isinf(x): return ("-" if x < 0 else "") + "(1./0.)"
return self.render_cast(["({:.1f})".format(x) if x == int(x) and dtypes.is_float(var_dtype) else f"({x})"]*var_dtype.sz, 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:
prg = "#version 330\nprecision highp float;\nprecision highp int;\nin vec2 uv;\nuniform int width;\n"
prg += "\n".join([f"uniform {self.sampler_prefix[dtype]}sampler2D {name};" for name,dtype in bufs if name != "data0"])
prg += f"\nout {'int' if bufs[0][1] == dtypes.bool else self.type_map[bufs[0][1]]} out_data;\n"
return prg + "\nvoid main() {\n" + "\n".join(kernel) + "\n}"
def render_cast(self, x:List[str], var_dtype:DType, bitcast=False) -> str:
if self.type_map[var_dtype]: return f"{self.type_map[var_dtype]}({x[0]})"
raise NotImplementedError(f"no cast for {var_dtype}")
def render_load(self, output_dtype, buf_name, buf_dtype, idx, local=False) -> str:
x_calc = f"float(int({idx})%textureSize({buf_name}, 0).x)"
y_calc = f"float(int({idx})/textureSize({buf_name}, 0).x)"
out_val = f"texture({buf_name}, vec2(float({x_calc} + {self.fragment_center_offset}f)/float(textureSize({buf_name}, 0).x),\
float({y_calc} + {self.fragment_center_offset}f)/float(textureSize({buf_name}, 0).y))).r"
return f"{self.render_cast([out_val], output_dtype)}"
def render_store(self, buf_name:str, buf_dtype:DType, var_name:str, var_dtype:DType, idx, local=False) -> str:
return f"out_data = {'int' if buf_dtype == dtypes.bool else self.type_map[buf_dtype]}({var_name});"

52
extra/backends/ops_webgl.py Normal file
View File

@@ -0,0 +1,52 @@
import numpy as np
import functools
from tinygrad.dtype import dtypes, ImageDType
from tinygrad.device import Compiled, Allocator
from tinygrad.codegen.kernel import LinearizerOptions, OptOps
from tinygrad.renderer.cstyle import uops_to_cstyle
from tinygrad.renderer.cstyle import GLSLLanguage
import moderngl
ctx = moderngl.create_standalone_context()
max_dims = 4096
dtype_map = { dtypes.float64: "f8", dtypes.float: "f4", dtypes.half: "f2", dtypes.int32: "i4", dtypes.uint32: "u4", dtypes.bool: "i1"}
vertex_shader="#version 330\nprecision highp float;\nin vec2 in_position;in vec2 in_uv;out vec2 uv;void main(){\
gl_Position=vec4(in_position,0.0,1.0);uv=in_uv;}"
class WebGLProgram:
def __init__(self, name: str, prg: str, bufs:int=0, vars:int=0):
self.name, self.prg = name, ctx.program(vertex_shader=vertex_shader, fragment_shader=prg)
def __call__(self, *bufs, global_size, local_size=None, vals=(), wait=False):
vert = ctx.buffer(np.asarray([-1, 1, -1, -1, 1, 1, 1, -1], dtype='f4').tobytes())
uv = ctx.buffer(np.asarray([0, 1, 0, 0, 1, 1, 1, 0], dtype='f4').tobytes())
self.vao = ctx.vertex_array(self.prg, [])
self.vao.bind(self.prg["in_position"].location if "in_position" in self.prg else 0, buffer=vert, cls='f', fmt='2f4')
self.vao.bind(self.prg["in_uv"].location if "in_uv" in self.prg else 1, buffer=uv, cls='f', fmt='2f4')
self.vao.vertices = vert.size//4//2
self.fbo = ctx.framebuffer(color_attachments=[bufs[0]])
for i, x in enumerate(bufs[1:], start=1):
if f"data{i}" in self.prg:
self.prg[f"data{i}"] = i
x.use(i)
if ("width" in self.prg): self.prg["width"].value = self.fbo.size[0]
ctx.viewport = (0, 0, self.fbo.size[0], self.fbo.size[1])
self.fbo.use()
self.vao.render(mode=moderngl.TRIANGLE_STRIP)
class RawWebGLAllocator(Allocator):
def _alloc_image(self, dtype:ImageDType):
tex = ctx.texture(dtype.shape, 1, dtype=dtype_map[dtype.base])
tex.filter = (moderngl.NEAREST, moderngl.NEAREST)
return tex
def copyin(self, dest:moderngl.Texture, src: memoryview): dest.write(src)
def copyout(self, dest:memoryview, src: moderngl.Texture):
src.read_into(dest)
return dest
class WebGlDevice(Compiled):
def __init__(self, device:str):
super().__init__(RawWebGLAllocator(),
LinearizerOptions(device="WEBGL", global_max=[4096*4096,1,1], unsupported_opts=[OptOps.UPCAST, OptOps.UPCASTMID],
supports_float4=False, supports_float4_alu=False, has_local=False, has_shared=False, dont_use_locals=True),
functools.partial(uops_to_cstyle, GLSLLanguage()), lambda x: x, WebGLProgram)

195
extra/export_model.py
View File

@@ -3,9 +3,21 @@ from tinygrad.dtype import DType
from tinygrad.tensor import Device, Tensor
from tinygrad.jit import TinyJit
from tinygrad.nn.state import get_state_dict
from tinygrad.dtype import dtypes
import json
EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CLANG", "CUDA", "GPU"]
EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "WEBGL", "CLANG", "CUDA", "GPU"]
web_utils = {
"getTensorBuffer":
"""const getTensorBuffer = (safetensorBuffer, tensorMetadata) => {
return safetensorBuffer.subarray(...tensorMetadata.data_offsets);
}""",
"getTensorMetadata": """const getTensorMetadata = (safetensorBuffer) => {
const metadataLength = Number(new DataView(safetensorBuffer.buffer).getBigUint64(0, true));
const metadata = JSON.parse(new TextDecoder("utf8").decode(safetensorBuffer.subarray(8, 8 + metadataLength)));
return Object.fromEntries(Object.entries(metadata).filter(([k, v]) => k !== "__metadata__").map(([k, v]) => [k, {...v, data_offsets: v.data_offsets.map(x => 8 + metadataLength + x)}]));
};"""
}
def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
@@ -66,6 +78,173 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
cprog += [f"void net({inputs}, {outputs}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
return '\n'.join(cprog)
def export_model_webgl(functions, statements, bufs, bufs_to_save, weight_names, input_names, output_names) -> str:
header = f"""
function setupNet(gl, safetensor) {{
function createShaderProgram(gl, code) {{
const vertexShader = loadShader(gl, gl.VERTEX_SHADER, '#version 300 es\\nin vec2 in_position;in vec2 in_uv;out vec2 uv;void main(){{gl_Position=vec4(in_position,0.0,1.0);uv=in_uv;}}');
const fragmentShader = loadShader(gl, gl.FRAGMENT_SHADER, code);
const shaderProgram = gl.createProgram();
gl.attachShader(shaderProgram, vertexShader);
gl.attachShader(shaderProgram, fragmentShader);
gl.linkProgram(shaderProgram);
if (!gl.getProgramParameter(shaderProgram, gl.LINK_STATUS)) {{
console.log(`Unable to initialize the shader program: ${{gl.getProgramInfoLog(shaderProgram)}}`);
return null;
}}
return shaderProgram;
}}
function loadShader(gl, type, source) {{
const shader = gl.createShader(type);
gl.shaderSource(shader, source);
gl.compileShader(shader);
if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {{
console.log(`An error occurred compiling the shaders: ${{gl.getShaderInfoLog(shader)}}`);
gl.deleteShader(shader);
return null;
}}
return shader;
}}
function setupVertexData(gl, program, vertices) {{
let vao = gl.createVertexArray();
gl.bindVertexArray(vao);
let vertexBuffer = gl.createBuffer();
gl.bindBuffer(gl.ARRAY_BUFFER, vertexBuffer);
gl.bufferData(gl.ARRAY_BUFFER, new Float32Array(vertices), gl.STATIC_DRAW);
const positionLocation = gl.getAttribLocation(program, 'in_position');
const uvLocation = gl.getAttribLocation(program, 'in_uv');
gl.enableVertexAttribArray(positionLocation);
gl.vertexAttribPointer(positionLocation, 2, gl.FLOAT, false, 4 * 4, 0);
gl.enableVertexAttribArray(uvLocation);
gl.vertexAttribPointer(uvLocation, 2, gl.FLOAT, false, 4 * 4, 2 * 4);
gl.bindVertexArray(null);
return vao;
}}
function runProgram(gl, kernelName, program, textures) {{
let framebuffer = gl.createFramebuffer();
gl.bindFramebuffer(gl.FRAMEBUFFER, framebuffer);
gl.framebufferTexture2D(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, textures[0].tex, 0);
gl.useProgram(program);
gl.uniform1i(gl.getUniformLocation(program, "width"), textures[0].width);
const vao = setupVertexData(gl, program, [-1, 1, 0, 1, -1, -1, 0, 0, 1, 1, 1, 1, 1, -1, 1, 0]);
gl.bindVertexArray(vao);
// Texture 0 is the framebuffer texture, so we skip that
for (let i = 1; i < textures.length; i++) {{
gl.activeTexture(gl.TEXTURE0 + i-1);
gl.bindTexture(gl.TEXTURE_2D, textures[i].tex);
gl.uniform1i(gl.getUniformLocation(program, 'data' + i), i-1);
}}
gl.viewport(0, 0, textures[0].width, textures[0].height);
gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4);
gl.bindFramebuffer(gl.FRAMEBUFFER, null);
for (let i = 1; i < textures.length; i++) {{
gl.activeTexture(gl.TEXTURE0 + i-1);
gl.bindTexture(gl.TEXTURE_2D, null);
}}
console.log("Finished running: " + kernelName);
}}
function limitTextureDims(size, threshold) {{
if (size <= threshold) {{ return [size, 1] }};
for (let i = 2; i < threshold + 1; i++) {{
if ((size % i == 0) && (Math.floor(size / i) <= threshold)) {{
return [Math.floor(size / i), i];
}}
}}
return [size, 1];
}}
function updateTextureData(gl, texture, data, isHalf) {{
gl.bindTexture(gl.TEXTURE_2D, texture.tex);
gl.texSubImage2D(gl.TEXTURE_2D, 0, 0, 0, texture.width, texture.height, gl.RED, (isHalf) ? gl.HALF_FLOAT : gl.FLOAT, data);
gl.bindTexture(gl.TEXTURE_2D, null);
}}
function readTextureData(gl, texture) {{
const framebuffer = gl.createFramebuffer();
gl.bindFramebuffer(gl.FRAMEBUFFER, framebuffer);
gl.framebufferTexture2D(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, texture.tex, 0);
if (gl.checkFramebufferStatus(gl.FRAMEBUFFER) !== gl.FRAMEBUFFER_COMPLETE) {{
throw new Error('Framebuffer not complete');
}}
let data = new Float32Array(texture.width * texture.height);
gl.readPixels(0, 0, texture.width, texture.height, gl.RED, gl.FLOAT, data);
gl.bindFramebuffer(gl.FRAMEBUFFER, null);
gl.deleteFramebuffer(framebuffer);
return data;
}}
function createTexture(gl, size, isHalf, tensorBuffer) {{
const texture = gl.createTexture();
gl.bindTexture(gl.TEXTURE_2D, texture);
const internalFormat = gl.RGBA;
const texSize = limitTextureDims(size, gl.getParameter(gl.MAX_TEXTURE_SIZE));
let weights;
if (tensorBuffer != null) {{
if (!isHalf)
weights = new Float32Array(tensorBuffer.buffer, tensorBuffer.byteOffset, tensorBuffer.byteLength / Float32Array.BYTES_PER_ELEMENT);
else
weights = new Uint16Array(tensorBuffer.buffer, tensorBuffer.byteOffset, tensorBuffer.byteLength / Uint16Array.BYTES_PER_ELEMENT);
}} else {{
if (!isHalf)
weights = new Float32Array(size).fill(0.0);
else
weights = new Uint16Array(size).fill(0.0);
}}
if (size != weights.length)
console.log("Weights length: " + weights.length + ", texsize: " + texSize[0]*texSize[1]);
gl.texImage2D(gl.TEXTURE_2D, 0, (isHalf) ? gl.R16F : gl.R32F, texSize[0], texSize[1], 0, gl.RED, (isHalf) ? gl.HALF_FLOAT : gl.FLOAT, weights);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
gl.bindTexture(gl.TEXTURE_2D, null);
return {{ tex: texture, width: texSize[0], height: texSize[1] }};
}}
{web_utils["getTensorBuffer"]}
{web_utils["getTensorMetadata"]}
const metadata = getTensorMetadata(safetensor);
"""
textures = '\n '.join([f"const {name} = " + (f"createTexture(gl, {size/(2 if dtype == dtypes.half else 4)}, {'true' if dtype == dtypes.half else 'false'});" if _key not in weight_names else f"createTexture(gl, {size/(2 if dtype == dtypes.half else 4)}, {'true' if dtype == dtypes.half else 'false'}, getTensorBuffer(safetensor, metadata['{weight_names[_key]}']))") + ";" for name,(size,dtype,_key) in bufs.items()])
kernels = '\n\n'.join([f"const {key} = `{code.replace(key, 'main').replace('version 330', 'version 300 es')}`;" for key, code in functions.items()])
kernel_names = ', '.join([name for (name, _args, _global_size, _local_size) in statements])
kernel_calls = '\n '.join([f"runProgram(gl, '{name}', programs[{i}], [{', '.join(args)}]);" for i, (name, args, _global_size, _local_size) in enumerate(statements) ])
copy_inputs = "\n".join([f'updateTextureData(gl, {name}, _{name}, {"true" if dtype == dtypes.half else "false"});' for name,(size,dtype,_key) in bufs.items() if "input" in name])
entry_point = f"""
return function({",".join([f"_{name}" for name,(size,dtype,_key) in bufs.items() if "input" in name])}) {{
const ext = gl.getExtension('EXT_color_buffer_float');
{copy_inputs}
{kernel_calls}
return readTextureData(gl, output0);
}}
"""
programs = f"let programs = [{kernel_names}].map((code) => createShaderProgram(gl, code));"
return f"{header}\n{kernels}\n{textures}\n{programs}\n{entry_point}}}"
def export_model_webgpu(functions, statements, bufs, bufs_to_save, weight_names, input_names, output_names) -> Tuple[str,int,int]:
kernel_code = '\n\n'.join([f"const {key} = `{code.replace(key, 'main')}`;" for key, code in functions.items()])
kernel_names = ', '.join([name for (name, _args, _global_size, _local_size) in statements])
@@ -78,15 +257,9 @@ def export_model_webgpu(functions, statements, bufs, bufs_to_save, weight_names,
output_readers = '\n '.join([f"await gpuReadBuffer{i}.mapAsync(GPUMapMode.READ);\n const resultBuffer{i} = new Float32Array(gpuReadBuffer{i}.size);\n resultBuffer{i}.set(new Float32Array(gpuReadBuffer{i}.getMappedRange()));\n gpuReadBuffer{i}.unmap();" for i in range(len(output_names))])
output_return = '[{}]'.format(",".join([f'resultBuffer{i}' for i in range(len(output_names))]))
return f"""
const getTensorMetadata = (safetensorBuffer) => {{
const metadataLength = Number(new DataView(safetensorBuffer.buffer).getBigUint64(0, true));
const metadata = JSON.parse(new TextDecoder("utf8").decode(safetensorBuffer.subarray(8, 8 + metadataLength)));
return Object.fromEntries(Object.entries(metadata).filter(([k, v]) => k !== "__metadata__").map(([k, v]) => [k, {{...v, data_offsets: v.data_offsets.map(x => 8 + metadataLength + x)}}]));
}};
{web_utils["getTensorBuffer"]}
const getTensorBuffer = (safetensorBuffer, tensorMetadata) => {{
return safetensorBuffer.subarray(...tensorMetadata.data_offsets);
}}
{web_utils["getTensorMetadata"]}
const createEmptyBuf = (device, size) => {{
return device.createBuffer({{size, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST }});
@@ -137,7 +310,7 @@ const setupNet = async (device, safetensor) => {{
""" + f"\n\nconst loadNet = async (device) => {{ return await fetch('net.safetensors').then(x => x.arrayBuffer()).then(x => setupNet(device, new Uint8Array(x))); }}"
def export_model(model, target:str, *inputs):
assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, "only WEBGPU, CLANG, CUDA, GPU, METAL are supported"
assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, "only WEBGPU, WEBGL, CLANG, CUDA, GPU, METAL are supported"
run,special_names = jit_model(model, *inputs)
functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
state = get_state_dict(model)
@@ -149,6 +322,8 @@ def export_model(model, target:str, *inputs):
prg = export_model_clang(functions, statements, bufs, bufs_to_save, input_names, output_names)
elif target == "webgpu":
prg = export_model_webgpu(functions, statements, bufs, bufs_to_save, weight_names, input_names, output_names)
elif target == "webgl":
prg = export_model_webgl(functions, statements, bufs, bufs_to_save, weight_names, input_names, output_names)
else:
prg = json.dumps({
"backend": Device.DEFAULT,

4
test/test_dtype.py
View File

@@ -13,7 +13,7 @@ core_dtypes = list(DTYPES_DICT.values())
floats = [dt for dt in core_dtypes if dtypes.is_float(dt)]
def is_dtype_supported(dtype: DType, device: str = Device.DEFAULT):
if dtype == dtypes.bfloat16: return False # numpy doesn't support bf16, tested separately in TestBFloat16DType
if device == "WEBGPU": return dtype in [dtypes.float, dtypes.int32, dtypes.uint32]
if device in ["WEBGPU", "WEBGL"]: return dtype in [dtypes.float, dtypes.int32, dtypes.uint32]
if device == "TORCH": return dtype not in [dtypes.uint16, dtypes.uint32, dtypes.uint64]
# for CI GPU, cl_khr_fp16 isn't supported
# for CI LLVM, it segfaults because it can't link to the casting function
@@ -90,6 +90,7 @@ class TestDType(unittest.TestCase):
get_available_cast_dtypes(self.DTYPE)
))
def test_bitcast(self):
if Device.DEFAULT == "WEBGL": raise unittest.SkipTest("no bitcast in WebGL GLSL")
if self.DTYPE == dtypes.bool: raise unittest.SkipTest("no bools in bitcast")
list(map(
lambda dtype:
@@ -160,6 +161,7 @@ class TestUint8Dtype(TestDType):
def test_uint8_to_int8_overflow(self):
_test_op(lambda: Tensor([255, 254, 253, 252], dtype=dtypes.uint8).cast(dtypes.int8), dtypes.int8, [-1, -2, -3, -4])
@unittest.skipIf(Device.DEFAULT == "WEBGL", "No bitcast on WebGL")
class TestBitCast(unittest.TestCase):
def test_shape_change_bitcast(self):
with self.assertRaises(AssertionError):

6
test/test_ops.py
View File

@@ -600,10 +600,12 @@ class TestOps(unittest.TestCase):
helper_test_op([()], lambda x: x.mean())
def test_mean_axis(self):
helper_test_op([(3,4,5,6)], lambda x: x.mean(axis=(1,2)), lambda x: Tensor.mean(x, axis=(1,2)))
@unittest.skipIf(Device.DEFAULT == "WEBGL" and CI, "Only broken on CI")
def test_std(self):
helper_test_op([(45, 65, 85)], lambda x: torch.std(x), lambda x: Tensor.std(x))
helper_test_op([(45, 65, 85)], lambda x: torch.std(x, dim=None, correction=0), lambda x: Tensor.std(x, correction=0))
helper_test_op([(45, 65, 85)], lambda x: torch.std(x, dim=None, correction=5), lambda x: Tensor.std(x, correction=5))
@unittest.skipIf(Device.DEFAULT == "WEBGL" and CI, "Only broken on CI")
def test_std_axis(self):
helper_test_op([(45, 65, 85)], lambda x: torch.std(x, dim=0), lambda x: Tensor.std(x, axis=0))
helper_test_op([(45, 65, 85)], lambda x: torch.std(x, dim=2), lambda x: Tensor.std(x, axis=2))
@@ -613,6 +615,7 @@ class TestOps(unittest.TestCase):
helper_test_op([(45, 65, 85)], lambda x: torch.std(x, correction=0, dim=2), lambda x: Tensor.std(x, axis=2, correction=0))
helper_test_op([(45, 65, 85)], lambda x: torch.std(x, correction=0, dim=[1, 2]), lambda x: Tensor.std(x, axis=[1, 2], correction=0))
helper_test_op([(45, 65, 85)], lambda x: torch.std(x, correction=0, dim=None), lambda x: Tensor.std(x, axis=None, correction=0))
@unittest.skipIf(Device.DEFAULT == "WEBGL" and CI, "Only broken on CI")
def test_std_keepdim(self):
helper_test_op([(45, 65, 85)], lambda x: torch.std(x, dim=None, keepdim=True), lambda x: Tensor.std(x, keepdim=True))
helper_test_op([(45, 65, 85)], lambda x: torch.std(x, dim=0, keepdim=True, correction=0),
@@ -729,7 +732,7 @@ class TestOps(unittest.TestCase):
np.testing.assert_allclose(a[:, 2:0:-1], t[:, 2:0:-1].numpy())
np.testing.assert_allclose(a[:, 2:0:-1, 3:1:-2], t[:, 2:0:-1, 3:1:-2].numpy())
np.testing.assert_allclose(a[4:0:-3, 2:0:-1, -1:-5:-2], t[4:0:-3, 2:0:-1, -1:-5:-2].numpy())
if Device.DEFAULT != "CPU":
if Device.DEFAULT not in ["CPU"]:
# broken
np.testing.assert_allclose(a[2:5:-1, :, :], t[2:5:-1, :, :].numpy()) # shape = (0, 10, 10)
np.testing.assert_allclose(a[:, 2:5:-1, :], t[:, 2:5:-1, :].numpy()) # shape = (0, 10, 10)
@@ -781,6 +784,7 @@ class TestOps(unittest.TestCase):
helper_test_op([(3,3)], lambda x: torch.nn.functional.pad(x, (0,0,3,4), value=1), lambda x: x.pad(((3,4), None), value=1))
helper_test_op([(3,3)], lambda x: torch.nn.functional.pad(x, (0,0,0,0), value=1), lambda x: x.pad((None, None), value=1))
@unittest.skipIf(Device.DEFAULT == "WEBGL", "incorrect result")
def test_pad_slice(self):
for value in 0., 3.456:
helper_test_op([(1)], lambda x: torch.nn.functional.pad(x,(1,0), value=value)[0], lambda x: x.pad(((1,0),), value=value)[0])