Function Overloads

[k2d222]

Function overloads allow redefining a function with different input/output parameters. It is nice to reuse the same function name for different parameters, e.g. implementing a function for f32, vec2, vec3, vec4... most of the numeric built-in functions do that.

What the spec does

WGSL spec has a concept of overload, which manifests itself in functions via function overloads and type parameters.

Function overloads are variants of the same function that take a different number of arguments or template parameters.

Function type parameters are similar-looking to generics, except that they are overloads, meaning each possible variant is implemented separately.

How we could implement it

This proposal is connected to generics, but is distinct. To support this, a WESL linker would need to implement the overload resolution algorithm which roughly is this:

• find all function overloads (function declarations with the same name)
• keep only those satisfied by the static types on the caller side (need to evaluate static types! not trivial)
• find the best candidate that has the lowest conversion rank

Side-Quest: why are "type parameters" not just generics? Why do we need overloads?

@const @must_use fn abs(e: T ) -> T

• S is AbstractInt, AbstractFloat, i32, u32, f32, or f16
• T is S, or vecN<S>

Here's how it's implemented with overloads:

fn abs<T: AbstractFloat>(v: vec3<T>) -> vec3<T> {
    return max(v, -v);
}
fn abs<T: AbstractFloat>(s: T) -> T {
    if s > 0 { return s; } else { return -s; }
}

// usage
let x: f32 = -5.0;
let _ = abs(x); // 5.0f
let y = vec3<f32>(3.0, -10.0, -4.0);
let _ = abs(y); // vec3f(3.0, 10.0, 4.0)

How could you implement abs in one go for both vec<T> and T? You can't, even with generics, unless you write complex traits that allow scalars to behave like vectors. See this example with module generics:

mod ScalarLike {
    alias T;
    fn zero() -> T;
    fn lt(lhs: T, rhs: T) -> bool;  // lhs < rhs
    fn sub(lhs: T, rhs: T) -> T   // lhs - rhs
}

mod ArrayLike<S: ScalarLike> {
    alias T;
    fn len(arr: T) -> u32;  //  arrayLength(arr)
    fn get(arr: T, index: u32) -> S::T;  // arr[index]
    fn set(arr: T, index: u32, value: S::T): T; // arr[index] = value
}

fn abs<S: ScalarLike, A: ArrayLike<S>>(arr: A::T) -> A::T {
    var res = arr;
    for (var i = 0u; i < A::len(arr); i++) {
        var scalar = A::get(arr, i);
        if S::lt(scalar, S::zero()) {
            scalar  = S::sub(S::zero(), scalar);
        }
        res = A::set(res, i, scalar);
    }
    return res;
}

...and then implement the concretizations...

@override(ScalarLike)
mod F32LikeScalar {
    alias T = f32;
    fn lt(lhs: T, rhs: T) -> bool { return lhs < rhs; }
    fn zero() -> T { return 0.0f; }
    fn sub(lhs: T, rhs: T) -> T { return lhs - rhs; }
}

@override(ArrayLike<S>)
mod ScalarLikeArray<S: ScalarLike> {
    alias T = S::T;
    fn len(arr: T) -> u32 { return 1u; }
    fn get(arr: T, index: u32) -> T { return arr; }
    fn set(arr: T, index: u32, value: S::T) -> T { return value; }
}

... and finally, usage of abs()...

let x: f32 = -5.0;
let _ = abs<F32LikeScalar, ScalarLikeArray<F32LikeScalar>>(x); // 5.0f
let y = vec3<f32>(3.0, -10.0, -4.0);
let _ = abs<F32LikeScalar, Vec3LikeArray<F32LikeScalar>>(y); // vec3f(3.0, 10.0, 4.0)

This is nicely extensible but terribly, terribly verbose!

[ncthbrt]

I wonder if we could simplify this by using generic constraints? Like instead of resolving types and using inference, we could match on the generic arguments?

[k2d222]

I wonder if we could simplify this by using generic constraints? Like instead of resolving types and using inference, we could match on the generic arguments?

aah, do you mean something roughly like that, with type algebra?

fn abs<S: AbstractFloat, T: S || vec3<S>>(e: T) -> T {
    @if(T == Vec3<S>) {
        // implementation with vec3
    }
    @if(T == S) {
        // implementation with scalars
    }
}

I love this one. It expresses exactly what's in the spec.

[k2d222]

Might be worth opening another issue for algebraic data types and type constraints

[ncthbrt]

I was thinking more overloads still but that the matching is done on the generic constraints

[ncthbrt]

But that is a cool idea too!

[mighdoll]

Do you have a use case in your own shader development for function overloads? I'd suggest we make a separate bug for the use case, and we can work with @sdedovic to see if it's the same as the one lygia describes (or make a separate case for the lygia issue). That'll give us some concrete examples to check our potential designs against.