WebGPUComputePipelines.js 1.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  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. get( computeNode ) {
  12. let pipeline = this.pipelines.get( computeNode );
  13. // @TODO: Reuse compute pipeline if possible, introduce WebGPUComputePipeline
  14. if ( pipeline === undefined ) {
  15. const device = this.device;
  16. // get shader
  17. const nodeBuilder = this.nodes.get( computeNode );
  18. const computeShader = nodeBuilder.computeShader;
  19. const shader = {
  20. computeShader
  21. };
  22. // programmable stage
  23. let stageCompute = this.stages.compute.get( shader );
  24. if ( stageCompute === undefined ) {
  25. stageCompute = new WebGPUProgrammableStage( device, computeShader, 'compute' );
  26. this.stages.compute.set( shader, stageCompute );
  27. }
  28. pipeline = device.createComputePipeline( {
  29. compute: stageCompute.stage,
  30. layout: 'auto'
  31. } );
  32. this.pipelines.set( computeNode, pipeline );
  33. }
  34. return pipeline;
  35. }
  36. dispose() {
  37. this.pipelines = new WeakMap();
  38. this.stages = {
  39. compute: new WeakMap()
  40. };
  41. }
  42. }
  43. export default WebGPUComputePipelines;