WGSLNodeBuilder.js 23 KB

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