ComputeNode.js 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. import Node, { addNodeClass } from '../core/Node.js';
  2. import { NodeUpdateType } from '../core/constants.js';
  3. import { addNodeElement, nodeObject } from '../shadernode/ShaderNode.js';
  4. class ComputeNode extends Node {
  5. constructor( computeNode, count, workgroupSize = [ 64 ] ) {
  6. super( 'void' );
  7. this.isComputeNode = true;
  8. this.computeNode = computeNode;
  9. this.count = count;
  10. this.workgroupSize = workgroupSize;
  11. this.dispatchCount = 0;
  12. this.updateType = NodeUpdateType.OBJECT;
  13. this.updateDispatchCount();
  14. }
  15. updateDispatchCount() {
  16. const { count, workgroupSize } = this;
  17. let size = workgroupSize[ 0 ];
  18. for ( let i = 1; i < workgroupSize.length; i ++ )
  19. size *= workgroupSize[ i ];
  20. this.dispatchCount = Math.ceil( count / size );
  21. }
  22. onInit() { }
  23. update( { renderer } ) {
  24. renderer.compute( this );
  25. }
  26. generate( builder ) {
  27. const { shaderStage } = builder;
  28. if ( shaderStage === 'compute' ) {
  29. const snippet = this.computeNode.build( builder, 'void' );
  30. if ( snippet !== '' ) {
  31. builder.addLineFlowCode( snippet );
  32. }
  33. }
  34. }
  35. }
  36. export default ComputeNode;
  37. export const compute = ( node, count, workgroupSize ) => nodeObject( new ComputeNode( nodeObject( node ), count, workgroupSize ) );
  38. addNodeElement( 'compute', compute );
  39. addNodeClass( ComputeNode );