Invoking raising functions in parameter expressions

Mojo
,
1

Interested in thinking about how the type system could be extended to support raising functions in type parameter computations. Also curious about the relationship between comptime parameter expressions and other static systems within the compiler (and external tooling like LSPs).

Below, an example where raising within a parameter expression might be useful.

Example: array broadcasting with compile time checking

Suppose one wanted to define a array broadcasting system which checked shape compatibility at comptime. One might try the following:

alias Shape = List[Int]


@value
struct Rep[dtype: DType, shape: Shape]:
    pass

# Broadcasting of arrays may fail statically!
fn broadcast(s1: Shape, s2: Shape) raises -> Shape:
    # imagine we check for compatibility here ...
    # if compatibility check fails, we'd like to raise
    return s1

# Now, say we wanted to add arrays with broadcasting turned on.
# The output shape depends on the input shapes (it will be 1 of the input shapes)
# but broadcasting may fail due to shape incompatibility.
fn add[
    D1: DType, D2: DType, S1: Shape, S2: Shape
](self: Rep[D1, S1], v: Rep[D2, S2],) -> Rep[D1, broadcast(S1, S2)]: # problem! raising in parameter expression
    return Rep[D1, broadcast(S1, S2)]()


def main():
    var r1 = Rep[DType.float32, [1, 2]]()
    var r2 = Rep[DType.float32, [1, 2, 3]]()
    var c = add(r1, r2)

The general idea here being that one might want to compute and check static information as part of the computation of a parameter expression, and possibly fail.

Above: broadcasting (following NumPy rules: Broadcasting — NumPy v2.2 Manual) is something that one could check at comptime.

Interaction with existing proposals

It’s not clear to me yet how this feature interacts with other proposed features like modular/mojo/proposals/requires-clause.md at main · modular/modular · GitHub (and likely, discussing that proposal should be another topic). If we could all have our language design utopias, I’m of the opinion that a comptime system which allows evaluation of the full language via the comptime interpreter would be an ideal to strive for! Then one doesn’t introduce feature biformity between comptime and runtime. This is obviously an ideal (!) and may have to be tarnished for e.g. fast compilation.

It is also not clear to me how compile time parameter expression evaluation should interact with useful tooling like LSPs. For instance, it seems beneficial if an LSP could surface a violation of broadcasting in a system such as the one above! But it’s not clear to me how this works in the Mojo compiler today (parser vs. elaborator vs. comptime vs. runtime) – or how to usefully think about this.

2

Thank you for creating this thread McCoy.

There are three different phases of evaluation of static information in the Mojo compiler:

  1. Parser time. This is when basic type checking and overload resolution happens. This phase is where the compiler can generate good error messages because it has full location information (including source range information) and all checks are purely local - so you don’t need a stack trace to explain things.

  2. Early dataflow checks. Checks like use-of-uninitialized variables and unused-store checking happens right after the parser in an MLIR pass. These give generally good error messages because they are function local and the pass in question has the info it needs to explain them.

  3. “Elaborator” or “instantiation” errors like failed constrain clauses. This checking happens later in the pipeline the comptime interpreter is run and function bodies are cloned based on the parameters that are passed into them, @parameter if’s are evaluated etc. The error messages from this are necessarily evil, because they need to include a stack trace like C++ template error messages.

It is important to know that there is no comptime interpreter in the parser - this is an inherent and intentional limitation. We want the output of the parser to be target independent, this is critical for our heterogenous compute goals, so we can’t fold things like sizeof in the parser, which means that any parser-time interpreter would necessarily be limited (keep in mind that Int is target specific size, so shift right has target specific behavior).

Generally speaking, we want to push errors up this list for quality reasons. The requires clause is specifically designed to replace some constrained clauses (#3) by pushing them into a form that can be checked by the parser (#1). This allows it to inform overload resolution, and makes the errors better.

What this means is that the type checker handles function calls lexically in phase 1:

struct TypeWithParam[a: Int]: ...
fn get_one() -> Int: return 1

var a: TypeWithParam[1]
var b: TypeWithParam[1] # Same type as a

var c: TypeWithParam[get_one()] # Parser sees a different type from b.
var d: TypeWithParam[get_one()] # Parser sees the same type as c

At parse time, Mojo don’t see into functions like get_one(), so it doesn’t know that it returns 1. In fact, this could be defined in a different package, and the author of the function might want to change what it returns, so it is good that mojo doesn’t just aggressively look through functions. Practically speaking, it also can’t, because most functions are way more complicated than this. This is unavoidable, but can also lead to annoying type incompatibilities: the solution to this is the rebind function in the standard library. rebind[TypeWithParam[1]](c) tells the Mojo parser “Look, I know that c has the same type as TypeWithParam[1] after elaboration, just go with it”. The types are checked by the elaborator so if they differ, it will explode with a stack trace.

There is one more thing to know in this space, Mojo has a feature called always_inline("builtin") which DOES force inline the body, which can allow some “obvious” things to work:

var e: TypeWithParam[2-1] # Same type as a.

In this case, minus on IntLiteral is always_inline("builtin") so Mojo is able to see into it, and able to do constant folding. This is an exception rather than a rule though. This functionality is extremely limited by design, we don’t want to rely on it globally.

Given that as background, let’s get back to your idea. We could allow this:

fn something() raises -> Int: ...

var f: TypeWithParam[something()]

However, we can’t know what the function call does until elaboration, which happens well after type checking and overload resolution is done, so it can’t influence overload resolution. We’d have two options to consider in how to handle this at elaboration time if an error is thrown: 1) we can just fail to compile the program and throw a stack trace. 2) We can reject it entirely as we do now.

#1 could be useful in some cases, but I think a better solution would be for us to add the forthcoming requires clause to see how much of the use-case it covers. If you really want this behavior, you can work around this by doing this for now:

fn something() raises -> Int: ...

fn something_wrapper() -> Int: 
   try:
      something()
   except:
      abort()

var g: TypeWithParam[something_wrapper()]

-Chris

7 Likes
3

@clattner Neat, thanks for the detailed response – clarifying!

Do you have any thoughts about the interaction between the elaborator and language server? Pragmatically, I’d expect the today answer to be “no need or desire for interaction” (makes a ton of sense for lots of practical reasons) – and possibly most of the useful cases would be covered by a requires like feature at parse time (cases which get “best of all worlds” solutions).

I indulge in fantasy, especially when it comes to languages – so I’ll indulge in a bit below the line (what if the answer was: “it might be interesting to have some interaction?”)

A few sketchy thoughts, predicated on following the fork (1) in

We’d have two options to consider

Presumably, the language server today doesn’t interact with stage (2) and (3) – because maybe today that requires invoking the full compiler pipeline? That’s a practical design: you do not want to be running the full compiler every time you write code :sweat_smile: (for obvious reasons)

On the other hand, dependently typed or interactive systems do approach this question: they expose systems which are analogous to the comptime interpreter (like Lean4 – where the LSP is constantly running an elaborator).

I think the way that they do this is by incrementalizing elaboration or comptime code execution – a fascinating topic (but of course, a topic for research, not production).

I do wonder (fantasy) what the language might look like if this fork was embraced. Comptime execution systems often look like some form of dependent types, and I’m of the (strong opinion, held loosely) that dependently typed design reduces feature biformity in static systems.

requires strikes me as a bit of biformity, as presumably one has to learn a restricted language (the set of things one can do at parse stage within requires). I’m not at all an extremist on this stuff, and there’s plenty of practical reasons to take fork (2) regardless – but I wondered if you felt a similar way about biformity here?

1 Like
4

I agree that it might be better to try to make the elaborator run in the LSP. Zig had what you could call a lack of that, since the LSP basically didn’t handle comptime, and it led to a horrible developer experience. I know that might be a substantial bump to the CPU usage of the LSP, but I think that it’s a worthwhile tradeoff (granted, I have CPU to spare). Especially as we look towards reflection, I can’t think of a good way to run that inside of what the parser currently expects.

Also, if the LSP is going to get elaboration capabilities, I think that having an “I’m compiling for this cpu/accelerator combo” option might be helpful, since then it could actually elaborate things and display feedback on @parameter if statements similar to what clangd does with areas disabled by the C preprocessor, but a “portable default” would also be useful for library authors.

1 Like
5

Currently the LSP supports phase 1 and 2, but not phase 3. That’s a pretty good thing for compile time responsiveness / performance of the LSP. I think that pushing things to parse time (eg requires) is the way to go for usability for lots and lots of reasons.

-Chris

6

This topic was automatically closed 7 days after the last reply. New replies are no longer allowed.