歡迎來到 WebGPU 的世界

WebGPU 是一門神奇的技術,在瀏覽器支持率 0%,標準還沒有定稿的情況下,就已經被 Three.js 和 Babylon.js 等主流 3D 和遊戲框架支持了。而且被 Tensorflow.js 用來加速手機端的深度學習,比起 WebGL 能帶來 20~30 倍的顯著提升。

在主流框架中 WebGPU 的例子

1、在 Three.js 中使用 WebGPU

使用 Three.js 的封裝,我們可以直接生成 WebGPU 的調用。

我們照貓畫虎引入 WebGPU 相關的庫:

   import * as THREE from 'three';
   import * as Nodes from 'three-nodes/Nodes.js';
   import { add, mul } from 'three-nodes/ShaderNode.js';

            import WebGPU from './jsm/capabilities/WebGPU.js';
   import WebGPURenderer from './jsm/renderers/webgpu/WebGPURenderer.js';
...

剩下就跟普通的 WebGL 代碼寫起來差不多:

   async function init() {

    if ( WebGPU.isAvailable() === false ) {
     document.body.appendChild( WebGPU.getErrorMessage() );
     throw new Error( 'No WebGPU support' );
    }

    const container = document.createElement( 'div' );
    document.body.appendChild( container );

    camera = new THREE.PerspectiveCamera( 45, window.innerWidth / window.innerHeight, 1, 4000 );
    camera.position.set( 0, 200, 1200 );

    scene = new THREE.Scene();
...

只不過渲染器使用 WebGPURenderer:

    renderer = new WebGPURenderer();
    renderer.setPixelRatio( window.devicePixelRatio );
    renderer.setSize( window.innerWidth, window.innerHeight );
    container.appendChild( renderer.domElement );
...

如果封裝的不能滿足需求了,我們可以使用 WGSL 語言進行擴展:

    material = new Nodes.MeshBasicNodeMaterial();
    material.colorNode = desaturateWGSLNode.call( { color: new Nodes.TextureNode( texture ) } );
    materials.push( material );

    const getWGSLTextureSample = new Nodes.FunctionNode( `
     fn getWGSLTextureSample( tex: texture_2d<f32>, tex_sampler: sampler, uv:vec2<f32> ) -> vec4<f32> {
      return textureSample( tex, tex_sampler, uv ) * vec4<f32>( 0.0, 1.0, 0.0, 1.0 );
     }
    ` );

    const textureNode = new Nodes.TextureNode( texture );

    material = new Nodes.MeshBasicNodeMaterial();
    material.colorNode = getWGSLTextureSample.call( { tex: textureNode, tex_sampler: textureNode, uv: new Nodes.UVNode() } );
    materials.push( material );

WGSL 是 WebGPU 進行 GPU 指令編程的語言。類似於 OpenGL 的 GLSL, Direct3D 的 HLSL。

我們來看一個完整的例子,顯示一個跳舞的小人,也不過 100 多行代碼:

<!DOCTYPE html>
<html lang="en">
 <head>
  <title>three.js - WebGPU - Skinning</title>
  <meta charset="utf-8">
  <meta >
  <link type="text/css" rel="stylesheet" href="main.css">
  <meta http-equiv="origin-trial" content="AoS1pSJwCV3KRe73TO0YgJkK9FZ/qhmvKeafztp0ofiE8uoGrnKzfxGVKKICvoBfL8dgE0zpkp2g/oEJNS0fDgkAAABeeyJvcmlnaW4iOiJodHRwczovL3RocmVlanMub3JnOjQ0MyIsImZlYXR1cmUiOiJXZWJHUFUiLCJleHBpcnkiOjE2NTI4MzE5OTksImlzU3ViZG9tYWluIjp0cnVlfQ==">
 </head>
 <body>
  <div id="info">
   <a href="https://threejs.org" target="_blank" rel="noopener">three.js</a> WebGPU - Skinning
  </div>
  <script async src="https://unpkg.com/es-module-shims@1.3.6/dist/es-module-shims.js"></script>
  <script type="importmap">
   {
    "imports"{
     "three""../build/three.module.js",
     "three-nodes/""./jsm/nodes/"
    }
   }
  </script>

  <script type="module">
   import * as THREE from 'three';
   import * as Nodes from 'three-nodes/Nodes.js';
   import { FBXLoader } from './jsm/loaders/FBXLoader.js';
   import WebGPU from './jsm/capabilities/WebGPU.js';
   import WebGPURenderer from './jsm/renderers/webgpu/WebGPURenderer.js';
   import LightsNode from 'three-nodes/lights/LightsNode.js';

   let camera, scene, renderer;
   let mixer, clock;
   init().then( animate ).catch( error );

   async function init() {
    if ( WebGPU.isAvailable() === false ) {
     document.body.appendChild( WebGPU.getErrorMessage() );
     throw new Error( 'No WebGPU support' );
    }
    camera = new THREE.PerspectiveCamera( 50, window.innerWidth / window.innerHeight, 1, 1000 );
    camera.position.set( 100, 200, 300 );
    scene = new THREE.Scene();
    camera.lookAt( 0, 100, 0 );
    clock = new THREE.Clock();

    // 光照
    const light = new THREE.PointLight( 0xffffff );
    camera.add( light );
    scene.add( camera );
    const lightNode = new LightsNode().fromLights( [ light ] );
    const loader = new FBXLoader();
    loader.load( 'models/fbx/Samba Dancing.fbx'function ( object ) {
     mixer = new THREE.AnimationMixer( object );
     const action = mixer.clipAction( object.animations[ 0 ] );
     action.play();
     object.traverse( function ( child ) {
      if ( child.isMesh ) {
       child.material = new Nodes.MeshStandardNodeMaterial();
       child.material.lightNode = lightNode;
      }
     } );
     scene.add( object );
    } );

    // 渲染
    renderer = new WebGPURenderer();
    renderer.setPixelRatio( window.devicePixelRatio );
    renderer.setSize( window.innerWidth, window.innerHeight );
    document.body.appendChild( renderer.domElement );
    window.addEventListener( 'resize', onWindowResize );
    return renderer.init();
   }

   function onWindowResize() {
    camera.aspect = window.innerWidth / window.innerHeight;
    camera.updateProjectionMatrix();
    renderer.setSize( window.innerWidth, window.innerHeight );
   }

   function animate() {
    requestAnimationFrame( animate );
    const delta = clock.getDelta();
    if ( mixer ) mixer.update( delta );
    renderer.render( scene, camera );
   }

   function error( error ) {
    console.error( error );
   }
  </script>
 </body>
</html>

2、在 Babylon.js 中使用 WebGPU

Babylon.js 的封裝與 Three.js 大同小異,我們來看個 PlayGround 的效果:

不同之處在於處理 WebGPU 的支持情況時,Babylon.js 並不判斷整體上支不支持 WebGPU,而是隻看具體功能。

比如上面的例子,只判斷是不是支持計算着色器。

    const supportCS = engine.getCaps().supportComputeShaders;

不過目前在 macOS 上,只有 WebGPU 支持計算着色器。

如果我們把環境切換成 WebGL2,就變成下面這樣了:

順便說一句,Babylon.js 判斷 WebGL2 和 WebGL 時也是同樣的邏輯,有高就用高。

如果對於着色器不熟悉,Babylon.js 提供了練習 Vertex Shader 和 Pixel Shader 的環境:https://cyos.babylonjs.com/ , 帶語法高亮和預覽。

針對需要通過寫手機應用的場景,Babylon.js 提供了與 React Native 結合的能力:

3、用 WebGPU 進行深度學習加速

除了 3D 界面和遊戲,深度學習的推理器也是 GPU 的重度用戶。所以 Tensorflow.js 也在還落不了地的時候就支持了 WebGPU。實在是計算着色器太重要了。

寫出來的加速代碼就像下面一樣,很多算子的實現最終是由 WGSL 代碼來實現的,最終會轉換成 GPU 的指令。

  getUserCode(): string {
    const rank = this.xShape.length;
    const type = getCoordsDataType(rank);
    const start = this.xShape.map((_, i) =`uniforms.pad${i}[0]`).join(',');
    const end = this.xShape
                    .map(
                        (_, i) =`uniforms.pad${i}[0] + uniforms.xShape${
                            rank > 1 ? `[${i}]` : ''}`)
                    .join(',');
    const startValue = rank > 1 ? `${type}(${start})` : `${start}`;
    const endValue = rank > 1 ? `${type}(${end})` : `${end}`;

    const leftPadCondition = rank > 1 ? `any(outC < start)` : `outC < start`;
    const rightPadCondition = rank > 1 ? `any(outC >= end)` : `outC >= end`;

    const unpackedCoords = rank > 1 ?
        ['coords[0]''coords[1]''coords[2]''coords[3]'].slice(0, rank) :
        'coords';

    const userCode = `
      ${getMainHeaderAndGlobalIndexString()}
        if (index < uniforms.size) {
          let start = ${startValue};
          let end = ${endValue};
          let outC = getCoordsFromIndex(index);
          if (${leftPadCondition} || ${rightPadCondition}) {
            setOutputAtIndex(index, uniforms.constantValue);
          } else {
            let coords = outC - start;
            setOutputAtIndex(index, getX(${unpackedCoords}));
          }
        }
      }
    `;
    return userCode;
  }

無框架手寫 WebGPU 代碼

通過框架,我們可以迅速地跟上技術的前沿。但是,框架的封裝也容易讓我們迷失對於技術本質的把握。

現在我們來看看如何手寫 WebGPU 代碼。

1、從 Canvas 說起

不管是 WebGL 還是 WebGPU,都是對於 Canvas 的擴展。做爲 HTML 5 的重要新增功能,大家對於 2D 的 Canvas 應該都不陌生。

比如我們要畫一個三角形,就可以調用 lineTo API 來實現:

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>Canvas</title>
</head>
<body>
    <canvas id="webcanvas" width="200" height="200" style="background-color: #eee"></canvas>
    <script>
        const canvas=document.getElementById('webcanvas');
        const ctx=canvas.getContext('2d');

        ctx.beginPath();
        ctx.moveTo(75,50);
        ctx.lineTo(100,75);
        ctx.lineTo(100,25);
        ctx.fill();
    </script>
</body>

畫出來的結果如下:

我們要修改畫出來的圖的顏色怎麼辦?
ctx 有 fillStyle 屬性,支持 CSS 的顏色字符串。

比如我們設成紅色,可以這麼寫:

ctx.fillStyle = 'red';

也可以這麼寫:

ctx.fillStyle = '#F00';

還可以這麼寫:

ctx.fillStyle = 'rgb(255,0,0,1)';

2、從 2D 到 3D

從 2D Canvas 到 3D WebGL 的最大跨越,就是從調用 API,到完全不同於 JavaScript 的新語言 GLSL 的出場。

第一步的步子我們邁得小一點,不畫三角形了,只畫一個點。

<!DOCTYPE html>
<html lang="en">

<head>
    <meta charset="UTF-8">
    <title>Test OpenGL for a point</title>
</head>

<body>
    <canvas id="webgl" width="500" height="500" style="background-color: blue"></canvas>
    <script>
        const canvas = document.getElementById('webgl');
        const gl = canvas.getContext('webgl');

        const program = gl.createProgram();

        const vertexShaderSource = `
           void main(){
              gl_PointSize=sqrt(20.0);
              gl_Position =vec4(0.0,0.0,0.0,1.0);
           }`;

        const vertexShader = gl.createShader(gl.VERTEX_SHADER);
        gl.shaderSource(vertexShader, vertexShaderSource);
        gl.compileShader(vertexShader);
        gl.attachShader(program, vertexShader);

        const fragShaderSource = `
          void main(){
            gl_FragColor = vec4(1.0,0.0,0.0,1.0);
          }
        `;

        const fragmentShader = gl.createShader(gl.FRAGMENT_SHADER);
        gl.shaderSource(fragmentShader, fragShaderSource);
        gl.compileShader(fragmentShader);
        gl.attachShader(program, fragmentShader);

        gl.linkProgram(program);
        gl.useProgram(program);

        gl.drawArrays(gl.POINTS, 0, 1);

    </script>

</body>

</html>

getContext 時將 2d 換成 webgl。
我們可以加一行console.log(gl)來看下 gl 是什麼東西:

我們可以看到,它是一個 WebGLRenderingContext 對象。
順便說一句,之前我們拿到的 2D 的 Context 是 CanvasRenderingContext2D。

下面就引入了兩段程序中的程序,第一段叫做頂點着色器,用於頂點的座標信息。第二段叫做片元着色器,用於配置如何進行一些屬性的操作,在本例中我們做一個最基本的操作,改顏色。

我們先看頂點着色器的代碼:

           void main(){
              gl_PointSize=sqrt(20.0);
              gl_Position =vec4(0.0,0.0,0.0,1.0);
           }

像其他語言一樣,glsl 中的代碼也需要一個入口函數。

gl_PointSize 是一個系統變量,用於存儲點的大小。我特意給大小加個了 sqrt 函數,給大家展示 glsl 的庫函數。

gl_Position 用於存儲起點的位置。vec4 是由 4 個元素構成的向量。

GLSL 的數據類型很豐富,包括標量、向量、數組、矩陣、結構體和採樣器等。

標量有布爾型 bool, 有符號整數 int, 無符號整數 uint 和浮點數 float 4 種類型。

類型的使用方式跟 C 語言一樣,比如我們用 float 來定義浮點變量。

                float pointSize = sqrt(20.0);
                gl_PointSize=pointSize;

GLSL 沒有 double 這樣表示雙精度的類型。在頂點着色器中是沒有精度設置的。
但是在片元着色器中有精度的設置,需要指定低精度 lowp, 中精度 mediump 和高精度 highp. 一般採用中精度:

            void main(){    
                mediump vec4 pointColor;
                pointColor.r = 1.0;
                pointColor.a = 1.0;
                gl_FragColor = pointColor;
            }

GLSL 因爲是基於 C 語言設計的,不支持泛型,所以每種向量同時有 4 種子類型的。

以四元組 vec4 爲例,有 4 種類型:

另外還有 vec2, vec3 各有 4 種子類型,以此類推。

在 GLSL 裏面,四元向量最常用的用途有兩種,在頂點着色器裏充當座標,和在片元着色器裏充當顏色。

當 vec4 作爲座標使用時,我們可以用 x,y,z,w 屬性來對應 4 個維度。

我們來看個例子:

                vec4 pos;
                pos.x = 0.0;
                pos.y = 0.0;
                pos.z = 0.0;
                pos.w = 1.0;
                gl_Position = pos;

同樣,我們在片元着色器裏面表示紅色的時候只用指令 r 和 a 兩個屬性,g,b 讓它們默認是 0:

            void main(){    
                mediump vec4 pointColor;
                pointColor.r = 1.0;
                pointColor.a = 1.0;
                gl_FragColor = pointColor;
            }

有了頂點着色器和片元着色器的 GLSL 代碼之後,我們將其進行編程,並 attach 到 program 上面。

最後再 link 和 use 這個 program,就可以調用 drawArrays 來進行繪製了。

3、更現代的 GPU 編程方法

跨越了從 Canvas API 到 GLSL 的鴻溝了之後,最後到 WebGPU 這一步相對就容易一些了。

我們要熟悉的是以 Vulkan 爲代表的更現代的 GPU 的編程方法。

渲染管線不再是唯一,我們可以使用更通用的計算管線了。也不再有頂點着色器和片元着色器那麼嚴格的限制。

另外最重要的一點是,爲了提升 GPU 執行效率,WebGPU 不再是像 WebGL 一樣基本每一步都要由 CPU 來控制,我們使用 commandEncoder 將所有 GPU 指令打包在一起,一次性執行。

我們先看一下完整代碼有個印象:

<!DOCTYPE html>
<html lang="en">

<head>
    <meta charset="UTF-8">
    <title>Test WebGPU</title>
</head>

<body>
    <canvas id="webgpu" width="500" height="500" style="background-color: blue"></canvas>
    <script>
        async function testGPU() {
            const canvas = document.getElementById('webgpu');
            const gpuContext = canvas.getContext('webgpu');

            const adapter = await navigator.gpu.requestAdapter();
            const device = await adapter.requestDevice();

            presentationFormat = gpuContext.getPreferredFormat(adapter);

            gpuContext.configure({
                device,
                format: presentationFormat
            });

            const triangleVertWGSL = `
            @stage(vertex)
            fn main(@builtin(vertex_index) VertexIndex : u32)
             -> @builtin(position) vec4<f32> {
                var pos = array<vec2<f32>, 3>(
                vec2<f32>(0.0, 0.5),
                vec2<f32>(-0.5, -0.5),
                vec2<f32>(0.5, -0.5));

                return vec4<f32>(pos[VertexIndex], 0.0, 1.0);
            }
            `;

            const redFragWGSL = `
            @stage(fragment)
                fn main() -> @location(0) vec4<f32> {
                return vec4<f32>(1.0, 0.0, 0.0, 1.0);
            }
            `

            const commandEncoder = device.createCommandEncoder();
            const textureView = gpuContext.getCurrentTexture().createView();

            const pipeline = device.createRenderPipeline({
                vertex: {
                    module: device.createShaderModule({
                        code: triangleVertWGSL,
                    }),
                    entryPoint: 'main',
                },
                fragment: {
                    module: device.createShaderModule({
                        code: redFragWGSL,
                    }),
                    entryPoint: 'main',
                    targets: [
                        {
                            format: presentationFormat,
                        },
                    ],
                },
                primitive: {
                    topology: 'triangle-list',
                },
            });

            const renderPassDescriptor = {
                colorAttachments: [
                    {
                        view: textureView,
                        loadValue: { r: 1.0, g: 1.0, b: 1.0, a: 1.0 }, 
                        storeOp: 'store',
                    },
                ],
            };
            const passEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
            console.log(passEncoder);
            passEncoder.setPipeline(pipeline);
            passEncoder.draw(3, 1, 0, 0);
            passEncoder.end();

            device.queue.submit([commandEncoder.finish()]);
        }

        testGPU();

    </script>

</body>

</html>

因爲瀏覽器還沒有支持,所以我們需要像 Chrome Canary 這樣的支持最新技術的瀏覽器。而且還要打開支持的開關,比如在 Chrome Canary 裏是 enable-unsafe-webgpu.

三角形畫出來的結果如下:

現在的 Context 從 WebGL 的 WebGLRenderingContext 變成了 GPUCanvasContext。

WGSL 語言的語法更像 Rust,vec4 這樣的容器可以用泛型的寫法綁定類型:

            @stage(vertex)
            fn main(@builtin(vertex_index) VertexIndex : u32)
             -> @builtin(position) vec4<f32> {
                var pos = array<vec2<f32>, 3>(
                vec2<f32>(0.0, 0.5),
                vec2<f32>(-0.5, -0.5),
                vec2<f32>(0.5, -0.5));

                return vec4<f32>(pos[VertexIndex], 0.0, 1.0);
            }

對比下 Rust 的代碼看看像不像:

fn fib2(n: i32) -> i64 {
    if n <= 2 {
        return 1i64
    } else {
        return fib2(n - 1) + fib2(n - 2)
    }
}

WGSL 是爲了規避知識產權問題發明的新語言,本質上它和 GLSL,HLSL 等語言一樣,都可以編譯成 Vulkan 的 SPIR-V 二進制格式:
.

Vulkan 不限制使用什麼樣的語言,既可以使用 GLSL, HLSL,也可以使用 Open CL 或者是 Open CL 的高級封裝 SYCL。

轉換成 SPIR-V 格式之後,可以轉成 iOS 上的 Metal Shading Language,也可以轉成 Windows Direct 12 上用的 DXIL。

WebGPU 沒有這麼自由,發明了一門新語言 WGSL,不過其思想都是基於 SPIR-V 的。

在 WebGPU 和 WGSL 還未定版,資料還比較缺乏的情況下,我們可以先學習 Vulkan 相關的知識,然後遷移到 WebGPU 上來。本質上是同樣的東西,只是封裝略有不同。

我們之前學習的 GLSL 的知識同樣用得上,而且在這種類 Rust 風格中可以寫得更爽一些。

比如同樣是給片元用的顏色值,在保留了 vec4 可以繼續使用 r,g,b,a 分量的好處之外,因爲指定了 f32 的精度,就不需要 mediump 了。而且,類型可以自動推斷,我們直接給個 var 就好了:

            @stage(fragment)
                fn main() -> @location(0) vec4<f32> {
                var triColor = vec4<f32>(0.0,0.0,0.0,0.0);
                triColor.r = 1.0;
                triColor.a = 1.0;    
                return triColor;
            }

有了作爲功能核心的 WGSL,剩下的工作主要就是組裝了。

我們把指令打包在 CommandEncoder 中,然後通過 beginRenderPass 來創建一個渲染 Pass,再給這個 Pass 設置一個渲染的流水線,添加相應的 draw 操作,最後提交到 GPU 設備的隊列中,就大功告成了。

小結

相對於基於 OpenGL ES 2.0 的 WebGL 1.0,WebGPU 更接近於 Vulkan 這樣更能發揮 GPU 能力的新 API,可以更有效地發揮出新的 GPU 的能力。就像渲染上 Three.js 和 Babylon.js 給我們展示的那樣和計算上 Tensorflow.js 的飛躍一樣。

雖然瀏覽器還不支持,但是不成熟的主要是封裝,底層的 Vulkan 和 Metal 技術已經非常成熟,並且廣泛被客戶端所使用了。

WebGPU 這個能力暴露給 H5 和小程序之後,將給元宇宙等熱門應用插上性能倍增的翅膀。結合 WebXR 等支持率更成問題的新技術一起,成爲未來幾年前端的主要工具。

Alibaba F2E 阿里巴巴前端官方公衆號

本文由 Readfog 進行 AMP 轉碼,版權歸原作者所有。
來源https://mp.weixin.qq.com/s/SZvqpGJKGNeQfIa9fcp8XA