WGSLNodeBuilder.js 23 KB

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