WGSLNodeBuilder.js 22 KB

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