Skip to content
This repository has been archived by the owner on Jun 25, 2024. It is now read-only.

Commit

Permalink
performance: epic improvement of PixelSorter
Browse files Browse the repository at this point in the history 
thank god......
  • Loading branch information
@0b5vr
0b5vr committed Apr 7, 2023
1 parent 2420149 commit e78e23a
Show file tree
Hide file tree
Showing 4 changed files with 93 additions and 84 deletions.
58 changes: 28 additions & 30 deletions src/nodes/PostStack/PixelSorter/PixelSorter.ts
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import { ALIGN_SIZE, pixelSorterFrag } from './shaders/pixelSorterFrag';
import { Blit } from '../../../heck/components/Blit';
import { BufferTextureRenderTarget } from '../../../heck/BufferTextureRenderTarget';
import { GLTextureFormatStuffRG16F, GLTextureFormatStuffRGBA8 } from '../../../gl/glSetTexture';
import { GL_NEAREST, GL_TEXTURE_2D } from '../../../gl/constants';
import { Material } from '../../../heck/Material';
import { Quad } from '../../../heck/components/Quad';
Expand All @@ -15,8 +16,6 @@ import { quadGeometry } from '../../../globals/quadGeometry';
import { quadVert } from '../../../shaders/common/quadVert';
import { resizeObservers } from '../../../globals/globalObservers';

const TARGET_WIDTH = 2048;

export interface PixelSorterOptions {
input: BufferTextureRenderTarget;
target: RenderTarget;
Expand Down Expand Up @@ -49,26 +48,32 @@ export class PixelSorter extends SceneNode {

// -- buffers ----------------------------------------------------------------------------------
const swap = new Swap(
new BufferTextureRenderTarget( width, height ),
new BufferTextureRenderTarget( width, height ),
new BufferTextureRenderTarget( width, height, 1, GLTextureFormatStuffRGBA8 ),
new BufferTextureRenderTarget( width, height, 1, GLTextureFormatStuffRGBA8 ),
);

const bufferIndex = new BufferTextureRenderTarget( width, height );
const indexSwap = new Swap(
new BufferTextureRenderTarget( width, height, 1, GLTextureFormatStuffRG16F ),
new BufferTextureRenderTarget( width, height, 1, GLTextureFormatStuffRG16F ),
);

resizeObservers.push( ( [ width, height ] ) => {
swap.i.resize( width, height );
swap.o.resize( width, height );
bufferIndex.resize( width, height );
indexSwap.i.resize( width, height );
indexSwap.o.resize( width, height );
} );

glTextureFilter( swap.i.texture, GL_NEAREST );
glTextureFilter( swap.o.texture, GL_NEAREST );
glTextureFilter( bufferIndex.texture, GL_NEAREST );
glTextureFilter( indexSwap.i.texture, GL_NEAREST );
glTextureFilter( indexSwap.o.texture, GL_NEAREST );

if ( import.meta.env.DEV ) {
swap.i.name = 'PixelSorter/swap0';
swap.o.name = 'PixelSorter/swap1';
bufferIndex.name = 'PixelSorter/index';
indexSwap.i.name = 'PixelSorter/indexSwap0';
indexSwap.o.name = 'PixelSorter/indexSwap1';
}

// -- bypass -----------------------------------------------------------------------------------
Expand All @@ -84,32 +89,24 @@ export class PixelSorter extends SceneNode {
nodeBypass.children.push( blitBypass );

// -- calc index -------------------------------------------------------------------------------
let indexWidth = 1;
const indexMaterials: Material[] = [];

while ( indexWidth < TARGET_WIDTH ) {
const isLast = ( indexWidth * 8 >= TARGET_WIDTH );

const indexMaterials: Material[] = [ 1, 16, 256 ].map( ( indexWidth ) => {
const material = new Material(
quadVert,
pixelSorterIndexFrag,
pixelSorterIndexFrag( indexWidth ),
{ initOptions: { geometry: quadGeometry, target: dummyRenderTarget1 } },
);

material.addUniform( 'indexWidth', '1f', indexWidth );
material.addUniformTextures( 'sampler0', GL_TEXTURE_2D, input.texture );
material.addUniformTextures( 'sampler1', GL_TEXTURE_2D, swap.o.texture );
material.addUniformTextures( 'sampler1', GL_TEXTURE_2D, indexSwap.o.texture );

if ( import.meta.hot ) {
import.meta.hot.accept( './shaders/pixelSorterIndexFrag', ( { pixelSorterIndexFrag } ) => {
material.replaceShader( quadVert, pixelSorterIndexFrag );
material.replaceShader( quadVert, pixelSorterIndexFrag( indexWidth ) );
} );
}

indexMaterials.push( material );

const quad = new Quad( {
target: isLast ? bufferIndex : swap.i,
target: indexSwap.i,
material,
} );

Expand All @@ -119,18 +116,18 @@ export class PixelSorter extends SceneNode {

nodeMain.children.push( quad );

swap.swap();
indexSwap.swap();

indexWidth *= 8;
}
return material;
} );

// -- sort -------------------------------------------------------------------------------------
let dir = 1.0;
let dir = 2.0;
let comp = 1.0;
let isFirst = true;

while ( dir < ALIGN_SIZE ) {
const isFirst = dir === 1.0;
const isLast = ( dir === ALIGN_SIZE / 2.0 ) && ( comp === 1.0 );
while ( dir <= ALIGN_SIZE ) {
const isLast = ( dir === ALIGN_SIZE ) && ( comp === 1.0 );

const material = new Material(
quadVert,
Expand All @@ -148,7 +145,7 @@ export class PixelSorter extends SceneNode {
material.addUniformTextures(
'sampler1',
GL_TEXTURE_2D,
bufferIndex.texture,
indexSwap.o.texture,
);

if ( import.meta.hot ) {
Expand All @@ -171,11 +168,12 @@ export class PixelSorter extends SceneNode {
swap.swap();

if ( comp === 1.0 ) {
dir *= 2.0;
comp = dir;
dir *= 2.0;
} else {
comp /= 2.0;
}
isFirst = false;
}

// -- update uniform ---------------------------------------------------------------------------
Expand Down
63 changes: 36 additions & 27 deletions src/nodes/PostStack/PixelSorter/shaders/pixelSorterFrag.ts
View file
Original file line number Diff line number Diff line change
@@ -1,59 +1,68 @@
import { add, assign, build, def, defInNamed, defOut, defUniformNamed, div, dot, floor, glFragCoord, ifThen, insert, lt, mad, main, mix, mod, mul, neg, retFn, step, sub, subAssign, sw, texture, vec2, vec4 } from '../../../../shaders/shaderBuilder';
import { add, addAssign, assign, build, def, defOut, defUniformNamed, div, dot, eq, floor, glFragCoord, ifThen, insert, int, ivec2, lt, main, mix, mod, mul, neg, retFn, step, sub, subAssign, sw, texelFetch, vec2, vec4 } from '../../../../shaders/shaderBuilder';

const LUMA = vec4( 0.2126, 0.7152, 0.0722, 0.0 );
export const ALIGN_SIZE = 64.0;

export const ALIGN_SIZE = 256;

export const pixelSorterFrag = build( () => {
insert( 'precision highp float;' );

const vUv = defInNamed( 'vec2', 'vUv' );
const fragColor = defOut( 'vec4' );

const comp = defUniformNamed( 'float', 'comp' );
const dir = defUniformNamed( 'float', 'dir' );
const resolution = defUniformNamed( 'vec2', 'resolution' );
const sampler0 = defUniformNamed( 'sampler2D', 'sampler0' );
const sampler1 = defUniformNamed( 'sampler2D', 'sampler1' );

main( () => {
const texIndex = def( 'vec4', texture( sampler1, vUv ) );
const coord = ivec2( sw( glFragCoord, 'xy' ) );
const texIndex = def( 'vec4', texelFetch( sampler1, coord, int( 0 ) ) );
const tex = texelFetch( sampler0, coord, int( 0 ) );

assign( fragColor, tex );

ifThen( lt( sw( texIndex, 'x' ), 0.5 ), () => {
assign( fragColor, texture( sampler0, vUv ) );
retFn();
} );
ifThen( eq( sw( texIndex, 'x' ), 0.0 ), () => retFn() );

const index = sw( texIndex, 'x' );
const width = add( index, sw( texIndex, 'y' ), -1.0 );

const coordOrigin = def( 'vec2', sw( glFragCoord, 'xy' ) );
subAssign( sw( coordOrigin, 'x' ), sub( index, 1.0 ) );
ifThen( lt( width, div( ALIGN_SIZE, comp ) ), () => retFn() );

const segOrigin = def( 'vec2', vec2( coord ) );
subAssign( sw( segOrigin, 'x' ), sub( index, 1.0 ) );

const alignedIndex = def( 'float', floor( div( index, add( 1.0, width ), 1.0 / ALIGN_SIZE ) ) );
const alignDelta = def( 'vec2', vec2( div( width, ALIGN_SIZE ), 0.0 ) );
const alignDelta = def( 'float', div( width, ALIGN_SIZE ) );

const isCompHigher = def( 'float', step(
mod( alignedIndex, mul( 2.0, comp ) ),
sub( comp, 1.0 ),
const isRight = def( 'float', step(
1.0,
mod( div( alignedIndex, comp ), 2.0 ),
) );

const compOffset = mix( neg( comp ), comp, isCompHigher );
const compDelta = mul( comp, alignDelta );
const coordA = def( 'vec2', segOrigin );
addAssign( sw( coordA, 'x' ), mul( alignDelta, add( alignedIndex, 0.5 ) ) );
const coordB = def( 'vec2', coordA );
addAssign( sw( coordB, 'x' ), mix( compDelta, neg( compDelta ), isRight ) );

const coordA = mad( alignDelta, alignedIndex, coordOrigin );
const coordB = mad( alignDelta, compOffset, coordA );

const texA = texture( sampler0, div( coordA, resolution ) );
const texB = texture( sampler0, div( coordB, resolution ) );
const texA = texelFetch( sampler0, ivec2( coordA ), int( 0 ) );
const texB = texelFetch( sampler0, ivec2( coordB ), int( 0 ) );

const valueA = dot( texA, LUMA );
const valueB = dot( texB, LUMA );

const shouldSwap = def( 'float', step(
mod( div( alignedIndex, mul( 2.0, dir ) ), 2.0 ),
0.999,
) );
const shouldSwap2 = mod( add( shouldSwap, isCompHigher, step( valueB, valueA ) ), 2.0 ); // TODO
const shouldRightBeSmaller = step(
1.0,
mod( div( alignedIndex, dir ), 2.0 ),
);

// should we swap the pixel?
const shouldSwap = mod( add(
shouldRightBeSmaller,
isRight,
step( valueB, valueA ),
), 2.0 );

assign( fragColor, mix( texA, texB, shouldSwap2 ) );
assign( fragColor, mix( tex, texB, shouldSwap ) );
} );
} );
55 changes: 28 additions & 27 deletions src/nodes/PostStack/PixelSorter/shaders/pixelSorterIndexFrag.ts
View file
Original file line number Diff line number Diff line change
@@ -1,55 +1,56 @@
import { GLSLExpression, add, assign, build, defInNamed, defOut, defUniformNamed, div, dot, float, forLoop, insert, main, min, mix, mul, step, sub, sw, texture, vec2, vec4 } from '../../../../shaders/shaderBuilder';
import { GLSLExpression, add, assign, build, def, defOut, defUniformNamed, dot, float, forLoop, glFragCoord, gte, insert, int, ivec2, lt, main, min, mul, or, step, sub, sw, tern, texelFetch, vec4 } from '../../../../shaders/shaderBuilder';

const LUMA = vec4( 0.2126, 0.7152, 0.0722, 0.0 );

export const pixelSorterIndexFrag = build( () => {
export const pixelSorterIndexFrag = ( indexWidth: number ): string => build( () => {
insert( 'precision highp float;' );

const vUv = defInNamed( 'vec2', 'vUv' );
const fragColor = defOut( 'vec4' );

const threshold = defUniformNamed( 'float', 'threshold' );
const indexWidth = defUniformNamed( 'float', 'indexWidth' );
const resolution = defUniformNamed( 'vec2', 'resolution' );
const sampler0 = defUniformNamed( 'sampler2D', 'sampler0' );
const sampler1 = defUniformNamed( 'sampler2D', 'sampler1' );

function getValue( uv: GLSLExpression<'vec2'> ): GLSLExpression<'vec4'> {
const xMax = def( 'int' );

function getValue( coord: GLSLExpression<'ivec2'> ): GLSLExpression<'vec4'> {
// distance to the nearest wall, width of its current section, vec2( left, right )
const isEdge = add( step( sw( uv, 'x' ), 0.0 ), step( 1.0, sw( uv, 'x' ) ) );

return mix(
mix(
texture( sampler1, uv ),
vec4(
vec2( mul( step( dot( texture( sampler0, uv ), LUMA ), threshold ), 1E4 ) ),
0.0,
0.0,
),
step( indexWidth, 1.0 ),
),
vec4( 0.0 ),
isEdge,
);
const isEdge = or( lt( sw( coord, 'x' ), int( 0 ) ), gte( sw( coord, 'x' ), xMax ) );

// If it's the first round, set 0.0 or 1E4 using threshold
// If it's following rounds, look up existing left / right
const v = indexWidth === 1.0
? vec4( mul( step( dot(
texelFetch( sampler0, coord, int( 0 ) ),
LUMA,
), threshold ), 1E4 ) )
: texelFetch( sampler1, coord, int( 0 ) );

return tern( isEdge, vec4( 0.0 ), v );
}

main( () => {
assign( fragColor, getValue( vUv ) );
assign( xMax, int( sw( resolution, 'x' ) ) );

const coord = ivec2( sw( glFragCoord, 'xy' ) );
assign( fragColor, getValue( coord ) );

forLoop( 7, ( i ) => {
forLoop( 15, ( i ) => {
const offset = mul( add( float( i ), 1 ), indexWidth );
const uvOffset = div( vec2( offset, 0 ), resolution );

const uvtLeft = sub( vUv, uvOffset );
const coordLeft = sub( coord, ivec2( offset, 0 ) );
const valueLeft = sw( getValue( coordLeft ), 'x' );
assign( sw( fragColor, 'x' ), min(
sw( fragColor, 'x' ),
add( sw( getValue( uvtLeft ), 'x' ), offset ),
add( valueLeft, offset ),
) );

const uvtRight = add( vUv, uvOffset );
const coordRight = add( coord, ivec2( offset, 0 ) );
const valueRight = sw( getValue( coordRight ), 'y' );
assign( sw( fragColor, 'y' ), min(
sw( fragColor, 'y' ),
add( sw( getValue( uvtRight ), 'y' ), offset ),
add( valueRight, offset ),
) );
} );
} );
Expand Down
1 change: 1 addition & 0 deletions src/shaders/shaderBuilder.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -415,6 +415,7 @@ export const rshift: {

export const neg: {
( x: Exf ): Ex<'float'>;
( x: Ex<'int'> ): Ex<'int'>;
<T extends GLSLGenType>( x: Ex<T> ): Ex<T>;
} = ( x: string | number ) => (
`(-${ x })`
Expand Down

0 comments on commit e78e23a

Please sign in to comment.