render-pass.ts

import { useFrame } from '@react-three/fiber'
import { MutableRefObject, useMemo, useRef } from 'react'
import {
  BufferAttribute,
  BufferGeometry,
  FloatType,
  HalfFloatType,
  LinearFilter,
  Mesh,
  NearestFilter,
  OrthographicCamera,
  Scene,
  ShaderMaterial,
  Texture,
  Uniform,
  WebGLRenderTarget
} from 'three'
import pass from './pass.vert'

interface PropsBase {
  type?: 'forward' | 'feedback'
  name?: string
  fragmentShader: string
  width?: number
  height?: number
  needsUpdate?: MutableRefObject<boolean>
  textureFilter?: 'linear' | 'nearest'
  textureType?: 'float' | 'half'
}
/**
 * A shader pass that renders to a texture.
 */
interface PropsForward extends PropsBase {
  type?: 'forward'
  uniforms?: Record<string, THREE.IUniform>
}
/** Backbuffer feedback pass. */
interface PropsFeedback extends PropsBase {
  type?: 'feedback'
  uniforms?: (self: Texture) => Record<string, THREE.IUniform>
}
type Props = PropsForward | PropsFeedback

function createScreenQuadGeometry() {
  const geometry = new BufferGeometry()
  const vertices = new Float32Array([-1, -1, 3, -1, -1, 3])
  geometry.setAttribute('position', new BufferAttribute(vertices, 2))
  const uvs = new Float32Array([0, 0, 2, 0, 0, 2])
  geometry.setAttribute('uv', new BufferAttribute(uvs, 2))
  return geometry
}

export function useRenderPass({
  type = 'forward',
  textureFilter = 'linear',
  textureType = 'half',
  name = 'renderPass',
  fragmentShader,
  width: w,
  height: h,
  uniforms,
  needsUpdate
}: Props) {
  const shader = useRef<ShaderMaterial>()
  const target = useRef<WebGLRenderTarget>()
  const targetB = useRef<WebGLRenderTarget>()
  const location = useRef('')
  const uniform = useRef(new Uniform<Texture|null>(null))
  const meshRef = useRef<Mesh>()
  const geometry = useMemo(createScreenQuadGeometry, [])
  const scene = useMemo(() => new Scene(), [])
  const camera = useMemo(() => new OrthographicCamera(-1, 1, 1, -1, 0, 1), [])
  const frame = useRef(0)
  useFrame(({ gl }) => {
    const { width, height } = gl.domElement
    if (!target.current) {
      // create render target
      target.current = new WebGLRenderTarget(w || width, h || height, {
        type: textureType === 'float' ? FloatType : HalfFloatType,
        minFilter: textureFilter === 'linear' ? LinearFilter : NearestFilter,
        magFilter: textureFilter === 'linear' ? LinearFilter : NearestFilter
      })
      uniform.current.value = target.current.texture
      // console.log('created render target', target.current)
    }
    if (!targetB.current && type === 'feedback') {
      // create render target
      targetB.current = new WebGLRenderTarget(w || width, h || height, {
        type: FloatType
      })
    }
    // main logic goes here
    if (!shader.current) {
      // create shader material
      if (type === 'forward') {
        shader.current = new ShaderMaterial({
          vertexShader: pass,
          fragmentShader,
          // @ts-ignore
          uniforms: uniforms || {}
        })
        // console.log('created shader', shaderRef.current)
      } else if (type === 'feedback') {
        // determine the location of the backbuffer in the uniforms
        const unis = (
          uniforms as (self: Texture) => Record<string, THREE.IUniform>
        )(targetB.current!.texture)
        location.current =
          Object.keys(unis).find(
            (key) => unis[key].value === targetB.current!.texture
          ) || ''
        shader.current = new ShaderMaterial({
          vertexShader: pass,
          fragmentShader,
          uniforms: unis
        })

        shader.current.uniforms[location.current].value =
          targetB.current!.texture

        // console.log('created feedback shader', location.current)
      }
    }
    if (!meshRef.current) {
      // create mesh
      meshRef.current = new Mesh(geometry, shader.current)
      meshRef.current.frustumCulled = false
      scene.add(meshRef.current)
      // console.log('created mesh', scene)
    }

    if (needsUpdate !== undefined && needsUpdate.current === false) {
      return
    } else if (needsUpdate !== undefined) {
      needsUpdate.current = false
      if (shader.current!) {
        shader.current!.needsUpdate = true
        shader.current!.uniformsNeedUpdate = true
      }
    }

    // render to target
    if (type === 'forward') {
      gl.setRenderTarget(target.current)
    } else if (type === 'feedback') {
      if (frame.current % 2 === 0) {
        // set uniforms
        uniform.current.value = target.current!.texture
        shader.current!.uniforms[location.current].value =
          targetB.current!.texture
        // render to target
        gl.setRenderTarget(target.current!)
      } else {
        // set uniforms
        uniform.current.value = targetB.current!.texture
        shader.current!.uniforms[location.current].value =
          target.current!.texture
        // render to targetB
        gl.setRenderTarget(targetB.current!)
      }
    }
    // console.log('render pass', name)
    gl.render(scene, camera)
    gl.setRenderTarget(null)
    frame.current++
  })
  // return the uniform
  return uniform.current
}

This GitHub gist describes an example implementation of a "render pass" using @react-three/fiber in Typescript.

Depending on the type of computation, you can make use of textures and render passes to do what compute shaders would. It is what I do in webGL due to compute shader not being available for webGL yet. In general here is a breakdown of what is needed:

• Create a texture for input data
• Create a texture for output data
• Create a render target and assign the texture
• Pass input texture as a uniform to the compute (fragment) shader
• Run the fragment program in a render pass (frame) and set the render target to the target texture. Use a full screen quad as the vertex buffer and don’t modify the vertex attributes in the vertex shader.
• use the output texture in other shaders