Functional alias to Base.Cartesian.@nif

[MilesCranmer]

You can use ntuple(f, Val(n)) as an alias of @ntuple(n, f). The reason this is preferable is because @ntuple requires n to be known at parse time, while ntuple(f, Val(n)) only requires this information at compile time. Therefore you can have n depend on types of arguments when using the functional alias, avoiding the need for @generated functions in package code.

However, there is no analog of Base.Cartesian.@nif, so you need to use @generated whenever you wish to have a type-dependent if statement.

This PR adds this alias as well as a test suite. Fixes #55087.

Use-cases

Some of my code in DynamicExpressions.jl needs to use @generated functions which roughly look like this:

@generated function foo(idx, operators::Tuple, data)
    N = length(operators)
    quote
        Base.Cartesian.@nif(
            $N,
            i -> i == idx,
            i -> eval_with(operators[i], data)
        )
    end
end

There's no way to do this without a @generated function, because N needs to be known to the macro at parse time. The code otherwise does not need to use @generated, and so likely incurs greater compilation effort than is actually needed.

The reason I need this if statement is so that the compiler is aware of exactly which operator is used at each branch – required for things for Enzyme to work. (If I just used getindex(t::Tuple, i::Int), the compiler would not be able to do constant propagation far enough).

So, rather than require users and package developers to use a @generated whenever they want to create a type-dependent @nif, this PR creates Base.Cartesian.nif which does the code generation internally, in an analogous way to how ntuple works. This is also nice because it has the non-@generated branch for improved inference in cases where the types are not fully known.

For a simple example, consider the following test (included in the test suite):

t = ("i am not an int", ntuple(d -> d, Val(10))...)

function extract_from_tuple(t::Tuple, i)
    Base.Cartesian.nif(
        d -> d == i,
        d -> t[d + 1],  # We skip the non-integer element
        Val(length(t) - 1)
    )
end

Normally, had we used getindex here, inference would have not been able to infer that the return type never includes the first element. But since we used an nif, the compiler knows all possible branches, and can inline d + 1, and therefore infer the correct type of every branch:

@test @inferred(extract_from_tuple(t, 3)) == 3

To solve this problem with Base.Cartesian.@nif, you would have to wrap the function body with @generated.

Tests

Base.Cartesian.@nif does not have any tests aside from the jldoctest. This PR also adds many tests of nif which therefore also test @nif.

[nsajko]

Is this meant for users or for internal use? That is, is the idea to make it public?

[nsajko]

As per #54544, it may make sense to replace the N::Int line with n = N::Int, and use n instead of N afterwards. This is because N::Int doesn't seem to help abstract type inference on its own, xref #38274.

[MilesCranmer]

Is this meant for users or for internal use? That is, is the idea to make it public?

I'm not sure. Currently, all of Base.Cartesian is not public: https://docs.julialang.org/en/v1/devdocs/cartesian/. However, ntuple which aliases Base.Cartesian.@ntuple, is exported.

So, I'm not sure what to do. Any advice? Should nif be marked public, and, if so, should nif be moved to Base.nif, rather than Base.Cartesian.nif?

Though regardless of Base.Cartesian being marked as internal, it does seem to get a lot of use: https://github.com/search?q=language:julia+/Base.Cartesian./&type=code (though oftentimes requiring @generated for use). Perhaps a good candidate for moving to public in the future at some point.

As per #54544, it may make sense to replace the N::Int line with n = N::Int, and use n instead of N afterwards. This is because N::Int doesn't seem to help abstract type inference on its own, xref #38274.

Will do

[MilesCranmer]

It seems like the compiler is too-often choosing the non-generated part of the function during testing... Which is breaking the inference test I wrote because the non-@generated if-statement over a tuple of varied types is not type stable. (The @generated one should be)

Any way to force @generated for the purposes of a unit-test? I could also just delete that test.

I could also remove the non-@generated branch.

[MilesCranmer]

Hm. It still seems to take the non-@generated branch even with that adjustment.

[MilesCranmer]

Ok, I got it. The fix was

-@inline function nif(condition, expression, ::Val{N}) where {N}
+@inline function nif(condition::F, expression::G, ::Val{N}) where {F,G,N}
     nif(condition, expression, expression, Val(N))
 end

[MilesCranmer]

No clue why the test is failing:

Test Failed at /Users/julia/.julia/scratchspaces/a66863c6-20e8-4ff4-8a62-49f30b1f605e/agent-cache/default-honeycrisp-XC9YQX9HH2.0/build/default-honeycrisp-XC9YQX9HH2-0/julialang/julia-master/julia-32100ec1d7/share/julia/test/cartesian.jl:588
  | Expression: nif(identity, identity, Val(-1))
  | Expected: ArgumentError("if statement length should be ≥ 0, got -1")
  | No exception thrown

that test works fine on my machine if I type it in the REPL. Could it be from the LazyString?

[MilesCranmer]

Argh. This seems to hurt inference quite a bit compared to using a normal Base.Cartesian.@nif inside a @generated function, with N interpolated into the @nif. I guess the result really depends on whether the compiler wants to inline or not?

Perhaps my functions are too complicated and the compiler avoids inlining:
https://github.com/SymbolicML/DynamicExpressions.jl/blob/67bfab09033d3192d79defb79e1c313c2da03900/src/Evaluate.jl#L302-L348

I am surprised because ntuple is used quite often, maybe it's because the function passed is quite simple, and the compiler always inlines it into the macro?

[nsajko]

Replacing the loop with recursion could fix the bad inference, I guess.

[MilesCranmer]

It seems even if I remove that branch, and have @generated covering the entire nif function, I still don't get the performance I had from just using Base.Cartesian.@nif directly. Not sure why. My guess is that it has to do with constant propagation through the anonymous functions passed to nif?