├── .gitignore
├── README.md
├── build.py
├── index.html
└── publish_gh_pages.sh


/.gitignore:
--------------------------------------------------------------------------------
1 | *~
2 | dist
3 | *.bak
4 | .gitignore
5 | .pkl*
6 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | TVM WebGPU Example
 2 | ==================
 3 | 
 4 | NOTE: This example is out dated, checkout https://github.com/mlc-ai/web-stable-diffusion instead
 5 | 
 6 | This is an example project for building a tvm webgpu-backed module and deploy it to the web.
 7 | 
 8 | ## Try it out
 9 | 
10 | - You can try it out at https://tqchen.com/tvm-webgpu-example/
11 | - You will need a WebGPU enabled browser
12 | 
13 | 
14 | ## Build from Source
15 | 
16 | You will need the latest version of TVM on the master.
17 | Please install the TVM via source build(by setting the PYTHONPATH).
18 | Checkout the steps under `tvm/web` to enable webgpu support.
19 | 
20 | Build the dependent files in the `tvm/web`
21 | ```bash
22 | cd /path/to/tvm/web
23 | make
24 | npm run bundle
25 | ```
26 | 
27 | Then run the build script to build the model specific data.
28 | 
29 | ```bash
30 | cd /path/to/tvm-webgpu-examples
31 | python3 build.py
32 | ```
33 | 
34 | You can now serve the content under `dist`
35 | ```bash
36 | cd /path/to/tvm-webgpu-examples/dist
37 | python3 http.server
38 | ```
39 | 
40 | You can also publish to github using the `publish_gh_pages` script.
41 | 


--------------------------------------------------------------------------------
/build.py:
--------------------------------------------------------------------------------
 1 | """Example script to create tvm wasm modules and deploy it."""
 2 | 
 3 | import argparse
 4 | import os
 5 | from tvm import relay
 6 | from tvm.contrib import util, emcc
 7 | from tvm import rpc
 8 | from tvm.contrib.download import download_testdata
 9 | import tvm
10 | from tvm._ffi import libinfo
11 | import shutil
12 | import logging
13 | import json
14 | 
15 | curr_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__)))
16 | 
17 | 
18 | def build_module(opts):
19 |     network = opts.network
20 |     build_dir = opts.out_dir
21 | 
22 |     dshape = (1, 3, 224, 224)
23 |     from mxnet.gluon.model_zoo.vision import get_model
24 |     block = get_model(network, pretrained=True)
25 | 
26 |     shape_dict = {"data": dshape}
27 |     mod, params = relay.frontend.from_mxnet(block, shape_dict)
28 |     func = mod["main"]
29 |     func = relay.Function(func.params,
30 |                           relay.nn.softmax(func.body), None,
31 |                           func.type_params, func.attrs)
32 | 
33 |     target_device = "webgpu -modle=1080ti"
34 |     target_host = "llvm -target=wasm32-unknown-unknown-wasm -system-lib"
35 | 
36 |     with relay.build_config(opt_level=3):
37 |         graph, lib, params = relay.build(
38 |             func, target=target_device, target_host=target_host, params=params)
39 | 
40 |     wasm_path = os.path.join(build_dir, f"{network}.wasm")
41 |     lib.export_library(wasm_path, emcc.create_tvmjs_wasm)
42 |     print(f"Exporting library to {build_dir}...")
43 | 
44 |     with open(os.path.join(build_dir, f"{network}.json"), "w") as f_graph_json:
45 |         f_graph_json.write(graph)
46 |     with open(os.path.join(build_dir, f"{network}.params"), "wb") as f_params:
47 |         f_params.write(relay.save_param_dict(params))
48 | 
49 | 
50 | def prepare_data(opts):
51 |     """Build auxiliary data file."""
52 |     # imagenet synset
53 |     synset_url = ''.join(['https://gist.githubusercontent.com/zhreshold/',
54 |                           '4d0b62f3d01426887599d4f7ede23ee5/raw/',
55 |                           '596b27d23537e5a1b5751d2b0481ef172f58b539/',
56 |                         'imagenet1000_clsid_to_human.txt'])
57 |     synset_name = 'imagenet1000_clsid_to_human.txt'
58 |     synset_path = download_testdata(synset_url, synset_name, module='data')
59 |     with open(synset_path) as f:
60 |         synset = eval(f.read())
61 |     build_dir = opts.out_dir
62 |     with open(os.path.join(build_dir, "imagenet1k_synset.json"), "w") as f_synset:
63 |         f_synset.write(json.dumps(synset))
64 |     # Test image
65 |     image_url = "https://homes.cs.washington.edu/~moreau/media/vta/cat.jpg"
66 |     image_fn = os.path.join(build_dir, "cat.png")
67 |     download_testdata(image_url, image_fn)
68 |     # copy runtime js files.
69 |     shutil.copyfile(libinfo.find_lib_path("tvmjs.bundle.js")[0],
70 |                     os.path.join(build_dir, "tvmjs.bundle.js"))
71 |     shutil.copyfile(libinfo.find_lib_path("tvmjs_runtime.wasi.js")[0],
72 |                     os.path.join(build_dir, "tvmjs_runtime.wasi.js"))
73 |     shutil.copyfile(os.path.join(curr_dir, "index.html"),
74 |                     os.path.join(build_dir, "index.html"))
75 | 
76 | 
77 | if __name__ == "__main__":
78 |     parser = argparse.ArgumentParser()
79 |     parser.add_argument("-o", "--out-dir", default=os.path.join(curr_dir, "dist"))
80 |     parser.add_argument("--network", default="mobilenet1.0")
81 |     parser.add_argument('-p', '--prepare', action='store_true')
82 | 
83 |     opts = parser.parse_args()
84 | 
85 |     build_dir = os.path.abspath(opts.out_dir)
86 |     if not os.path.isdir(build_dir):
87 |         os.makedirs(build_dir)
88 |         opts.out_dir = build_dir
89 | 
90 |     if opts.prepare:
91 |         prepare_data(opts)
92 |     else:
93 |         prepare_data(opts)
94 |         build_module(opts)
95 | 


--------------------------------------------------------------------------------
/index.html:
--------------------------------------------------------------------------------
  1 | <html>
  2 |   <head>
  3 |     <title>TVM WebGPU MobileNet Demo</title>
  4 |   </head>
  5 |   <script src="tvmjs_runtime.wasi.js"></script>
  6 |   <script src="tvmjs.bundle.js"></script>
  7 |   <script>
  8 |     function customLog(message) {
  9 |       console.log(message);
 10 |       const d = document.createElement("div");
 11 |       d.innerHTML = message;
 12 |       document.getElementById("log").appendChild(d);
 13 |     }
 14 | 
 15 |     function preprocImage(imageData) {
 16 |       tvmjs.assert(imageData instanceof ImageData, "Input must be ImageData.");
 17 |       tvmjs.assert(imageData.width == 224, "Width must be 224.");
 18 |       tvmjs.assert(imageData.height == 224, "Height must be 224.");
 19 | 
 20 |       const width = imageData.width;
 21 |       const height = imageData.height;
 22 |       const npixels = width * height;
 23 | 
 24 |       const rgbaU8 = imageData.data;
 25 |       tvmjs.assert(rgbaU8.length == npixels * 4, "Image should be RGBA.");
 26 | 
 27 |       // Drop alpha channel. Resnet does not need it.
 28 |       const rgbU8 = new Uint8Array(npixels * 3);
 29 |       for (let i = 0; i < npixels; ++i) {
 30 |         rgbU8[i * 3] = rgbaU8[i * 4];
 31 |         rgbU8[i * 3 + 1] = rgbaU8[i * 4 + 1];
 32 |         rgbU8[i * 3 + 2] = rgbaU8[i * 4 + 2];
 33 |       }
 34 | 
 35 |       // Cast to float and normalize.
 36 |       const rgbF32 = new Float32Array(npixels * 3);
 37 |       for (let i = 0; i < npixels; ++i) {
 38 |         rgbF32[i * 3] = (rgbU8[i * 3] - 123.0) / 58.395;
 39 |         rgbF32[i * 3 + 1] = (rgbU8[i * 3 + 1] - 117.0) / 57.12;
 40 |         rgbF32[i * 3 + 2] = (rgbU8[i * 3 + 2] - 104.0) / 57.375;
 41 |       }
 42 | 
 43 |       // Transpose. Resnet expects 3 greyscale images.
 44 |       const data = new Float32Array(npixels * 3);
 45 |       for (let i = 0; i < npixels; ++i) {
 46 |         data[i] = rgbF32[i * 3];
 47 |         data[npixels + i] = rgbF32[i * 3 + 1];
 48 |         data[npixels * 2 + i] = rgbF32[i * 3 + 2];
 49 |       }
 50 |       return data;
 51 |     }
 52 | 
 53 |     // Asynchrously load an image.
 54 |     function loadImage(uri) {
 55 |       return new Promise((resolve, reject) => {
 56 |         const image = new Image();
 57 |         image.src = uri;
 58 |         image.setAttribute("crossOrigin", "anonymous");
 59 |         image.onload = () => resolve(image);
 60 |         image.onerror = reject;
 61 |       });
 62 |     }
 63 |     // global classifier object
 64 |     classifier = {};
 65 | 
 66 |     // initialize the classifier.
 67 |     async function initModel(network, logger) {
 68 |       const wasmSource = await (
 69 |         await fetch("./" + network + ".wasm")
 70 |       ).arrayBuffer();
 71 |       const inst = await tvmjs.instantiate(
 72 |         new Uint8Array(wasmSource),
 73 |         new EmccWASI(),
 74 |         logger
 75 |       );
 76 | 
 77 |       const gpuDevice = await tvmjs.detectGPUDevice();
 78 |       if (gpuDevice === undefined) {
 79 |         logger(
 80 |           "Cannot find WebGPU device, make sure you use the browser that suports webGPU"
 81 |         );
 82 |         return;
 83 |       }
 84 |       inst.initWebGPU(gpuDevice);
 85 |       const graphJson = await (await fetch("./" + network + ".json")).text();
 86 |       const synset = await (await fetch("./imagenet1k_synset.json")).json();
 87 |       const paramsBinary = new Uint8Array(
 88 |         await (await fetch("./" + network + ".params")).arrayBuffer()
 89 |       );
 90 |       logger("Start to intialize the classifier with WebGPU...");
 91 |       ctx = inst.webgpu(0);
 92 |       const syslib = inst.systemLib();
 93 |       const executor = inst.createGraphRuntime(graphJson, syslib, ctx);
 94 |       executor.loadParams(paramsBinary);
 95 |       const inputData = inst.empty([1, 3, 224, 224], "float32", inst.cpu());
 96 |       const outputData = inst.empty([1, 1000], "float32", inst.cpu());
 97 |       const outputGPU = executor.getOutput(0);
 98 |       // run the first time to make sure all weights are populated.
 99 |       executor.run();
100 |       await ctx.sync();
101 |       logger("Finish initializing the classifier with WebGPU support");
102 | 
103 |       // the classify function
104 |       classifier.classify = async () => {
105 |         // preprocessing image
106 |         const imageURL = document.getElementById("imageURL").value;
107 |         const image = await loadImage(imageURL);
108 |         // resize the image to 224 224
109 |         const sourceWidth = image.width;
110 |         const sourceHeight = image.height;
111 |         const shortEdge = Math.min(image.width, image.height);
112 |         const yy = Math.floor((sourceHeight - shortEdge) / 2);
113 |         const xx = Math.floor((sourceWidth - shortEdge) / 2);
114 |         const imageCanvas = document.getElementById("canvas");
115 |         const imageCanvasContext = imageCanvas.getContext("2d");
116 |         imageCanvasContext.drawImage(
117 |           image,
118 |           xx,
119 |           yy,
120 |           shortEdge,
121 |           shortEdge,
122 |           0,
123 |           0,
124 |           224,
125 |           224
126 |         );
127 |         // preprocessing
128 |         const imageData = imageCanvasContext.getImageData(0, 0, 224, 224);
129 |         const processedImage = preprocImage(imageData);
130 |         // classify the image.
131 |         const tStart = performance.now();
132 |         inputData.copyFrom(processedImage);
133 |         executor.setInput("data", inputData);
134 |         executor.run();
135 |         outputData.copyFrom(outputGPU);
136 |         // need t sync before the data is avaialble.
137 |         await ctx.sync();
138 |         const tOnCopyCmplete = performance.now();
139 |         const standAloneRuns = await executor.benchmarkRuns(ctx, 1, 10);
140 |         const stats = Array.from(standAloneRuns).map(value => value.toFixed(2));
141 |         // report the result.
142 |         const sortedIndex = Array.from(outputData.toArray())
143 |           .map((value, index) => [value, index])
144 |           .sort(([a], [b]) => b - a)
145 |           .map(([, index]) => index);
146 |         clearLog();
147 |         if (typeof inst.lib.webGPUContext.device.defaultQueue.createFence  == "undefined") {
148 |           logger("WARNING: createFence is not supported, timing will be inaccurate.");
149 |         }
150 |         logger(
151 |           "Time cost(ms): end to end=" +
152 |             (tOnCopyCmplete - tStart).toFixed(2) +
153 |             ", GPU run benchmarks=" + stats.join(", ")
154 |         );
155 |         for (let i = 0; i < 5; ++i) {
156 |           logger("Top-" + (i + 1) + " " + synset[sortedIndex[i]]);
157 |         }
158 |       };
159 |     }
160 | 
161 |     async function classify() {
162 |       while (classifier.classify === undefined) {
163 |         await new Promise((r) => setTimeout(r, 10));
164 |       }
165 |       classifier.classify();
166 |     }
167 |     function clearLog() {
168 |       const node = document.getElementById("log");
169 |       while (node.hasChildNodes()) {
170 |         node.removeChild(node.lastChild);
171 |       }
172 |     }
173 | 
174 |     initModel("mobilenet1.0", customLog);
175 |   </script>
176 |   <body>
177 |     <h1>TVM WebGPU Demo</h1>
178 |     <ul>
179 |       <li>
180 |         You need a browser with web GPU capability.
181 |         Checkout instructions <a href="https://github.com/gpuweb/gpuweb/wiki/Implementation-Status">here</a> Tested on Chrome Canary
182 |       </li>
183 |       <li>End to end time cost includes data copy from/to the GPU.</li>
184 |       <li>GPU run benchmarks only invoke compute without data copy.
185 |         We can get a sense of copy overhead by comparing benchmark runs to the end to end timing.
186 |       </li>
187 |       <li>Try to click classify multiple times,
188 |         you will notice that the classifier becomes faster.
189 |         This could due to GPU driver execution stablizes.
190 |         The stablized GPU cost helps us to know what is the best performance
191 |         we can get if we run the model continuously(e.g. in an always on detetcor demo).
192 |       </li>
193 |     </ul>
194 |     <h2>Options</h2>
195 |     Image URL<input
196 |       name="imageURL"
197 |       id="imageURL"
198 |       type="text"
199 |       value="./cat.png"
200 |     /><br /><br />
201 |     <button onclick="classify()">Classify</button>
202 |     <button onclick="clearLog()">Clear Log</button><br />
203 |     <canvas id="canvas" width="224" height="224"></canvas>
204 |     <div id="result"></div>
205 |     <div id="log"></div>
206 |   </body>
207 | </html>
208 | 


--------------------------------------------------------------------------------
/publish_gh_pages.sh:
--------------------------------------------------------------------------------
 1 | #!/bin/bash
 2 | # Publish to gh-pages branch.
 3 | set -e
 4 | set -u
 5 | 
 6 | cleanup()
 7 | {
 8 |     git checkout master
 9 | }
10 | trap cleanup 0
11 | 
12 | git fetch
13 | cp .gitignore .gitignore.bak
14 | git checkout -B gh-pages origin/gh-pages
15 | git ls-files | xargs  rm -f
16 | 
17 | mv .gitignore.bak .gitignore
18 | cp -rf dist/* .
19 | 
20 | DATE=`date`
21 | git add --all && git commit -am "Build at ${DATE}"
22 | git push origin gh-pages
23 | git checkout master
24 | echo "Finish deployment at ${DATE}"
25 | 


--------------------------------------------------------------------------------