WGSLNodeBuilder.js 27 KB

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