WGSLNodeBuilder.js 23 KB

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