WebGPUComputePipelines.js 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. import WebGPUProgrammableStage from './WebGPUProgrammableStage.js';
  2. class WebGPUComputePipelines {
  3. constructor( device, nodes ) {
  4. this.device = device;
  5. this.nodes = nodes;
  6. this.pipelines = new WeakMap();
  7. this.stages = {
  8. compute: new WeakMap()
  9. };
  10. }
  11. has( computeNode ) {
  12. return this.pipelines.get( computeNode ) !== undefined;
  13. }
  14. get( computeNode ) {
  15. let pipeline = this.pipelines.get( computeNode );
  16. // @TODO: Reuse compute pipeline if possible, introduce WebGPUComputePipeline
  17. if ( pipeline === undefined ) {
  18. const device = this.device;
  19. // get shader
  20. const nodeBuilder = this.nodes.get( computeNode );
  21. const computeShader = nodeBuilder.computeShader;
  22. const shader = {
  23. computeShader
  24. };
  25. // programmable stage
  26. let stageCompute = this.stages.compute.get( shader );
  27. if ( stageCompute === undefined ) {
  28. stageCompute = new WebGPUProgrammableStage( device, computeShader, 'compute' );
  29. this.stages.compute.set( shader, stageCompute );
  30. }
  31. pipeline = device.createComputePipeline( {
  32. compute: stageCompute.stage,
  33. layout: 'auto'
  34. } );
  35. this.pipelines.set( computeNode, pipeline );
  36. }
  37. return pipeline;
  38. }
  39. dispose() {
  40. this.pipelines = new WeakMap();
  41. this.stages = {
  42. compute: new WeakMap()
  43. };
  44. }
  45. }
  46. export default WebGPUComputePipelines;