[RFC] Taichi on Javascript

[AmesingFlank]

I'm writing this RFC to share my thoughts and plans to build a version of Taichi completely embedded in Javascript, with the hope of getting production-level applications to use Taichi, and thereby increasing the impact of Taichi as a programming language.

Currently, Taichi uses Python as the host language. Thanks to Python's simplicity, Taichi enables programmers with little GPU experiences to have fun and experiment with GPU computing. Javascript shares the simplicity of Python (to some extent), and additionally, it has the advantage of being used extensively in production-level, user-facing applications. For this reason, if we offer Taichi as a (deployable) Javascript package, it would provide professional web developers with a brand new way of building and serving compute-heavy or graphics-heavy web programs. This would certainly boost the popularity of Taichi.

To be more specifically, I'm planning to build a Javascript library that enables this:

const ti = require("taichi")

x = ti.Vector.field(3, ti.f32, 512)

let k = ti.kernel(() => {
    for (i of ti.range(512)) {
        x[i] = x[i] + 1
    }
})

k()

Here, taichi would be a normal Javascript module, one that could be installed simply by running npm install taichi. Most importantly, it must be a pure Javascript module, meaning that it must have no binary code dependencies attached to it. This is of course vital if we want web developers to be able to use it in their own products. We will of course implement interop routines with normal Javascript arrays, HTML canvases, and maybe even provide a GGUI.

Use cases and Impacts

From the top of my head I can think of a lot of use cool cases for a Javascript version of Taichi

• Graphics applications on the web. No doubt there are lots and lots of graphics demands on web pages, and these are done most often via WebGL, which is flexible but cumbersome. As a language directly embedded in JS, this could be very appealing for beginner-level front-end engineers who don’t want to learn WebGL.

• Compute heavy applications on the web These are considerably rarer, but part of the reason for that is exactly because it’s difficult to develop high performance programs on the web -- and taichi just might change that.

• Other fancy apps:

  • A (much better) taichi zoo, where actually everything happens on the frontend, and the backend server does nothing but serving (instead of compiling).
  • A compute resource sharing platform: This is something I thought about a few years ago. We could build a site where people can submit compute tasks written in taichi, and other people could simple open up the website and contribute their (GPU) computing resources.
  • A transpiler between python-taichi and JS-taichi, which helps improve the impact of the python language.

Technological Considerations and Challenges

Obviously, this will be a huge project with many challenges. Here’re the things that have crossed my mind:

• Feasibility on the Javascript side In Javascript, you can obtain the source code of a function passed in as an argument. So this shouldn’t be a problem. In fact, there is a library called gpu.js[https://gpu.rocks/#/] which already uses this to translate Javascript to GPU code. That package is pretty cool, but Taichi is much, much more powerful of course. Also, Javascript might support decorators in the future. If this ECMAScript proposal goes through, we might be able to have

  
  const ti = require("taichi")

x = ti.Vector.field(3, ti.f32, 512)

@ti.kernel function k() { for (i of ti.range(512)) { x[i] = x[i] + 1 } }

k()

* No binary code dependencies?
Emscripten. I have already successfully compiled the entire Taichi project using emscripten. I haven't started trying to use the generated code, and there will probably be some troubles to overcome, but overall this is very feasible.

* GPU Computing
This is a big "X" of this project. WebGL doesn’t have good support for compute. WebGPU does, but browsers don’t have good support for WebGPU. If we only use CPU, e.g. WASM, then the performance probably won’t be enough of a selling point.
Although, I remain hopeful that not too long in the future, WebGPU will become widely adopted. I mean, WebGL can’t just live forever right? 
<p align="center">
  <img  width="300" src="https://user-images.githubusercontent.com/16066115/145727484-9c6693f7-a271-48da-bfaf-13786142557c.png" />
</p>

# Roadmap
* [ ] Javascript Frontend. ( i.e. JS source code -> CHI IR)
I believe the best way to do this is to use the [Typescript Compiler API](https://github.com/microsoft/TypeScript/wiki/Using-the-Compiler-API).
* [ ] Fully emscripten Taichi
After the Javascript frontend, we will use a emscriptened Taichi to translate CHI IR to SPIR-V
* [ ] SPIR-V to WGSL.
There are a few publicly available tools that does this, namely [naga](https://github.com/gfx-rs/naga) and [tint](https://dawn.googlesource.com/tint/). But in fact none of them works well enough and there're some painful problems. I won't go into the details here, instead, I will post a comment below to discuss this.
* [ ] A WebGPU runtime that ingests the WGSL.

# Proof of Concept

I have hacked up a WebGPU Taichi runtime and used Taichi AOT and Tint to generate some WGSL kernels. If you have Chrome 94+, you can see it here 

https://amesingflank.github.io/taichi.js/

---EDIT: 2022.Feb.13: 
I built a taichi.js zoo with an interactive code editor and a few examples ported from Python:
https://amesingflank.github.io/taichi.js/

---EDIT: 2022.April.21
Moved website to https://taichi-js.com

[AmesingFlank]

Regarding SPIR-V to WGSL, the biggest pain point is atomics. To understand the problem, please have a look at this discussion.

TLDR:

• In WGSL, atomic operations can only be applied to atomic<T> types. And there is no way to convert a T* to a atomic<T>*.
• If we wish to perform any atomic operations at all on a buffer, we would need to promote the entire buffer into atomic types, and always use atomic operations (even load/store). This not only inconvenient, but also inefficient.
• Neither tint nor naga handles this. They immediately give up when seeing an atomic.

I'd appreciate any help/suggestions on this, as this feels like a huge blocker. To make matters worse, it looks that the WebGPU committee has no plans to make atomics more flexible in the near future, as they have moved the discussion to post-v1.

[bobcao3]

Regarding SPIR-V to WGSL, the biggest pain point is atomics. To understand the problem, please have a look at this discussion.

TLDR:

* In WGSL, atomic operations can only be applied to `atomic<T>` types. And there is no way to convert a `T*` to a `atomic<T>*`.

* If we wish to perform any atomic operations at all on a buffer, we would need to promote the entire buffer into atomic types, and always use atomic operations (even load/store). This not only inconvenient, but also inefficient.

* Neither tint or naga handles this. They immediately give up when seeing an atomic.

I'd appreciate any help/suggestions on this, as this feels like a huge blocker. To make matters worse, it looks that the WebGPU committee has no plans to make atomics more flexible in the near future, as they have moved the discussion to post-v1.

WebGPU is not a stable API by any means and the atomic issue also originates from their persistence to use WGSL instead of a common standard. This does not sound like an addressable issue in any way, we might want to straight up limit atomic capabilities on WGSL all together. This is my personal opinion, but I feel like if the platform did not allow flexibility, we should not try to "fix" it by hacking our ways around it

[bobcao3]

Would it be reasonable to make this work two-parted. The first part is the front-end, and we can use bindings such as https://github.com/maierfelix/nvk to allow at least node.js front-end to work with Vulkan. The second part would be the backend where we figure out WebGPU and WGSL

[AmesingFlank]

This does not sound like an addressable issue in any way, we might want to straight up limit atomic capabilities on WGSL all together. This is my personal opinion, but I feel like if the platform did not allow flexibility, we should not try to "fix" it by hacking our ways around it

I more or less agree. I hate the idea of hacking around it either. But this would essentially mean we would support no atomics at all, and that's not great. Virtually all of the physical simulation programs in taichi requires atomics. Not having atomics is a huge limitation.

Would it be reasonable to make this work two-parted. The first part is the front-end, and we can use bindings such as https://github.com/maierfelix/nvk to allow at least node.js front-end to work with Vulkan. The second part would be the backend where we figure out WebGPU and WGSL

That makes sense. Although I wouldn't invest too much time into a node-only version if it turns out to be a lot of work. My reasoning being that a node-only version doesn't add that much value to our existing Python version, while a version that runs in browsers can open up a lot of possibilities.

[eatcosmos]

https://amesingflank.github.io/taichi.js/ compile and run not work?

[AmesingFlank]

https://amesingflank.github.io/taichi.js/ compile and run not work?

Hi, thanks for reporting. Could you please tell me your browser version? And could you paste the code that you put in the code editor? Are there any error messages in the Logs?

[eatcosmos]

https://amesingflank.github.io/taichi.js/ compile and run not work?

Hi, thanks for reporting. Could you please tell me your browser version? And could you paste the code that you put in the code editor? Are there any error messages in the Logs?

96.0.4664.45（） （32 ）

19 error(s) generated while compiling the shader:
27:28 error: float literals must not be suffixed with 'f'
  let tmp1 : f32 = 1.000000f;
                           ^

31:28 error: float literals must not be suffixed with 'f'
  let tmp5 : f32 = 0.150000f;
                           ^

32:28 error: float literals must not be suffixed with 'f'
  let tmp6 : f32 = 0.450000f;
                           ^

33:28 error: float literals must not be suffixed with 'f'
  let tmp7 : f32 = 0.500000f;
                           ^
......

[AmesingFlank]

@eatcosmos Thanks! I will try to fix this.

[AmesingFlank]

@eatcosmos Also, could you try upgrading to chrome v98?

[eatcosmos]

@eatcosmos Also, could you try upgrading to chrome v98?

v98 is ok, load image quickly

[eatcosmos]

What about webgl Oasis Engines, is three any relation between Taichi.js and Oasis or Three.js in technology？ Oasis Engine - Mobile first web graphic engine 支付宝App 3D动画和小游戏背后的故事_前端_曾柏然_InfoQ精选文章

[bobcao3]

What about webgl Oasis Engines, is three any relation between Taichi.js and Oasis or Three.js in technology？ Oasis Engine - Mobile first web graphic engine 支付宝App 3D动画和小游戏背后的故事_前端_曾柏然_InfoQ精选文章

There isn't much in common at all. Taichi.js is a compiler codegen and JIT runtime system that compiles your JS code into GPU kernels. These examples are graphics and rendering engines, they are very different things. So, no relation in them at all.

[eatcosmos]

What about webgl Oasis Engines, is three any relation between Taichi.js and Oasis or Three.js in technology？ Oasis Engine - Mobile first web graphic engine 支付宝App 3D动画和小游戏背后的故事_前端_曾柏然_InfoQ精选文章

There isn't much in common at all. Taichi.js is a compiler codegen and JIT runtime system that compiles your JS code into GPU kernels. These examples are graphics and rendering engines, they are very different things. So, no relation in them at all.

run taichi.js demo in mobile chrome	run three.js demo in mobile chrome	run oasis.js demo in mobile chrome

---

Test taichi.js three.js oasis.js, taichi.js cannot show demo in mobile chrome, what need to do to let taichi.js run directly in mobile chrome? Is this feature is taichi.js's objective in the future?

[AmesingFlank]

@eatcosmos Taichi.js runs on WebGPU, which is not supported by mobile chrome yet.

Of course I hope that in the future, mobile chrome/safari will support WebGPU and thus Taichi.js. But at this point there's no telling how long that would take.

[eatcosmos]

@eatcosmos Taichi.js runs on WebGPU, which is not supported by mobile chrome yet.

Of course I hope that in the future, mobile chrome/safari will support WebGPU and thus Taichi.js. But at this point there's no telling how long that would take.

but mobile has no gpu, three.js could work, three.js use webgl , taichi.js not use webgl? Then will taichi.js do the webgl's work , so taichi.js can work in mobile chrome without using webgl?

[AmesingFlank]

@eatcosmos Three.js uses WebGL, which is a legacy API, and is well-supported by mobile browsers. Taichi.js uses a new API called WebGPU, which is aimed to replace WebGL, but has a long way to go from actually doing so.

[k-ye]

FYI you can track Chrome's support status of WebGPU in https://chromestatus.com/feature/6213121689518080. It doesn't mention mobile Chrome yet, but this is almost a guaranteed thing.

[eatcosmos]

what is the process of producing taichi.js？ taichi.js is converted or rewrite from taichi-dev c++ source code？ or converted to wasm then convert wasm to taichi.js source code？

[AmesingFlank]

We use emscripten to compile taichi's C++ code base into javascript, paired with a WebGPU runtime written in Typescript.

[eatcosmos]

thank you for your explanation，maybe webgl is also a good choice for now to cover most devices

[AmesingFlank]

The problem with WebGL is that it doesn't support compute shaders. It is a legacy API and we expect that in due time, WebGPU will receive fairly common support.