Python: SPG Separable Blur

Two chained SPG nodes blur LdrColor, first along x and then along y, and publish the result as LdrBlurred. The tap weights are built in the launch script rather than on the GPU, because they depend only on the radius.

The radius is a USD attribute the run rewrites between renders: at radius 0 the weight table is {1.0} and the blur is the identity.

Running

uv run main.py                                # CUDA
uv run main.py --scene blur_scene_slang.usda  # Slang

The first step compiles the kernel and can take up to 5 minutes on a cold shader cache. A successful run writes five PNGs to _output/: input.png, the rendered LdrColor; blurred-one-pass.png, after the horizontal pass only; blurred.png, after both passes at radius 8; and, once the radius is set to 0, radius-0-input.png and radius-0.png, that frame’s LdrColor and the published AOV, which match pixel for pixel.

The Weights

Each of the two blur nodes builds its table for the first radius and again when the radius changes, rather than once per frame: the radius changes once during the run, so each node builds its table twice.

-- One row of Gaussian taps for the stated radius, normalised so the blur keeps
-- the image's brightness. It depends only on the radius, so it is built
-- here rather than recomputed by every GPU thread.
local function gaussianWeights(radius)
    if radius == 0 then
        return cuda.array({ 1.0 }, cuda.float)
    end

    local sigma = radius / 3.0
    local taps, total = {}, 0.0
    for t = -radius, radius do
        local w = math.exp(-0.5 * (t / sigma) ^ 2)
        taps[#taps + 1] = w
        total = total + w
    end
    for i = 1, #taps do
        taps[i] = taps[i] / total
    end
    return cuda.array(taps, cuda.float)
end

The Two Launches

The horizontal pass takes one thread per output pixel, which is SPG’s default launch geometry:

function blurHorizontal(inputs, outputs)
    assert(inputs["Image"].rank == 2, "Input must be a 2D image")

    local height = inputs["Image"].shape[1]
    local width = inputs["Image"].shape[2]
    local radius = inputs["radius"].value

    outputs["Blurred"] = cuda.image(width, height, cuda.uchar4)

    return cuda.kernel({
        args = {
            cuda.int(width),
            cuda.int(height),
            cuda.int(radius),
            -- Built once rather than on every frame: cuda.static caches the
            -- call against its arguments, and this script runs per frame.
            cuda.static(gaussianWeights, radius),
            cuda.TextureObject(inputs["Image"]),
            cuda.SurfaceObject(outputs["Blurred"]),
        },
        -- No block and no grid. One thread per output pixel is what SPG derives
        -- from the output's shape, and that is the mapping this kernel wants.
    })
end

The vertical pass gives each thread a whole column to walk, so the iteration domain is the width alone and the geometry has to be stated:

function blurVertical(inputs, outputs)
    assert(inputs["Image"].rank == 2, "Input must be a 2D image")

    local height = inputs["Image"].shape[1]
    local width = inputs["Image"].shape[2]
    local radius = inputs["radius"].value

    outputs["Blurred"] = cuda.image(width, height, cuda.uchar4)

    return cuda.kernel({
        args = {
            cuda.int(width),
            cuda.int(height),
            cuda.int(radius),
            cuda.static(gaussianWeights, radius),
            cuda.TextureObject(inputs["Image"]),
            cuda.SurfaceObject(outputs["Blurred"]),
        },
        -- One thread per column, so the domain is the width alone. The derived
        -- geometry would cover width x height and launch a thread for every
        -- pixel, each of which would then walk the whole column.
        block = { BLOCK },
        grid = { math.ceil(width / BLOCK) },
    })
end

Sharing Code Between the Passes

Both passes run the same tap loop. On Slang it lives in its own module beside the shader, reached with import BlurTaps;:

// The tap loop both blur passes run, kept in its own module so neither entry
// point restates it. This file declares no entry point and is not a node: it is
// reached only through `import BlurTaps;` in BlurKernel.slang, which resolves to
// this file because it sits beside it.
//
// `public` is what makes a symbol visible to an importing module.
public float3 blurTap(Texture2D<float4> image,
                      StructuredBuffer<float> weights,
                      int radius,
                      int width,
                      int height,
                      int2 pixel,
                      int2 direction)
{
    float3 sum = float3(0.0, 0.0, 0.0);
    for (int t = -radius; t <= radius; ++t)
    {
        int2 at = clamp(pixel + t * direction, int2(0, 0), int2(width - 1, height - 1));
        sum += weights[t + radius] * image.Load(int3(at, 0)).rgb;
    }
    return sum;
}

CUDA has no equivalent. One include directory is searched and it is never the .cu’s own, so the two CUDA passes share the loop through a __device__ function in the same file. Refer to Compile It Yourself.

Changing the Radius

inputs:radius is connected to nothing, so writing it on the Shader prim takes effect on the next step with no reset:

def set_radius(stage, ordinal: int, radius: int) -> None:
    """Write inputs:radius on both blur nodes and publish the edit.

    The attribute is connected to nothing, so the value on the prim is the value
    the node reads. Each node is written separately because each carries its own
    copy of the attribute.
    """
    value = np.array([radius], dtype=np.int32)
    tensor = ovstage.make_dltensor(
        value, dtype=ovstage.numpy_to_dldatatype(value.dtype, lanes=1), shape=[1], ndim=1
    )
    with ovstage.PathDictionary(stage) as paths:
        attribute = paths.intern_token("inputs:radius")
        for node in BLUR_NODES:
            path_list = paths.create_path_list_from_strings([node])
            with stage.query_from_path_list(path_list) as query:
                stage.write_attribute(
                    query, attribute, ordinal=ordinal, tensors=tensor, is_array=False
                ).wait()
            paths.destroy_path_list(path_list)
    stage.advance_write_floor(ordinal, ovstage.Scope.ALL).wait()

Refer to Upload Your Own Data, Cache Work Across Frames, Control the Launch Geometry and Change Values and Rewire for the same material in prose.