2
0

WGSLNodeBuilder.js 23 KB

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