Proposal: Unified "conditional expression" for Boolable or SIMD masks

Mojo
1

Today, `math.max()` has separate implementations for scalar and SIMD values.

Conceptually, both are doing this in the 2-input case:

a if a > b else b

However, scalar comparisons produce a Bool/Boolable, while SIMD comparisons produce a mask.

Instead of having parallel implementations for Boolable conditions vs. SIMD masks, it would be better to be able to write generic elementwise algorithms that work for either scalar or SIMD types.

Two possible approaches:

Option A: stdlib select()

select(cond, a, b)

where select() accepts either a Boolable condition or a SIMD mask.

Then math.max() could conceptually be expressed as:

select(a > b, a, b)

for both scalar and SIMD inputs.

The implementation could be something like:

def select(cond, a, b):
    comptime if conforms_to(typeof(cond), Boolable):
        return a if cond else b
    else:
        return cond.select(a, b)

Option B: overloadable conditional expressions

Allow:

a if cond else b

to be customized by the condition type. Maybe an __if__ method on the condition type?

This would allow the same source expression to naturally support both scalar and SIMD mask conditions.

I’m new to Mojo and this is my first forum post, so apologies if there’s already an established solution or prior discussion that I missed.

2

I didn’t see a way to be generic for the path of using select on the condition. So, this is the best I found for the lib solution:

def select[C: Boolable, T: Copyable](condition: C, yes: T, no: T) -> T:

    return yes.copy() if condition else no.copy()

def select[dtype: DType, width: SIMDLength](

    condition: SIMD[DType.bool, width],

    yes: SIMD[dtype, width],

    no: SIMD[dtype, width],

) -> SIMD[dtype, width]:

    return condition.select(yes, no)
3

fyi the stdlib already has math.select(cond, a, b) for the SIMD side of this, so your option A is kind of half done. whats missing is just the scalar overload for Boolable conditions, which would be a pretty nice small PR. welcome to the forum btw

4

Thanks! I missed that free function in math. But now I see that SIMD doesn’t behave like numpy in terms of comparison operators. Seems like comparison operators reduce and/or only work for 1 element. So, what I proposed isn’t that useful, I think.

IMHO, it would have been nicer if all SIMD operators that would normally apply to scalar were elementwise. Explicit reduction more like numpy. i.e. these which are currently not element-wise:

  • Comparison: ==, !=, <, >, <=, >, >= (so, not today’s Comparable)
  • Conversion: int, float (so, not today’s Intable or Floatable - maintains SIMD length)
  • if-else/select: proposal here.

OK, not quite all. A few exceptions like __hash__ would need to be scalar.

Only after you have unification of comparison would it be useful to think about the proposal here. Also, you’d need a way to allow for SIMD-like, normal scalar type, etc. And, some way to make sure SIMD lengths are the same on this generic thing (I only see a way with actual SIMD struct). Only then, could write scalar code without data-conditional statements that might work for SIMD-like types. My proposal was just one of those steps.

5

This was intentional. After some discussion, we decided that there is no good “least surprising” between SIMD.__eq__ producing a mask (as a SIMD programmer would expect) and it producing a boolean (as almost everyone else would expect). As a result, SIMD has SIMD.eq, SIMD.leq, etc which perform various elementwise comparisons, it simply doesn’t implement __leq__, and SIMD.__eq__ is a reduction.