86 lines
1.5 KiB
JavaScript
86 lines
1.5 KiB
JavaScript
|
import Node, { addNodeClass } from '../core/Node.js';
|
||
|
|
import { NodeUpdateType } from '../core/constants.js';
|
||
|
|
import { addNodeElement, nodeObject } from '../shadernode/ShaderNode.js';
|
||
|
|
|
||
|
|
class ComputeNode extends Node {
|
||
|
|
|
||
|
|
constructor( computeNode, count, workgroupSize = [ 64 ] ) {
|
||
|
|
|
||
|
|
super( 'void' );
|
||
|
|
|
||
|
|
this.isComputeNode = true;
|
||
|
|
|
||
|
|
this.computeNode = computeNode;
|
||
|
|
|
||
|
|
this.count = count;
|
||
|
|
this.workgroupSize = workgroupSize;
|
||
|
|
this.dispatchCount = 0;
|
||
|
|
|
||
|
|
this.version = 1;
|
||
|
|
this.updateBeforeType = NodeUpdateType.OBJECT;
|
||
|
|
|
||
|
|
this.updateDispatchCount();
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
dispose() {
|
||
|
|
|
||
|
|
this.dispatchEvent( { type: 'dispose' } );
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
set needsUpdate( value ) {
|
||
|
|
|
||
|
|
if ( value === true ) this.version ++;
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
updateDispatchCount() {
|
||
|
|
|
||
|
|
const { count, workgroupSize } = this;
|
||
|
|
|
||
|
|
let size = workgroupSize[ 0 ];
|
||
|
|
|
||
|
|
for ( let i = 1; i < workgroupSize.length; i ++ )
|
||
|
|
size *= workgroupSize[ i ];
|
||
|
|
|
||
|
|
this.dispatchCount = Math.ceil( count / size );
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
onInit() { }
|
||
|
|
|
||
|
|
updateBefore( { renderer } ) {
|
||
|
|
|
||
|
|
renderer.compute( this );
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
generate( builder ) {
|
||
|
|
|
||
|
|
const { shaderStage } = builder;
|
||
|
|
|
||
|
|
if ( shaderStage === 'compute' ) {
|
||
|
|
|
||
|
|
const snippet = this.computeNode.build( builder, 'void' );
|
||
|
|
|
||
|
|
if ( snippet !== '' ) {
|
||
|
|
|
||
|
|
builder.addLineFlowCode( snippet );
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
export default ComputeNode;
|
||
|
|
|
||
|
|
export const compute = ( node, count, workgroupSize ) => nodeObject( new ComputeNode( nodeObject( node ), count, workgroupSize ) );
|
||
|
|
|
||
|
|
addNodeElement( 'compute', compute );
|
||
|
|
|
||
|
|
addNodeClass( 'ComputeNode', ComputeNode );
|