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 on nothing but the radius.

The radius is a USD attribute the run rewrites between renders, which is what makes the checks exact: at radius 0 the weight table is {1.0}, the blur is the identity, and the published AOV has to equal LdrColor byte for byte.

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 a minute on a cold shader cache. A successful run writes _output/input.png and _output/blurred.png and prints:

radius 0, pixels differing from the input: 0
radius 8, sharpest edge: 122 -> 24
radius 8, sharpest edge after one pass: 55
radius 8, mean brightness drift: 0.003
weight tables built: 2 (radii 8, 0)

The Weights

Built once per distinct radius rather than once per frame. The warning is what makes that countable: over some fifty rendered frames it appears twice.

-- One row of Gaussian taps for the stated radius, normalised so the blur keeps
-- the image's brightness. It depends on nothing but the radius, so it is built
-- here rather than recomputed by every GPU thread.
--
-- The warning is what makes the caching visible: this line appears once per
-- distinct radius over a run, not once per frame.
local function gaussianWeights(radius)
    warning("blur: building the weight table for radius " .. 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 what SPG derives when the launch script states nothing:

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.

Reading Back the Half-Finished Blur

The first pass feeds the second pass and a RenderVar of its own, so all three stages are readable from the host. The run checks that the sharpest edge in the image falls at each step, which only holds if the fan-out reached both consumers. The three figures move with the render; their order does not.

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.