Skip to content

Commit

Permalink
wgpu shader update
Browse files Browse the repository at this point in the history
  • Loading branch information
@hana-alice
hana-alice committed Jun 24, 2023
1 parent 5097a82 commit 5cb71b8
Show file tree
Hide file tree
Showing 4 changed files with 149 additions and 73 deletions.
136 changes: 85 additions & 51 deletions cocos/gfx/webgpu/webgpu-define.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
*/

import { WEBGPU } from 'internal:constants';
import { gfx, webgpuAdapter, glslalgWasmModule, promiseForWebGPUInstantiation } from '../../webgpu/instantiated';
import { gfx, webgpuAdapter, glslangWasmModule, promiseForWebGPUInstantiation, spirvOptModule } from '../../webgpu/instantiated';
import {
Texture, CommandBuffer, DescriptorSet, Device, InputAssembler, Buffer, Shader
} from './override';
Expand Down Expand Up @@ -230,73 +230,112 @@ WEBGPU && promiseForWebGPUInstantiation.then(() => {
it = sampReg.exec(code);
}

let funcReg = /\s([\S]+)\(([^)]+)\)[\s].?{/g;
let funcIter = funcReg.exec(code);
const builtinSample = ['texture', 'textureSize', 'texelFetch', 'textureLod'];
const replaceBultin = function (samplerName:string, samplerType: string, target:string) {
builtinSample.forEach((sampleFunc) => {
const builtinSampleReg = new RegExp(`${sampleFunc}\\s*\\(\\s*${samplerName}\\s*,`);
let builtinFuncIter = builtinSampleReg.exec(target);
while(builtinFuncIter) {
target = target.replace(builtinFuncIter[0], `${sampleFunc}(sampler${samplerType}(${samplerName}, ${samplerName}_sampler),`);
builtinFuncIter = builtinSampleReg.exec(target);
}
});
return target;
}

let funcReg = /\s([\S]+)\s*\(([^)]+)\)[\s].?{/g;
let funcIter = funcReg.exec(code);
const funcSet = new Set<string>();
const paramTypeMap = new Map<string, string>();
while(funcIter) {
paramTypeMap.clear();

const params = funcIter[2];
let paramsRes = params.slice();
if (params.includes('sampler')) {
const paramIndexSet = new Set<number>();
let singleParamReg = new RegExp(`,?(\\S+)\\s+\\b([^,)]+)\\b`);
let singleParamReg = new RegExp(`\\W?sampler(\\w+)\\s+\\b([^,)]+)\\b`);
// const paramReg = /[^,]+?sampler(\S+)\s+\b([^,)]+)\b/g;
let paramIter = singleParamReg.exec(params);
let paramIter = singleParamReg.exec(paramsRes);
let index = 0;

while(paramIter) {
if(paramIter[0].includes('sampler')) {
const samplerType = paramIter[1].replace('sampler', '');
const samplerType = paramIter[1];
const paramName = paramIter[2];
const originParamReg = new RegExp(`sampler(${samplerType})\\s+(${paramName})`)
paramsRes = paramsRes.replace(originParamReg,'sampler $2, texture$1 $2_sampler');
paramsRes = paramsRes.replace(singleParamReg,'texture$1 $2, sampler $2_sampler');
paramIndexSet.add(index);
paramTypeMap.set(paramName, samplerType);
}
++index;
paramIter = singleParamReg.exec(params);

paramIter = singleParamReg.exec(paramsRes);
}
code = code.replace(params, paramsRes);

const funcName = funcIter[1];
const funcSampleReg = new RegExp(`${funcName}\\s+?\\(([^,\\)]+)\\)[^(\\s*{)]`, 'g');
let funcSampleIter = funcSampleReg.exec(code);
while (funcSampleIter) {
const funcParams = funcSampleIter[1];
let funcParamsRes = funcParams.slice();
const singleParamReg = new RegExp(`,?\\s+\\b([^,)]+)\\b`);
let index = 0;
let pIter = singleParamReg.exec(funcParams);
while(pIter) {
if(paramIndexSet.has(index)) {
funcParamsRes = funcParamsRes.replace(pIter[1], `${pIter[1]}, ${pIter[1]}_sampler`);
// function may overload
if(!funcSet.has(funcName)) {
const funcSampleReg = new RegExp(`${funcName}\\s*?\\(([,|\\s]?\\w+\\s*?[,|\\)]){1,10}`, 'g');
let funcSampleIter = funcSampleReg.exec(code);
const samplerSet = new Set<string>();
while (funcSampleIter) {
samplerSet.clear();
const func = funcSampleIter[0];
let funcParamsRes = func.slice();
const singleParamReg = new RegExp(/[,|\(|\s]\b(\w+)\b/g);
let index = 0;
// const matches = funcParamsRes.match(singleParamReg)
let pIter = singleParamReg.exec(funcParamsRes);
while(pIter) {
if(paramIndexSet.has(index)) {
samplerSet.add(pIter![1]);
}
pIter = singleParamReg.exec(funcParamsRes);
++index;
}
pIter = singleParamReg.exec(funcParams);
++index;
samplerSet.forEach((samplerParam) => {
funcParamsRes = funcParamsRes.replace(samplerParam, `${samplerParam}, ${samplerParam}_sampler`);
});
code = code.replace(funcSampleIter[0], funcParamsRes);
funcSampleIter = funcSampleReg.exec(code);
}
code = code.replace(funcSampleIter[0], `${funcName}(${funcParamsRes})`);
}
code = code.replace(funcIter[0], `${funcIter[1]}(${paramsRes}) {`);
funcSet.add(funcIter[1]);
}
funcIter = funcReg.exec(code);
}

['texture', 'textureSize', 'texelFetch', 'textureLod'].forEach((sampleFunc) => {
const builtinSampleReg = new RegExp(`${sampleFunc}\\s*\\(\\s*([^,\\)]+)`);
let builtinFuncIter = builtinSampleReg.exec(code);
while(builtinFuncIter) {
const sampleName = builtinFuncIter[1];
let samplerType = '2D';
if(referredMap.has(sampleName)) {
samplerType = referredMap.get(sampleName)!;
} else {
samplerType
let count = 1;
let startIndex = code.indexOf(funcIter[1], funcIter.index);
startIndex = code.indexOf('{', startIndex) + 1;
let endIndex = 0;
while (count) {
if (code.at(startIndex) === '{') {
++count;
} else if (code.at(startIndex) === '}') {
--count;
}

if (count === 0) {
endIndex = startIndex;
break;
}

const nextLeft = code.indexOf('{', startIndex + 1);
const nextRight = code.indexOf('}', startIndex + 1);
startIndex = nextLeft === -1 ? nextRight : Math.min(nextLeft, nextRight);
}
const funcBody = code.slice(funcIter.index, endIndex);
let newFunc = funcBody;
paramTypeMap.forEach((type, name) => {
newFunc = replaceBultin(name, type, newFunc);
});

code = code.replace(builtinFuncIter[0], `${sampleFunc}(sampler${sampleName}`)
code = code.replace(funcBody, newFunc);
funcSet.add(funcIter[1]);
}
});
funcIter = funcReg.exec(code);
}

funcSet.forEach((sampleFunc) => {
const samplerReg = new RegExp(`${sampleFunc}\\s+?\( )`)
referredMap.forEach((type, name) => {
code = replaceBultin(name, type, code);
});

return code;
Expand All @@ -313,15 +352,10 @@ WEBGPU && promiseForWebGPUInstantiation.then(() => {
shaderInfo.stages[i].source = seperateCombinedSamplerTexture(shaderInfo.stages[i].source);
const stageStr = shaderInfo.stages[i].stage === ShaderStageFlagBit.VERTEX ? 'vertex'
: shaderInfo.stages[i].stage === ShaderStageFlagBit.FRAGMENT ? 'fragment' : 'compute';
// const sourceCode = `#version 450\n${shaderInfo.stages[i].source}`;
let sourceCode;
if(i === 0) {
sourceCode = '#version 450\n#define CC_DEVICE_SUPPORT_FLOAT_TEXTURE 1\n#define CC_ENABLE_CLUSTERED_LIGHT_CULLING 0\n#define CC_DEVICE_MAX_VERTEX_UNIFORM_VECTORS 256\n#define CC_DEVICE_MAX_FRAGMENT_UNIFORM_VECTORS 256\n#define CC_DEVICE_CAN_BENEFIT_FROM_INPUT_ATTACHMENT 1\n#define CC_PLATFORM_ANDROID_AND_WEBGL 0\n#define CC_ENABLE_WEBGL_HIGHP_STRUCT_VALUES 0\n#define CC_JOINT_UNIFORM_CAPACITY 60\n#define CC_EFFECT_USED_VERTEX_UNIFORM_VECTORS 56\n#define CC_EFFECT_USED_FRAGMENT_UNIFORM_VECTORS 1\n#define USE_LOCAL 0\n#define SAMPLE_FROM_RT 0\n#define USE_PIXEL_ALIGNMENT 0\n#define CC_USE_EMBEDDED_ALPHA 0\n#define USE_ALPHA_TEST 0\n#define USE_TEXTURE 0\n#define IS_GRAY 0\n\nprecision highp float;\nlayout(set = 0, binding = 0) uniform CCGlobal {\n highp vec4 cc_time;\n mediump vec4 cc_screenSize;\n mediump vec4 cc_nativeSize;\n mediump vec4 cc_probeInfo;\n mediump vec4 cc_debug_view_mode;\n};\nlayout(set = 0, binding = 1) uniform CCCamera {\n highp mat4 cc_matView;\n highp mat4 cc_matViewInv;\n highp mat4 cc_matProj;\n highp mat4 cc_matProjInv;\n highp mat4 cc_matViewProj;\n highp mat4 cc_matViewProjInv;\n highp vec4 cc_cameraPos;\n mediump vec4 cc_surfaceTransform;\n mediump vec4 cc_screenScale;\n mediump vec4 cc_exposure;\n mediump vec4 cc_mainLitDir;\n mediump vec4 cc_mainLitColor;\n mediump vec4 cc_ambientSky;\n mediump vec4 cc_ambientGround;\n mediump vec4 cc_fogColor;\n mediump vec4 cc_fogBase;\n mediump vec4 cc_fogAdd;\n mediump vec4 cc_nearFar;\n mediump vec4 cc_viewPort;\n};\n#if USE_LOCAL\n layout(set = 2, binding = 0) uniform CCLocal {\n highp mat4 cc_matWorld;\n highp mat4 cc_matWorldIT;\n highp vec4 cc_lightingMapUVParam;\n highp vec4 cc_localShadowBias;\n highp vec4 cc_reflectionProbeData1;\n highp vec4 cc_reflectionProbeData2;\n highp vec4 cc_reflectionProbeBlendData1;\n highp vec4 cc_reflectionProbeBlendData2;\n };\n#endif\n#if SAMPLE_FROM_RT\n #define QUATER_PI 0.78539816340\n #define HALF_PI 1.57079632679\n #define PI 3.14159265359\n #define PI2 6.28318530718\n #define PI4 12.5663706144\n #define INV_QUATER_PI 1.27323954474\n #define INV_HALF_PI 0.63661977237\n #define INV_PI 0.31830988618\n #define INV_PI2 0.15915494309\n #define INV_PI4 0.07957747155\n #define EPSILON 1e-6\n #define EPSILON_LOWP 1e-4\n #define LOG2 1.442695\n #define EXP_VALUE 2.71828183f\n #define FP_MAX 65504.0\n #define FP_SCALE 0.0009765625\n #define FP_SCALE_INV 1024.0\n #define GRAY_VECTOR vec3(0.299, 0.587, 0.114)\n #define LIGHT_MAP_TYPE_DISABLED 0\n #define LIGHT_MAP_TYPE_ALL_IN_ONE 1\n #define LIGHT_MAP_TYPE_INDIRECT_OCCLUSION 2\n #define REFLECTION_PROBE_TYPE_NONE 0\n #define REFLECTION_PROBE_TYPE_CUBE 1\n #define REFLECTION_PROBE_TYPE_PLANAR 2\n #define REFLECTION_PROBE_TYPE_BLEND 3\n #define REFLECTION_PROBE_TYPE_BLEND_AND_SKYBOX 4\n #define LIGHT_TYPE_DIRECTIONAL 0.0\n #define LIGHT_TYPE_SPHERE 1.0\n #define LIGHT_TYPE_SPOT 2.0\n #define LIGHT_TYPE_POINT 3.0\n #define LIGHT_TYPE_RANGED_DIRECTIONAL 4.0\n #define IS_DIRECTIONAL_LIGHT(light_type) (abs(float(light_type) - float(LIGHT_TYPE_DIRECTIONAL)) < EPSILON_LOWP)\n #define IS_SPHERE_LIGHT(light_type) (abs(float(light_type) - float(LIGHT_TYPE_SPHERE)) < EPSILON_LOWP)\n #define IS_SPOT_LIGHT(light_type) (abs(float(light_type) - float(LIGHT_TYPE_SPOT)) < EPSILON_LOWP)\n #define IS_POINT_LIGHT(light_type) (abs(float(light_type) - float(LIGHT_TYPE_POINT)) < EPSILON_LOWP)\n #define IS_RANGED_DIRECTIONAL_LIGHT(light_type) (abs(float(light_type) - float(LIGHT_TYPE_RANGED_DIRECTIONAL)) < EPSILON_LOWP)\n #define TONE_MAPPING_ACES 0\n #define TONE_MAPPING_LINEAR 1\n #define SURFACES_MAX_TRANSMIT_DEPTH_VALUE 999999.0\n#endif\nlayout(location = 0) in vec3 a_position;\nlayout(location = 1) in vec2 a_texCoord;\nlayout(location = 2) in vec4 a_color;\nlayout(location = 0) out vec4 color;\nlayout(location = 1) out vec2 uv0;\nvec4 vert () {\n vec4 pos = vec4(a_position, 1);\n #if USE_LOCAL\n pos = cc_matWorld * pos;\n #endif\n #if USE_PIXEL_ALIGNMENT\n pos = cc_matView * pos;\n pos.xyz = floor(pos.xyz);\n pos = cc_matProj * pos;\n #else\n pos = cc_matViewProj * pos;\n #endif\n uv0 = a_texCoord;\n #if SAMPLE_FROM_RT\n uv0 = cc_cameraPos.w > 1.0 ? vec2(uv0.x, 1.0 - uv0.y) : uv0;\n #endif\n color = a_color;\n return pos;\n}\nvoid main() { gl_Position = vert(); }';
} else {
sourceCode = '#version 450\n#define CC_DEVICE_SUPPORT_FLOAT_TEXTURE 1\n#define CC_ENABLE_CLUSTERED_LIGHT_CULLING 0\n#define CC_DEVICE_MAX_VERTEX_UNIFORM_VECTORS 256\n#define CC_DEVICE_MAX_FRAGMENT_UNIFORM_VECTORS 256\n#define CC_DEVICE_CAN_BENEFIT_FROM_INPUT_ATTACHMENT 1\n#define CC_PLATFORM_ANDROID_AND_WEBGL 0\n#define CC_ENABLE_WEBGL_HIGHP_STRUCT_VALUES 0\n#define CC_JOINT_UNIFORM_CAPACITY 60\n#define CC_EFFECT_USED_VERTEX_UNIFORM_VECTORS 56\n#define CC_EFFECT_USED_FRAGMENT_UNIFORM_VECTORS 1\n#define USE_LOCAL 0\n#define SAMPLE_FROM_RT 0\n#define USE_PIXEL_ALIGNMENT 0\n#define CC_USE_EMBEDDED_ALPHA 0\n#define USE_ALPHA_TEST 0\n#define USE_TEXTURE 0\n#define IS_GRAY 0\n\nprecision highp float;\nvec4 CCSampleWithAlphaSeparated(texture2D tex, sampler tex_sampler, vec2 uv) {\n#if CC_USE_EMBEDDED_ALPHA\n return vec4(texture(sampler2D(tex, tex_sampler), uv).rgb, texture(sampler2D(tex, tex_sampler), uv + vec2(0.0, 0.5)).r);\n#else\n return texture(sampler2D(tex, tex_sampler), uv);\n#endif\n}\n#if USE_ALPHA_TEST\n layout(set = 1, binding = 0) uniform ALPHA_TEST_DATA {\n float alphaThreshold;\n };\n#endif\nvoid ALPHA_TEST (in vec4 color) {\n #if USE_ALPHA_TEST\n if (color.a < alphaThreshold) discard;\n #endif\n}\nvoid ALPHA_TEST (in float alpha) {\n #if USE_ALPHA_TEST\n if (alpha < alphaThreshold) discard;\n #endif\n}\nlayout(location = 0) in vec4 color;\n#if USE_TEXTURE\n layout(location = 1) in vec2 uv0;\nlayout(set = 2, binding = 12 + 16) uniform texture2D _cc_spriteTexture;\nlayout(set = 2, binding = 12) uniform sampler _cc_spriteTexture_sampler;\n\n#endif\nvec4 frag () {\n vec4 o = vec4(1, 1, 1, 1);\n #if USE_TEXTURE\n o *= CCSampleWithAlphaSeparated(sampler2D(_cc_spriteTexture, _cc_spriteTexture_sampler), uv0);\n #if IS_GRAY\n float gray = 0.2126 * o.r + 0.7152 * o.g + 0.0722 * o.b;\n o.r = o.g = o.b = gray;\n #endif\n #endif\n o *= color;\n ALPHA_TEST(o);\n return o;\n}\nlayout(location = 0) out vec4 cc_FragColor;\nvoid main() { cc_FragColor = frag(); }';
}
const spv = glslalgWasmModule.glslang.compileGLSL(sourceCode, stageStr, true, '1.3');
spvDatas.push(spv);
const sourceCode = `#version 450\n${shaderInfo.stages[i].source}`;
const spv = glslangWasmModule.glslang.compileGLSL(sourceCode, stageStr, false, '1.3');
const spvOpted = spirvOptModule.spv.opt(spirvOptModule.spv.opt.SPV_ENV_WEBGPU_0, spv);
spvDatas.push(spvOpted);
}

const shader = this.createShaderNative(shaderInfo, spvDatas);
Expand Down
18 changes: 16 additions & 2 deletions cocos/webgpu/instantiated.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,15 +30,22 @@
import { WASM_SUPPORT_MODE, WEBGPU } from 'internal:constants';
import webgpuUrl from 'external:emscripten/webgpu/webgpu_wasm.wasm';
import glslangUrl from 'external:emscripten/webgpu/glslang.wasm';
import spirvOptUrl from 'external:emscripten/webgpu/spirv-tools.wasm';
import wasmDevice from 'external:emscripten/webgpu/webgpu_wasm.js';
import glslangLoader from 'external:emscripten/webgpu/glslang.js';
import spirvOptLoader from 'external:emscripten/webgpu/spirv-tools.js';
import { legacyCC } from '../core/global-exports';
import { WebAssemblySupportMode } from '../misc/webassembly-support';
import spirvTools from '../../native/external/emscripten/webgpu/spirv-tools';

export const glslalgWasmModule: any = {
export const glslangWasmModule: any = {
glslang: null,
};

export const spirvOptModule: any = {
spv: null,
};

export const gfx: any = legacyCC.gfx = {
wasmBinary: null,
nativeDevice: null,
Expand All @@ -54,7 +61,14 @@ export const promiseForWebGPUInstantiation = (() => {
// TODO: we need to support AsmJS fallback option
return Promise.all([
glslangLoader(new URL(glslangUrl, import.meta.url).href).then((res) => {
glslalgWasmModule.glslang = res;
glslangWasmModule.glslang = res;
}),
new Promise<void>((resolve) => {
spirvOptLoader(new URL(spirvOptUrl, import.meta.url).href).then((spv) => {
// spirvOptModule.spv = await spirvTools();
spirvOptModule.spv = spv;
resolve();
});
}),
new Promise<void>((resolve) => {
fetch(new URL(webgpuUrl, import.meta.url).href).then((response) => {
Expand Down
Loading

0 comments on commit 5cb71b8

Please sign in to comment.