WGSLNodeBuilder.js 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212
  1. import { NoColorSpace, FloatType } from 'three';
  2. import NodeUniformsGroup from '../../common/nodes/NodeUniformsGroup.js';
  3. import NodeSampler from '../../common/nodes/NodeSampler.js';
  4. import { NodeSampledTexture, NodeSampledCubeTexture, NodeSampledTexture3D } from '../../common/nodes/NodeSampledTexture.js';
  5. import NodeUniformBuffer from '../../common/nodes/NodeUniformBuffer.js';
  6. import NodeStorageBuffer from '../../common/nodes/NodeStorageBuffer.js';
  7. import { NodeBuilder, CodeNode } from '../../../nodes/Nodes.js';
  8. import { getFormat } from '../utils/WebGPUTextureUtils.js';
  9. import WGSLNodeParser from './WGSLNodeParser.js';
  10. import { GPUBufferBindingType, GPUStorageTextureAccess } from '../utils/WebGPUConstants.js';
  11. // GPUShaderStage is not defined in browsers not supporting WebGPU
  12. const GPUShaderStage = self.GPUShaderStage;
  13. const gpuShaderStageLib = {
  14. 'vertex': GPUShaderStage ? GPUShaderStage.VERTEX : 1,
  15. 'fragment': GPUShaderStage ? GPUShaderStage.FRAGMENT : 2,
  16. 'compute': GPUShaderStage ? GPUShaderStage.COMPUTE : 4
  17. };
  18. const supports = {
  19. instance: true,
  20. swizzleAssign: false,
  21. storageBuffer: true
  22. };
  23. const wgslFnOpLib = {
  24. '^^': 'threejs_xor'
  25. };
  26. const wgslTypeLib = {
  27. float: 'f32',
  28. int: 'i32',
  29. uint: 'u32',
  30. bool: 'bool',
  31. color: 'vec3<f32>',
  32. vec2: 'vec2<f32>',
  33. ivec2: 'vec2<i32>',
  34. uvec2: 'vec2<u32>',
  35. bvec2: 'vec2<bool>',
  36. vec3: 'vec3<f32>',
  37. ivec3: 'vec3<i32>',
  38. uvec3: 'vec3<u32>',
  39. bvec3: 'vec3<bool>',
  40. vec4: 'vec4<f32>',
  41. ivec4: 'vec4<i32>',
  42. uvec4: 'vec4<u32>',
  43. bvec4: 'vec4<bool>',
  44. mat2: 'mat2x2<f32>',
  45. imat2: 'mat2x2<i32>',
  46. umat2: 'mat2x2<u32>',
  47. bmat2: 'mat2x2<bool>',
  48. mat3: 'mat3x3<f32>',
  49. imat3: 'mat3x3<i32>',
  50. umat3: 'mat3x3<u32>',
  51. bmat3: 'mat3x3<bool>',
  52. mat4: 'mat4x4<f32>',
  53. imat4: 'mat4x4<i32>',
  54. umat4: 'mat4x4<u32>',
  55. bmat4: 'mat4x4<bool>'
  56. };
  57. const wgslMethods = {
  58. dFdx: 'dpdx',
  59. dFdy: '- dpdy',
  60. mod_float: 'threejs_mod_float',
  61. mod_vec2: 'threejs_mod_vec2',
  62. mod_vec3: 'threejs_mod_vec3',
  63. mod_vec4: 'threejs_mod_vec4',
  64. equals_bool: 'threejs_equals_bool',
  65. equals_bvec2: 'threejs_equals_bvec2',
  66. equals_bvec3: 'threejs_equals_bvec3',
  67. equals_bvec4: 'threejs_equals_bvec4',
  68. lessThanEqual: 'threejs_lessThanEqual',
  69. greaterThan: 'threejs_greaterThan',
  70. inversesqrt: 'inverseSqrt',
  71. bitcast: 'bitcast<f32>'
  72. };
  73. const wgslPolyfill = {
  74. threejs_xor: new CodeNode( `
  75. fn threejs_xor( a : bool, b : bool ) -> bool {
  76. return ( a || b ) && !( a && b );
  77. }
  78. ` ),
  79. lessThanEqual: new CodeNode( `
  80. fn threejs_lessThanEqual( a : vec3<f32>, b : vec3<f32> ) -> vec3<bool> {
  81. return vec3<bool>( a.x <= b.x, a.y <= b.y, a.z <= b.z );
  82. }
  83. ` ),
  84. greaterThan: new CodeNode( `
  85. fn threejs_greaterThan( a : vec3<f32>, b : vec3<f32> ) -> vec3<bool> {
  86. return vec3<bool>( a.x > b.x, a.y > b.y, a.z > b.z );
  87. }
  88. ` ),
  89. mod_float: new CodeNode( 'fn threejs_mod_float( x : f32, y : f32 ) -> f32 { return x - y * floor( x / y ); }' ),
  90. mod_vec2: new CodeNode( 'fn threejs_mod_vec2( x : vec2f, y : vec2f ) -> vec2f { return x - y * floor( x / y ); }' ),
  91. mod_vec3: new CodeNode( 'fn threejs_mod_vec3( x : vec3f, y : vec3f ) -> vec3f { return x - y * floor( x / y ); }' ),
  92. mod_vec4: new CodeNode( 'fn threejs_mod_vec4( x : vec4f, y : vec4f ) -> vec4f { return x - y * floor( x / y ); }' ),
  93. equals_bool: new CodeNode( 'fn threejs_equals_bool( a : bool, b : bool ) -> bool { return a == b; }' ),
  94. equals_bvec2: new CodeNode( 'fn threejs_equals_bvec2( a : vec2f, b : vec2f ) -> vec2<bool> { return vec2<bool>( a.x == b.x, a.y == b.y ); }' ),
  95. equals_bvec3: new CodeNode( 'fn threejs_equals_bvec3( a : vec3f, b : vec3f ) -> vec3<bool> { return vec3<bool>( a.x == b.x, a.y == b.y, a.z == b.z ); }' ),
  96. equals_bvec4: new CodeNode( 'fn threejs_equals_bvec4( a : vec4f, b : vec4f ) -> vec4<bool> { return vec4<bool>( a.x == b.x, a.y == b.y, a.z == b.z, a.w == b.w ); }' ),
  97. repeatWrapping: new CodeNode( `
  98. fn threejs_repeatWrapping( uv : vec2<f32>, dimension : vec2<u32> ) -> vec2<u32> {
  99. let uvScaled = vec2<u32>( uv * vec2<f32>( dimension ) );
  100. return ( ( uvScaled % dimension ) + dimension ) % dimension;
  101. }
  102. ` ),
  103. biquadraticTexture: new CodeNode( `
  104. fn threejs_biquadraticTexture( map : texture_2d<f32>, coord : vec2f, level : i32 ) -> vec4f {
  105. let res = vec2f( textureDimensions( map, level ) );
  106. let uvScaled = coord * res;
  107. let uvWrapping = ( ( uvScaled % res ) + res ) % res;
  108. // https://www.shadertoy.com/view/WtyXRy
  109. let uv = uvWrapping - 0.5;
  110. let iuv = floor( uv );
  111. let f = fract( uv );
  112. let rg1 = textureLoad( map, vec2i( iuv + vec2( 0.5, 0.5 ) ), level );
  113. let rg2 = textureLoad( map, vec2i( iuv + vec2( 1.5, 0.5 ) ), level );
  114. let rg3 = textureLoad( map, vec2i( iuv + vec2( 0.5, 1.5 ) ), level );
  115. let rg4 = textureLoad( map, vec2i( iuv + vec2( 1.5, 1.5 ) ), level );
  116. return mix( mix( rg1, rg2, f.x ), mix( rg3, rg4, f.x ), f.y );
  117. }
  118. ` )
  119. };
  120. class WGSLNodeBuilder extends NodeBuilder {
  121. constructor( object, renderer ) {
  122. super( object, renderer, new WGSLNodeParser() );
  123. this.uniformGroups = {};
  124. this.builtins = {};
  125. }
  126. needsColorSpaceToLinear( texture ) {
  127. return texture.isVideoTexture === true && texture.colorSpace !== NoColorSpace;
  128. }
  129. _generateTextureSample( texture, textureProperty, uvSnippet, depthSnippet, shaderStage = this.shaderStage ) {
  130. if ( shaderStage === 'fragment' ) {
  131. if ( depthSnippet ) {
  132. return `textureSample( ${ textureProperty }, ${ textureProperty }_sampler, ${ uvSnippet }, ${ depthSnippet } )`;
  133. } else {
  134. return `textureSample( ${ textureProperty }, ${ textureProperty }_sampler, ${ uvSnippet } )`;
  135. }
  136. } else if ( this.isFilteredTexture( texture ) ) {
  137. return this.generateFilteredTexture( texture, textureProperty, uvSnippet );
  138. } else {
  139. return this.generateTextureLod( texture, textureProperty, uvSnippet, '0' );
  140. }
  141. }
  142. _generateVideoSample( textureProperty, uvSnippet, shaderStage = this.shaderStage ) {
  143. if ( shaderStage === 'fragment' ) {
  144. return `textureSampleBaseClampToEdge( ${ textureProperty }, ${ textureProperty }_sampler, vec2<f32>( ${ uvSnippet }.x, 1.0 - ${ uvSnippet }.y ) )`;
  145. } else {
  146. console.error( `WebGPURenderer: THREE.VideoTexture does not support ${ shaderStage } shader.` );
  147. }
  148. }
  149. _generateTextureSampleLevel( texture, textureProperty, uvSnippet, levelSnippet, depthSnippet, shaderStage = this.shaderStage ) {
  150. if ( shaderStage === 'fragment' && this.isUnfilterable( texture ) === false ) {
  151. return `textureSampleLevel( ${ textureProperty }, ${ textureProperty }_sampler, ${ uvSnippet }, ${ levelSnippet } )`;
  152. } else if ( this.isFilteredTexture( texture ) ) {
  153. return this.generateFilteredTexture( texture, textureProperty, uvSnippet, levelSnippet );
  154. } else {
  155. return this.generateTextureLod( texture, textureProperty, uvSnippet, levelSnippet );
  156. }
  157. }
  158. generateFilteredTexture( texture, textureProperty, uvSnippet, levelSnippet = '0' ) {
  159. this._include( 'biquadraticTexture' );
  160. return `threejs_biquadraticTexture( ${ textureProperty }, ${ uvSnippet }, i32( ${ levelSnippet } ) )`;
  161. }
  162. generateTextureLod( texture, textureProperty, uvSnippet, levelSnippet = '0' ) {
  163. this._include( 'repeatWrapping' );
  164. const dimension = `textureDimensions( ${ textureProperty }, 0 )`;
  165. return `textureLoad( ${ textureProperty }, threejs_repeatWrapping( ${ uvSnippet }, ${ dimension } ), i32( ${ levelSnippet } ) )`;
  166. }
  167. generateTextureLoad( texture, textureProperty, uvIndexSnippet, depthSnippet, levelSnippet = '0u' ) {
  168. if ( depthSnippet ) {
  169. return `textureLoad( ${ textureProperty }, ${ uvIndexSnippet }, ${ depthSnippet }, ${ levelSnippet } )`;
  170. } else {
  171. return `textureLoad( ${ textureProperty }, ${ uvIndexSnippet }, ${ levelSnippet } )`;
  172. }
  173. }
  174. generateTextureStore( texture, textureProperty, uvIndexSnippet, valueSnippet ) {
  175. return `textureStore( ${ textureProperty }, ${ uvIndexSnippet }, ${ valueSnippet } )`;
  176. }
  177. isUnfilterable( texture ) {
  178. return this.getComponentTypeFromTexture( texture ) !== 'float' || ( texture.isDataTexture === true && texture.type === FloatType );
  179. }
  180. generateTexture( texture, textureProperty, uvSnippet, depthSnippet, shaderStage = this.shaderStage ) {
  181. let snippet = null;
  182. if ( texture.isVideoTexture === true ) {
  183. snippet = this._generateVideoSample( textureProperty, uvSnippet, shaderStage );
  184. } else if ( this.isUnfilterable( texture ) ) {
  185. snippet = this.generateTextureLod( texture, textureProperty, uvSnippet, '0', depthSnippet, shaderStage );
  186. } else {
  187. snippet = this._generateTextureSample( texture, textureProperty, uvSnippet, depthSnippet, shaderStage );
  188. }
  189. return snippet;
  190. }
  191. generateTextureGrad( texture, textureProperty, uvSnippet, gradSnippet, depthSnippet, shaderStage = this.shaderStage ) {
  192. if ( shaderStage === 'fragment' ) {
  193. // TODO handle i32 or u32 --> uvSnippet, array_index: A, ddx, ddy
  194. return `textureSampleGrad( ${ textureProperty }, ${ textureProperty }_sampler, ${ uvSnippet }, ${ gradSnippet[ 0 ] }, ${ gradSnippet[ 1 ] } )`;
  195. } else {
  196. console.error( `WebGPURenderer: THREE.TextureNode.gradient() does not support ${ shaderStage } shader.` );
  197. }
  198. }
  199. generateTextureCompare( texture, textureProperty, uvSnippet, compareSnippet, depthSnippet, shaderStage = this.shaderStage ) {
  200. if ( shaderStage === 'fragment' ) {
  201. return `textureSampleCompare( ${ textureProperty }, ${ textureProperty }_sampler, ${ uvSnippet }, ${ compareSnippet } )`;
  202. } else {
  203. console.error( `WebGPURenderer: THREE.DepthTexture.compareFunction() does not support ${ shaderStage } shader.` );
  204. }
  205. }
  206. generateTextureLevel( texture, textureProperty, uvSnippet, levelSnippet, depthSnippet, shaderStage = this.shaderStage ) {
  207. let snippet = null;
  208. if ( texture.isVideoTexture === true ) {
  209. snippet = this._generateVideoSample( textureProperty, uvSnippet, shaderStage );
  210. } else {
  211. snippet = this._generateTextureSampleLevel( texture, textureProperty, uvSnippet, levelSnippet, depthSnippet, shaderStage );
  212. }
  213. return snippet;
  214. }
  215. getPropertyName( node, shaderStage = this.shaderStage ) {
  216. if ( node.isNodeVarying === true && node.needsInterpolation === true ) {
  217. if ( shaderStage === 'vertex' ) {
  218. return `varyings.${ node.name }`;
  219. }
  220. } else if ( node.isNodeUniform === true ) {
  221. const name = node.name;
  222. const type = node.type;
  223. if ( type === 'texture' || type === 'cubeTexture' || type === 'storageTexture' || type === 'texture3D' ) {
  224. return name;
  225. } else if ( type === 'buffer' || type === 'storageBuffer' ) {
  226. return `NodeBuffer_${ node.id }.${name}`;
  227. } else {
  228. return node.groupNode.name + '.' + name;
  229. }
  230. }
  231. return super.getPropertyName( node );
  232. }
  233. getOutputStructName() {
  234. return 'output';
  235. }
  236. _getUniformGroupCount( shaderStage ) {
  237. return Object.keys( this.uniforms[ shaderStage ] ).length;
  238. }
  239. getFunctionOperator( op ) {
  240. const fnOp = wgslFnOpLib[ op ];
  241. if ( fnOp !== undefined ) {
  242. this._include( fnOp );
  243. return fnOp;
  244. }
  245. return null;
  246. }
  247. getStorageAccess( node ) {
  248. if ( node.isStorageTextureNode ) {
  249. switch ( node.access ) {
  250. case GPUStorageTextureAccess.ReadOnly:
  251. return 'read';
  252. case GPUStorageTextureAccess.WriteOnly:
  253. return 'write';
  254. default:
  255. return 'read_write';
  256. }
  257. } else {
  258. switch ( node.access ) {
  259. case GPUBufferBindingType.Storage:
  260. return 'read_write';
  261. case GPUBufferBindingType.ReadOnlyStorage:
  262. return 'read';
  263. default:
  264. return 'write';
  265. }
  266. }
  267. }
  268. getUniformFromNode( node, type, shaderStage, name = null ) {
  269. const uniformNode = super.getUniformFromNode( node, type, shaderStage, name );
  270. const nodeData = this.getDataFromNode( node, shaderStage, this.globalCache );
  271. if ( nodeData.uniformGPU === undefined ) {
  272. let uniformGPU;
  273. const group = node.groupNode;
  274. const groupName = group.name;
  275. const bindings = this.getBindGroupArray( groupName, shaderStage );
  276. if ( type === 'texture' || type === 'cubeTexture' || type === 'storageTexture' || type === 'texture3D' ) {
  277. let texture = null;
  278. if ( type === 'texture' || type === 'storageTexture' ) {
  279. texture = new NodeSampledTexture( uniformNode.name, uniformNode.node, group, node.access ? node.access : null );
  280. } else if ( type === 'cubeTexture' ) {
  281. texture = new NodeSampledCubeTexture( uniformNode.name, uniformNode.node, group, node.access ? node.access : null );
  282. } else if ( type === 'texture3D' ) {
  283. texture = new NodeSampledTexture3D( uniformNode.name, uniformNode.node, group, node.access ? node.access : null );
  284. }
  285. texture.store = node.isStorageTextureNode === true;
  286. texture.setVisibility( gpuShaderStageLib[ shaderStage ] );
  287. if ( shaderStage === 'fragment' && this.isUnfilterable( node.value ) === false && texture.store === false ) {
  288. const sampler = new NodeSampler( `${uniformNode.name}_sampler`, uniformNode.node, group );
  289. sampler.setVisibility( gpuShaderStageLib[ shaderStage ] );
  290. bindings.push( sampler, texture );
  291. uniformGPU = [ sampler, texture ];
  292. } else {
  293. bindings.push( texture );
  294. uniformGPU = [ texture ];
  295. }
  296. } else if ( type === 'buffer' || type === 'storageBuffer' ) {
  297. const bufferClass = type === 'storageBuffer' ? NodeStorageBuffer : NodeUniformBuffer;
  298. const buffer = new bufferClass( node, group );
  299. buffer.setVisibility( gpuShaderStageLib[ shaderStage ] );
  300. bindings.push( buffer );
  301. uniformGPU = buffer;
  302. } else {
  303. const uniformsStage = this.uniformGroups[ shaderStage ] || ( this.uniformGroups[ shaderStage ] = {} );
  304. let uniformsGroup = uniformsStage[ groupName ];
  305. if ( uniformsGroup === undefined ) {
  306. uniformsGroup = new NodeUniformsGroup( groupName, group );
  307. uniformsGroup.setVisibility( gpuShaderStageLib[ shaderStage ] );
  308. uniformsStage[ groupName ] = uniformsGroup;
  309. bindings.push( uniformsGroup );
  310. }
  311. uniformGPU = this.getNodeUniform( uniformNode, type );
  312. uniformsGroup.addUniform( uniformGPU );
  313. }
  314. nodeData.uniformGPU = uniformGPU;
  315. }
  316. return uniformNode;
  317. }
  318. getBuiltin( name, property, type, shaderStage = this.shaderStage ) {
  319. const map = this.builtins[ shaderStage ] || ( this.builtins[ shaderStage ] = new Map() );
  320. if ( map.has( name ) === false ) {
  321. map.set( name, {
  322. name,
  323. property,
  324. type
  325. } );
  326. }
  327. return property;
  328. }
  329. getVertexIndex() {
  330. if ( this.shaderStage === 'vertex' ) {
  331. return this.getBuiltin( 'vertex_index', 'vertexIndex', 'u32', 'attribute' );
  332. }
  333. return 'vertexIndex';
  334. }
  335. buildFunctionCode( shaderNode ) {
  336. const layout = shaderNode.layout;
  337. const flowData = this.flowShaderNode( shaderNode );
  338. const parameters = [];
  339. for ( const input of layout.inputs ) {
  340. parameters.push( input.name + ' : ' + this.getType( input.type ) );
  341. }
  342. //
  343. const code = `fn ${ layout.name }( ${ parameters.join( ', ' ) } ) -> ${ this.getType( layout.type ) } {
  344. ${ flowData.vars }
  345. ${ flowData.code }
  346. return ${ flowData.result };
  347. }`;
  348. //
  349. return code;
  350. }
  351. getInstanceIndex() {
  352. if ( this.shaderStage === 'vertex' ) {
  353. return this.getBuiltin( 'instance_index', 'instanceIndex', 'u32', 'attribute' );
  354. }
  355. return 'instanceIndex';
  356. }
  357. getDrawIndex() {
  358. return null;
  359. }
  360. getFrontFacing() {
  361. return this.getBuiltin( 'front_facing', 'isFront', 'bool' );
  362. }
  363. getFragCoord() {
  364. return this.getBuiltin( 'position', 'fragCoord', 'vec4<f32>' ) + '.xyz';
  365. }
  366. getFragDepth() {
  367. return 'output.' + this.getBuiltin( 'frag_depth', 'depth', 'f32', 'output' );
  368. }
  369. isFlipY() {
  370. return false;
  371. }
  372. getBuiltins( shaderStage ) {
  373. const snippets = [];
  374. const builtins = this.builtins[ shaderStage ];
  375. if ( builtins !== undefined ) {
  376. for ( const { name, property, type } of builtins.values() ) {
  377. snippets.push( `@builtin( ${name} ) ${property} : ${type}` );
  378. }
  379. }
  380. return snippets.join( ',\n\t' );
  381. }
  382. getAttributes( shaderStage ) {
  383. const snippets = [];
  384. if ( shaderStage === 'compute' ) {
  385. this.getBuiltin( 'global_invocation_id', 'id', 'vec3<u32>', 'attribute' );
  386. }
  387. if ( shaderStage === 'vertex' || shaderStage === 'compute' ) {
  388. const builtins = this.getBuiltins( 'attribute' );
  389. if ( builtins ) snippets.push( builtins );
  390. const attributes = this.getAttributesArray();
  391. for ( let index = 0, length = attributes.length; index < length; index ++ ) {
  392. const attribute = attributes[ index ];
  393. const name = attribute.name;
  394. const type = this.getType( attribute.type );
  395. snippets.push( `@location( ${index} ) ${ name } : ${ type }` );
  396. }
  397. }
  398. return snippets.join( ',\n\t' );
  399. }
  400. getStructMembers( struct ) {
  401. const snippets = [];
  402. const members = struct.getMemberTypes();
  403. for ( let i = 0; i < members.length; i ++ ) {
  404. const member = members[ i ];
  405. snippets.push( `\t@location( ${i} ) m${i} : ${ member }<f32>` );
  406. }
  407. const builtins = this.getBuiltins( 'output' );
  408. if ( builtins ) snippets.push( builtins );
  409. return snippets.join( ',\n' );
  410. }
  411. getStructs( shaderStage ) {
  412. const snippets = [];
  413. const structs = this.structs[ shaderStage ];
  414. for ( let index = 0, length = structs.length; index < length; index ++ ) {
  415. const struct = structs[ index ];
  416. const name = struct.name;
  417. let snippet = `\struct ${ name } {\n`;
  418. snippet += this.getStructMembers( struct );
  419. snippet += '\n}';
  420. snippets.push( snippet );
  421. snippets.push( `\nvar<private> output : ${ name };\n\n` );
  422. }
  423. return snippets.join( '\n\n' );
  424. }
  425. getVar( type, name ) {
  426. return `var ${ name } : ${ this.getType( type ) }`;
  427. }
  428. getVars( shaderStage ) {
  429. const snippets = [];
  430. const vars = this.vars[ shaderStage ];
  431. if ( vars !== undefined ) {
  432. for ( const variable of vars ) {
  433. snippets.push( `\t${ this.getVar( variable.type, variable.name ) };` );
  434. }
  435. }
  436. return `\n${ snippets.join( '\n' ) }\n`;
  437. }
  438. getVaryings( shaderStage ) {
  439. const snippets = [];
  440. if ( shaderStage === 'vertex' ) {
  441. this.getBuiltin( 'position', 'Vertex', 'vec4<f32>', 'vertex' );
  442. }
  443. if ( shaderStage === 'vertex' || shaderStage === 'fragment' ) {
  444. const varyings = this.varyings;
  445. const vars = this.vars[ shaderStage ];
  446. for ( let index = 0; index < varyings.length; index ++ ) {
  447. const varying = varyings[ index ];
  448. if ( varying.needsInterpolation ) {
  449. let attributesSnippet = `@location( ${index} )`;
  450. if ( /^(int|uint|ivec|uvec)/.test( varying.type ) ) {
  451. attributesSnippet += ' @interpolate( flat )';
  452. }
  453. snippets.push( `${ attributesSnippet } ${ varying.name } : ${ this.getType( varying.type ) }` );
  454. } else if ( shaderStage === 'vertex' && vars.includes( varying ) === false ) {
  455. vars.push( varying );
  456. }
  457. }
  458. }
  459. const builtins = this.getBuiltins( shaderStage );
  460. if ( builtins ) snippets.push( builtins );
  461. const code = snippets.join( ',\n\t' );
  462. return shaderStage === 'vertex' ? this._getWGSLStruct( 'VaryingsStruct', '\t' + code ) : code;
  463. }
  464. getUniforms( shaderStage ) {
  465. const uniforms = this.uniforms[ shaderStage ];
  466. const bindingSnippets = [];
  467. const bufferSnippets = [];
  468. const structSnippets = [];
  469. const uniformGroups = {};
  470. for ( const uniform of uniforms ) {
  471. const groundName = uniform.groupNode.name;
  472. const uniformIndexes = this.bindingsIndexes[ groundName ];
  473. if ( uniform.type === 'texture' || uniform.type === 'cubeTexture' || uniform.type === 'storageTexture' || uniform.type === 'texture3D' ) {
  474. const texture = uniform.node.value;
  475. if ( shaderStage === 'fragment' && this.isUnfilterable( texture ) === false && uniform.node.isStorageTextureNode !== true ) {
  476. if ( texture.isDepthTexture === true && texture.compareFunction !== null ) {
  477. bindingSnippets.push( `@binding( ${ uniformIndexes.binding ++ } ) @group( ${ uniformIndexes.group } ) var ${ uniform.name }_sampler : sampler_comparison;` );
  478. } else {
  479. bindingSnippets.push( `@binding( ${ uniformIndexes.binding ++ } ) @group( ${ uniformIndexes.group } ) var ${ uniform.name }_sampler : sampler;` );
  480. }
  481. }
  482. let textureType;
  483. if ( texture.isCubeTexture === true ) {
  484. textureType = 'texture_cube<f32>';
  485. } else if ( texture.isDataArrayTexture === true ) {
  486. textureType = 'texture_2d_array<f32>';
  487. } else if ( texture.isDepthTexture === true ) {
  488. textureType = 'texture_depth_2d';
  489. } else if ( texture.isVideoTexture === true ) {
  490. textureType = 'texture_external';
  491. } else if ( texture.isData3DTexture === true ) {
  492. textureType = 'texture_3d<f32>';
  493. } else if ( uniform.node.isStorageTextureNode === true ) {
  494. const format = getFormat( texture );
  495. const access = this.getStorageAccess( uniform.node );
  496. textureType = `texture_storage_2d<${ format }, ${ access }>`;
  497. } else {
  498. const componentPrefix = this.getComponentTypeFromTexture( texture ).charAt( 0 );
  499. textureType = `texture_2d<${ componentPrefix }32>`;
  500. }
  501. bindingSnippets.push( `@binding( ${ uniformIndexes.binding ++ } ) @group( ${ uniformIndexes.group } ) var ${ uniform.name } : ${ textureType };` );
  502. } else if ( uniform.type === 'buffer' || uniform.type === 'storageBuffer' ) {
  503. const bufferNode = uniform.node;
  504. const bufferType = this.getType( bufferNode.bufferType );
  505. const bufferCount = bufferNode.bufferCount;
  506. const bufferCountSnippet = bufferCount > 0 ? ', ' + bufferCount : '';
  507. const bufferSnippet = `\t${ uniform.name } : array< ${ bufferType }${ bufferCountSnippet } >\n`;
  508. const bufferAccessMode = bufferNode.isStorageBufferNode ? `storage, ${ this.getStorageAccess( bufferNode ) }` : 'uniform';
  509. bufferSnippets.push( this._getWGSLStructBinding( 'NodeBuffer_' + bufferNode.id, bufferSnippet, bufferAccessMode, uniformIndexes.binding ++, uniformIndexes.group ) );
  510. } else {
  511. const vectorType = this.getType( this.getVectorType( uniform.type ) );
  512. const groupName = uniform.groupNode.name;
  513. const group = uniformGroups[ groupName ] || ( uniformGroups[ groupName ] = {
  514. index: uniformIndexes.binding ++,
  515. id: uniformIndexes.group,
  516. snippets: []
  517. } );
  518. group.snippets.push( `\t${ uniform.name } : ${ vectorType }` );
  519. }
  520. }
  521. for ( const name in uniformGroups ) {
  522. const group = uniformGroups[ name ];
  523. structSnippets.push( this._getWGSLStructBinding( name, group.snippets.join( ',\n' ), 'uniform', group.index, group.id ) );
  524. }
  525. let code = bindingSnippets.join( '\n' );
  526. code += bufferSnippets.join( '\n' );
  527. code += structSnippets.join( '\n' );
  528. return code;
  529. }
  530. buildCode() {
  531. const shadersData = this.material !== null ? { fragment: {}, vertex: {} } : { compute: {} };
  532. for ( const shaderStage in shadersData ) {
  533. const stageData = shadersData[ shaderStage ];
  534. stageData.uniforms = this.getUniforms( shaderStage );
  535. stageData.attributes = this.getAttributes( shaderStage );
  536. stageData.varyings = this.getVaryings( shaderStage );
  537. stageData.structs = this.getStructs( shaderStage );
  538. stageData.vars = this.getVars( shaderStage );
  539. stageData.codes = this.getCodes( shaderStage );
  540. //
  541. let flow = '// code\n\n';
  542. flow += this.flowCode[ shaderStage ];
  543. const flowNodes = this.flowNodes[ shaderStage ];
  544. const mainNode = flowNodes[ flowNodes.length - 1 ];
  545. const outputNode = mainNode.outputNode;
  546. const isOutputStruct = ( outputNode !== undefined && outputNode.isOutputStructNode === true );
  547. for ( const node of flowNodes ) {
  548. const flowSlotData = this.getFlowData( node/*, shaderStage*/ );
  549. const slotName = node.name;
  550. if ( slotName ) {
  551. if ( flow.length > 0 ) flow += '\n';
  552. flow += `\t// flow -> ${ slotName }\n\t`;
  553. }
  554. flow += `${ flowSlotData.code }\n\t`;
  555. if ( node === mainNode && shaderStage !== 'compute' ) {
  556. flow += '// result\n\n\t';
  557. if ( shaderStage === 'vertex' ) {
  558. flow += `varyings.Vertex = ${ flowSlotData.result };`;
  559. } else if ( shaderStage === 'fragment' ) {
  560. if ( isOutputStruct ) {
  561. stageData.returnType = outputNode.nodeType;
  562. flow += `return ${ flowSlotData.result };`;
  563. } else {
  564. let structSnippet = '\t@location(0) color: vec4<f32>';
  565. const builtins = this.getBuiltins( 'output' );
  566. if ( builtins ) structSnippet += ',\n\t' + builtins;
  567. stageData.returnType = 'OutputStruct';
  568. stageData.structs += this._getWGSLStruct( 'OutputStruct', structSnippet );
  569. stageData.structs += '\nvar<private> output : OutputStruct;\n\n';
  570. flow += `output.color = ${ flowSlotData.result };\n\n\treturn output;`;
  571. }
  572. }
  573. }
  574. }
  575. stageData.flow = flow;
  576. }
  577. if ( this.material !== null ) {
  578. this.vertexShader = this._getWGSLVertexCode( shadersData.vertex );
  579. this.fragmentShader = this._getWGSLFragmentCode( shadersData.fragment );
  580. } else {
  581. this.computeShader = this._getWGSLComputeCode( shadersData.compute, ( this.object.workgroupSize || [ 64 ] ).join( ', ' ) );
  582. }
  583. }
  584. getMethod( method, output = null ) {
  585. let wgslMethod;
  586. if ( output !== null ) {
  587. wgslMethod = this._getWGSLMethod( method + '_' + output );
  588. }
  589. if ( wgslMethod === undefined ) {
  590. wgslMethod = this._getWGSLMethod( method );
  591. }
  592. return wgslMethod || method;
  593. }
  594. getType( type ) {
  595. return wgslTypeLib[ type ] || type;
  596. }
  597. isAvailable( name ) {
  598. let result = supports[ name ];
  599. if ( result === undefined ) {
  600. if ( name === 'float32Filterable' ) {
  601. result = this.renderer.hasFeature( 'float32-filterable' );
  602. }
  603. supports[ name ] = result;
  604. }
  605. return result;
  606. }
  607. _getWGSLMethod( method ) {
  608. if ( wgslPolyfill[ method ] !== undefined ) {
  609. this._include( method );
  610. }
  611. return wgslMethods[ method ];
  612. }
  613. _include( name ) {
  614. const codeNode = wgslPolyfill[ name ];
  615. codeNode.build( this );
  616. if ( this.currentFunctionNode !== null ) {
  617. this.currentFunctionNode.includes.push( codeNode );
  618. }
  619. return codeNode;
  620. }
  621. _getWGSLVertexCode( shaderData ) {
  622. return `${ this.getSignature() }
  623. // uniforms
  624. ${shaderData.uniforms}
  625. // varyings
  626. ${shaderData.varyings}
  627. var<private> varyings : VaryingsStruct;
  628. // codes
  629. ${shaderData.codes}
  630. @vertex
  631. fn main( ${shaderData.attributes} ) -> VaryingsStruct {
  632. // vars
  633. ${shaderData.vars}
  634. // flow
  635. ${shaderData.flow}
  636. return varyings;
  637. }
  638. `;
  639. }
  640. _getWGSLFragmentCode( shaderData ) {
  641. return `${ this.getSignature() }
  642. // uniforms
  643. ${shaderData.uniforms}
  644. // structs
  645. ${shaderData.structs}
  646. // codes
  647. ${shaderData.codes}
  648. @fragment
  649. fn main( ${shaderData.varyings} ) -> ${shaderData.returnType} {
  650. // vars
  651. ${shaderData.vars}
  652. // flow
  653. ${shaderData.flow}
  654. }
  655. `;
  656. }
  657. _getWGSLComputeCode( shaderData, workgroupSize ) {
  658. return `${ this.getSignature() }
  659. // system
  660. var<private> instanceIndex : u32;
  661. // uniforms
  662. ${shaderData.uniforms}
  663. // codes
  664. ${shaderData.codes}
  665. @compute @workgroup_size( ${workgroupSize} )
  666. fn main( ${shaderData.attributes} ) {
  667. // system
  668. instanceIndex = id.x;
  669. // vars
  670. ${shaderData.vars}
  671. // flow
  672. ${shaderData.flow}
  673. }
  674. `;
  675. }
  676. _getWGSLStruct( name, vars ) {
  677. return `
  678. struct ${name} {
  679. ${vars}
  680. };`;
  681. }
  682. _getWGSLStructBinding( name, vars, access, binding = 0, group = 0 ) {
  683. const structName = name + 'Struct';
  684. const structSnippet = this._getWGSLStruct( structName, vars );
  685. return `${structSnippet}
  686. @binding( ${binding} ) @group( ${group} )
  687. var<${access}> ${name} : ${structName};`;
  688. }
  689. }
  690. export default WGSLNodeBuilder;