WebGPUNodeBuilder.js 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679
  1. import WebGPUNodeUniformsGroup from './WebGPUNodeUniformsGroup.js';
  2. import {
  3. FloatNodeUniform, Vector2NodeUniform, Vector3NodeUniform, Vector4NodeUniform,
  4. ColorNodeUniform, Matrix3NodeUniform, Matrix4NodeUniform
  5. } from './WebGPUNodeUniform.js';
  6. import WebGPUNodeSampler from './WebGPUNodeSampler.js';
  7. import { WebGPUNodeSampledTexture, WebGPUNodeSampledCubeTexture } from './WebGPUNodeSampledTexture.js';
  8. import WebGPUUniformBuffer from '../WebGPUUniformBuffer.js';
  9. import { getVectorLength, getStrideLength } from '../WebGPUBufferUtils.js';
  10. import NodeBuilder from 'three-nodes/core/NodeBuilder.js';
  11. import WGSLNodeParser from 'three-nodes/parsers/WGSLNodeParser.js';
  12. import CodeNode from 'three-nodes/core/CodeNode.js';
  13. import { NodeMaterial } from 'three-nodes/materials/Materials.js';
  14. const supports = {
  15. instance: true
  16. };
  17. const wgslTypeLib = {
  18. float: 'f32',
  19. int: 'i32',
  20. uint: 'u32',
  21. bool: 'bool',
  22. vec2: 'vec2<f32>',
  23. ivec2: 'vec2<i32>',
  24. uvec2: 'vec2<u32>',
  25. bvec2: 'vec2<bool>',
  26. vec3: 'vec3<f32>',
  27. ivec3: 'vec3<i32>',
  28. uvec3: 'vec3<u32>',
  29. bvec3: 'vec3<bool>',
  30. vec4: 'vec4<f32>',
  31. ivec4: 'vec4<i32>',
  32. uvec4: 'vec4<u32>',
  33. bvec4: 'vec4<bool>',
  34. mat3: 'mat3x3<f32>',
  35. imat3: 'mat3x3<i32>',
  36. umat3: 'mat3x3<u32>',
  37. bmat3: 'mat3x3<bool>',
  38. mat4: 'mat4x4<f32>',
  39. imat4: 'mat4x4<i32>',
  40. umat4: 'mat4x4<u32>',
  41. bmat4: 'mat4x4<bool>'
  42. };
  43. const wgslMethods = {
  44. dFdx: 'dpdx',
  45. dFdy: 'dpdy',
  46. inversesqrt: 'inverseSqrt'
  47. };
  48. const wgslPolyfill = {
  49. lessThanEqual: new CodeNode( `
  50. fn lessThanEqual( a : vec3<f32>, b : vec3<f32> ) -> vec3<bool> {
  51. return vec3<bool>( a.x <= b.x, a.y <= b.y, a.z <= b.z );
  52. }
  53. ` ),
  54. mod: new CodeNode( `
  55. fn mod( x : f32, y : f32 ) -> f32 {
  56. return x - y * floor( x / y );
  57. }
  58. ` ),
  59. repeatWrapping: new CodeNode( `
  60. fn repeatWrapping( uv : vec2<f32>, dimension : vec2<i32> ) -> vec2<i32> {
  61. let uvScaled = vec2<i32>( uv * vec2<f32>( dimension ) );
  62. return ( ( uvScaled % dimension ) + dimension ) % dimension;
  63. }
  64. ` )
  65. };
  66. class WebGPUNodeBuilder extends NodeBuilder {
  67. constructor( object, renderer ) {
  68. super( object, renderer, new WGSLNodeParser() );
  69. this.lightNode = null;
  70. this.fogNode = null;
  71. this.bindings = { vertex: [], fragment: [] };
  72. this.bindingsOffset = { vertex: 0, fragment: 0 };
  73. this.uniformsGroup = {};
  74. this.builtins = new Set();
  75. }
  76. build() {
  77. NodeMaterial.fromMaterial( this.material ).build( this );
  78. return super.build();
  79. }
  80. addFlowCode( code ) {
  81. if ( ! /;\s*$/.test( code ) ) {
  82. code += ';';
  83. }
  84. super.addFlowCode( code + '\n\t' );
  85. }
  86. getSampler( textureProperty, uvSnippet, shaderStage = this.shaderStage ) {
  87. if ( shaderStage === 'fragment' ) {
  88. return `textureSample( ${textureProperty}, ${textureProperty}_sampler, ${uvSnippet} )`;
  89. } else {
  90. this._include( 'repeatWrapping' );
  91. const dimension = `textureDimensions( ${textureProperty}, 0 )`;
  92. return `textureLoad( ${textureProperty}, repeatWrapping( ${uvSnippet}, ${dimension} ), 0 )`;
  93. }
  94. }
  95. getTexture( textureProperty, uvSnippet, shaderStage = this.shaderStage ) {
  96. return this.getSampler( textureProperty, uvSnippet, shaderStage );
  97. }
  98. getCubeTexture( textureProperty, uvSnippet, shaderStage = this.shaderStage ) {
  99. return this.getSampler( textureProperty, uvSnippet, shaderStage );
  100. }
  101. getPropertyName( node, shaderStage = this.shaderStage ) {
  102. if ( node.isNodeVary === true ) {
  103. if ( shaderStage === 'vertex' ) {
  104. return `NodeVarys.${ node.name }`;
  105. }
  106. } else if ( node.isNodeUniform === true ) {
  107. const name = node.name;
  108. const type = node.type;
  109. if ( type === 'texture' || type === 'cubeTexture' ) {
  110. return name;
  111. } else if ( type === 'buffer' ) {
  112. return `NodeBuffer_${node.node.id}.${name}`;
  113. } else {
  114. return `NodeUniforms.${name}`;
  115. }
  116. }
  117. return super.getPropertyName( node );
  118. }
  119. getBindings() {
  120. const bindings = this.bindings;
  121. return [ ...bindings.vertex, ...bindings.fragment ];
  122. }
  123. getUniformFromNode( node, shaderStage, type ) {
  124. const uniformNode = super.getUniformFromNode( node, shaderStage, type );
  125. const nodeData = this.getDataFromNode( node, shaderStage );
  126. if ( nodeData.uniformGPU === undefined ) {
  127. let uniformGPU;
  128. const bindings = this.bindings[ shaderStage ];
  129. if ( type === 'texture' || type === 'cubeTexture' ) {
  130. const sampler = new WebGPUNodeSampler( `${uniformNode.name}_sampler`, uniformNode.node );
  131. let texture = null;
  132. if ( type === 'texture' ) {
  133. texture = new WebGPUNodeSampledTexture( uniformNode.name, uniformNode.node );
  134. } else if ( type === 'cubeTexture' ) {
  135. texture = new WebGPUNodeSampledCubeTexture( uniformNode.name, uniformNode.node );
  136. }
  137. // add first textures in sequence and group for last
  138. const lastBinding = bindings[ bindings.length - 1 ];
  139. const index = lastBinding && lastBinding.isUniformsGroup ? bindings.length - 1 : bindings.length;
  140. if ( shaderStage === 'fragment' ) {
  141. bindings.splice( index, 0, sampler, texture );
  142. uniformGPU = [ sampler, texture ];
  143. } else {
  144. bindings.splice( index, 0, texture );
  145. uniformGPU = [ texture ];
  146. }
  147. } else if ( type === 'buffer' ) {
  148. const buffer = new WebGPUUniformBuffer( 'NodeBuffer_' + node.id, node.value );
  149. // add first textures in sequence and group for last
  150. const lastBinding = bindings[ bindings.length - 1 ];
  151. const index = lastBinding && lastBinding.isUniformsGroup ? bindings.length - 1 : bindings.length;
  152. bindings.splice( index, 0, buffer );
  153. uniformGPU = buffer;
  154. } else {
  155. let uniformsGroup = this.uniformsGroup[ shaderStage ];
  156. if ( uniformsGroup === undefined ) {
  157. uniformsGroup = new WebGPUNodeUniformsGroup( shaderStage );
  158. this.uniformsGroup[ shaderStage ] = uniformsGroup;
  159. bindings.push( uniformsGroup );
  160. }
  161. if ( node.isArrayUniformNode === true ) {
  162. uniformGPU = [];
  163. for ( const uniformNode of node.nodes ) {
  164. const uniformNodeGPU = this._getNodeUniform( uniformNode, type );
  165. // fit bounds to buffer
  166. uniformNodeGPU.boundary = getVectorLength( uniformNodeGPU.itemSize );
  167. uniformNodeGPU.itemSize = getStrideLength( uniformNodeGPU.itemSize );
  168. uniformsGroup.addUniform( uniformNodeGPU );
  169. uniformGPU.push( uniformNodeGPU );
  170. }
  171. } else {
  172. uniformGPU = this._getNodeUniform( uniformNode, type );
  173. uniformsGroup.addUniform( uniformGPU );
  174. }
  175. }
  176. nodeData.uniformGPU = uniformGPU;
  177. if ( shaderStage === 'vertex' ) {
  178. this.bindingsOffset[ 'fragment' ] = bindings.length;
  179. }
  180. }
  181. return uniformNode;
  182. }
  183. isReference( type ) {
  184. return super.isReference( type ) || type === 'texture_2d' || type === 'texture_cube';
  185. }
  186. getInstanceIndex( shaderStage = this.shaderStage ) {
  187. this.builtins.add( 'instance_index' );
  188. if ( shaderStage === 'vertex' ) {
  189. return 'instanceIndex';
  190. }
  191. }
  192. getAttributes( shaderStage ) {
  193. const snippets = [];
  194. if ( shaderStage === 'vertex' ) {
  195. if ( this.builtins.has( 'instance_index' ) ) {
  196. snippets.push( `@builtin( instance_index ) instanceIndex : u32` );
  197. }
  198. const attributes = this.attributes;
  199. const length = attributes.length;
  200. for ( let index = 0; index < length; index ++ ) {
  201. const attribute = attributes[ index ];
  202. const name = attribute.name;
  203. const type = this.getType( attribute.type );
  204. snippets.push( `@location( ${index} ) ${ name } : ${ type }` );
  205. }
  206. }
  207. return snippets.join( ',\n\t' );
  208. }
  209. getVars( shaderStage ) {
  210. const snippets = [];
  211. const vars = this.vars[ shaderStage ];
  212. for ( let index = 0; index < vars.length; index ++ ) {
  213. const variable = vars[ index ];
  214. const name = variable.name;
  215. const type = this.getType( variable.type );
  216. snippets.push( `\tvar ${name} : ${type};` );
  217. }
  218. return `\n${ snippets.join( '\n' ) }\n`;
  219. }
  220. getVarys( shaderStage ) {
  221. const snippets = [];
  222. if ( shaderStage === 'vertex' ) {
  223. snippets.push( '@builtin( position ) Vertex: vec4<f32>' );
  224. const varys = this.varys;
  225. for ( let index = 0; index < varys.length; index ++ ) {
  226. const vary = varys[ index ];
  227. snippets.push( `@location( ${index} ) ${ vary.name } : ${ this.getType( vary.type ) }` );
  228. }
  229. } else if ( shaderStage === 'fragment' ) {
  230. const varys = this.varys;
  231. for ( let index = 0; index < varys.length; index ++ ) {
  232. const vary = varys[ index ];
  233. snippets.push( `@location( ${index} ) ${ vary.name } : ${ this.getType( vary.type ) }` );
  234. }
  235. }
  236. const code = snippets.join( ',\n\t' );
  237. return shaderStage === 'vertex' ? this._getWGSLStruct( 'NodeVarysStruct', '\t' + code ) : code;
  238. }
  239. getUniforms( shaderStage ) {
  240. const uniforms = this.uniforms[ shaderStage ];
  241. const bindingSnippets = [];
  242. const bufferSnippets = [];
  243. const groupSnippets = [];
  244. let index = this.bindingsOffset[ shaderStage ];
  245. for ( const uniform of uniforms ) {
  246. if ( uniform.type === 'texture' ) {
  247. if ( shaderStage === 'fragment' ) {
  248. bindingSnippets.push( `@group( 0 ) @binding( ${index ++} ) var ${uniform.name}_sampler : sampler;` );
  249. }
  250. bindingSnippets.push( `@group( 0 ) @binding( ${index ++} ) var ${uniform.name} : texture_2d<f32>;` );
  251. } else if ( uniform.type === 'cubeTexture' ) {
  252. if ( shaderStage === 'fragment' ) {
  253. bindingSnippets.push( `@group( 0 ) @binding( ${index ++} ) var ${uniform.name}_sampler : sampler;` );
  254. }
  255. bindingSnippets.push( `@group( 0 ) @binding( ${index ++} ) var ${uniform.name} : texture_cube<f32>;` );
  256. } else if ( uniform.type === 'buffer' ) {
  257. const bufferNode = uniform.node;
  258. const bufferType = this.getType( bufferNode.bufferType );
  259. const bufferCount = bufferNode.bufferCount;
  260. const bufferSnippet = `\t${uniform.name} : array< ${bufferType}, ${bufferCount} >\n`;
  261. bufferSnippets.push( this._getWGSLUniforms( 'NodeBuffer_' + bufferNode.id, bufferSnippet, index ++ ) );
  262. } else {
  263. const vectorType = this.getType( this.getVectorType( uniform.type ) );
  264. if ( Array.isArray( uniform.value ) === true ) {
  265. const length = uniform.value.length;
  266. groupSnippets.push( `uniform ${vectorType}[ ${length} ] ${uniform.name}` );
  267. } else {
  268. groupSnippets.push( `\t${uniform.name} : ${ vectorType}` );
  269. }
  270. }
  271. }
  272. let code = bindingSnippets.join( '\n' );
  273. code += bufferSnippets.join( '\n' );
  274. if ( groupSnippets.length > 0 ) {
  275. code += this._getWGSLUniforms( 'NodeUniforms', groupSnippets.join( ',\n' ), index ++ );
  276. }
  277. return code;
  278. }
  279. buildCode() {
  280. const shadersData = { fragment: {}, vertex: {} };
  281. for ( const shaderStage in shadersData ) {
  282. let flow = '// code\n';
  283. flow += `\t${ this.flowCode[ shaderStage ] }`;
  284. flow += '\n\t';
  285. const flowNodes = this.flowNodes[ shaderStage ];
  286. const mainNode = flowNodes[ flowNodes.length - 1 ];
  287. for ( const node of flowNodes ) {
  288. const flowSlotData = this.getFlowData( shaderStage, node );
  289. const slotName = node.name;
  290. if ( slotName ) {
  291. if ( flow.length > 0 ) flow += '\n';
  292. flow += `\t// FLOW -> ${ slotName }\n\t`;
  293. }
  294. flow += `${ flowSlotData.code }\n\t`;
  295. if ( node === mainNode ) {
  296. flow += '// FLOW RESULT\n\t';
  297. if ( shaderStage === 'vertex' ) {
  298. flow += 'NodeVarys.Vertex = ';
  299. } else if ( shaderStage === 'fragment' ) {
  300. flow += 'return ';
  301. }
  302. flow += `${ flowSlotData.result };`;
  303. }
  304. }
  305. const stageData = shadersData[ shaderStage ];
  306. stageData.uniforms = this.getUniforms( shaderStage );
  307. stageData.attributes = this.getAttributes( shaderStage );
  308. stageData.varys = this.getVarys( shaderStage );
  309. stageData.vars = this.getVars( shaderStage );
  310. stageData.codes = this.getCodes( shaderStage );
  311. stageData.flow = flow;
  312. }
  313. this.vertexShader = this._getWGSLVertexCode( shadersData.vertex );
  314. this.fragmentShader = this._getWGSLFragmentCode( shadersData.fragment );
  315. }
  316. getMethod( method ) {
  317. if ( wgslPolyfill[ method ] !== undefined ) {
  318. this._include( method );
  319. }
  320. return wgslMethods[ method ] || method;
  321. }
  322. getType( type ) {
  323. return wgslTypeLib[ type ] || type;
  324. }
  325. isAvailable( name ) {
  326. return supports[ name ] === true;
  327. }
  328. _include( name ) {
  329. wgslPolyfill[ name ].build( this );
  330. }
  331. _getNodeUniform( uniformNode, type ) {
  332. if ( type === 'float' ) return new FloatNodeUniform( uniformNode );
  333. if ( type === 'vec2' ) return new Vector2NodeUniform( uniformNode );
  334. if ( type === 'vec3' ) return new Vector3NodeUniform( uniformNode );
  335. if ( type === 'vec4' ) return new Vector4NodeUniform( uniformNode );
  336. if ( type === 'color' ) return new ColorNodeUniform( uniformNode );
  337. if ( type === 'mat3' ) return new Matrix3NodeUniform( uniformNode );
  338. if ( type === 'mat4' ) return new Matrix4NodeUniform( uniformNode );
  339. throw new Error( `Uniform "${type}" not declared.` );
  340. }
  341. _getWGSLVertexCode( shaderData ) {
  342. return `${ this.getSignature() }
  343. // uniforms
  344. ${shaderData.uniforms}
  345. // varys
  346. ${shaderData.varys}
  347. // codes
  348. ${shaderData.codes}
  349. @stage( vertex )
  350. fn main( ${shaderData.attributes} ) -> NodeVarysStruct {
  351. // system
  352. var NodeVarys: NodeVarysStruct;
  353. // vars
  354. ${shaderData.vars}
  355. // flow
  356. ${shaderData.flow}
  357. return NodeVarys;
  358. }
  359. `;
  360. }
  361. _getWGSLFragmentCode( shaderData ) {
  362. return `${ this.getSignature() }
  363. // uniforms
  364. ${shaderData.uniforms}
  365. // codes
  366. ${shaderData.codes}
  367. @stage( fragment )
  368. fn main( ${shaderData.varys} ) -> @location( 0 ) vec4<f32> {
  369. // vars
  370. ${shaderData.vars}
  371. // flow
  372. ${shaderData.flow}
  373. }
  374. `;
  375. }
  376. _getWGSLStruct( name, vars ) {
  377. return `
  378. struct ${name} {
  379. ${vars}
  380. };`;
  381. }
  382. _getWGSLUniforms( name, vars, binding = 0, group = 0 ) {
  383. const structName = name + 'Struct';
  384. const structSnippet = this._getWGSLStruct( structName, vars );
  385. return `${structSnippet}
  386. @binding( ${binding} ) @group( ${group} )
  387. var<uniform> ${name} : ${structName};`;
  388. }
  389. }
  390. export default WebGPUNodeBuilder;