Proposal for Returning Relevant Cursors

[SamirDroubi]

Problem:

We have been talking about having scheduling ops support the idea of "Returning relevant cursors". For example, when using stage_mem, we generate an allocation (which we might wanna further schedule to for example set the memory), and at least one of a load stage or store stage:

def foo(n: size, A: i8[n] @ DRAM):
    assert n % 4 == 0
    for io in seq(0, n / 4):
        for ii in seq(0, 4):
            A[4 * io + ii] += 0.2

foo = stage_mem(foo, "for ii in _:_", "A[io*4:io*4+4]", "tile")

def foo(n: size, A: i8[n] @ DRAM):
    assert n % 4 == 0
    for io in seq(0, n / 4):
        tile: i8[4] @ DRAM      <----- Allocation
        for i0 in seq(0, 4):        <----- Load Stage
            tile[i0] = A[i0 + 4 * io]
        for ii in seq(0, 4):
            tile[ii] += 0.2
        for i0 in seq(0, 4):        <------ Store Stage
            A[i0 + 4 * io] = tile[i0]

It would be convenient if scheduling operations returned cursors to those relevant pieces of the procedure. Currently the alternative to this is using the argument cursor as an anchor and using navigation to get to such cursors from the scheduling code. The goal is to offload that work from the user to the tool.

Proposals:
Proposal 1: Return a dataclass algonside the proc. The dataclass contains fields which are the relevant cursors.

Definition goes along with the scheduling operations definition, serves as a documentation of what can be returned:

@dataclass
class stage_mem_cursors:
        alloc: AllocCursor = InvalidCursor()
        load_stage: AssignCursor | LoopCursor = InvalidCursor()
        store_stage: AssignCursor | Reduce_Crursor | LoopCursor = InvalidCursor()

Usage:

foo, cursors = stage_mem(foo, "for ii in _:_", "A[io*4:io*4+4]", "tile")
load_stage = cursors.load_stage # Save for later
foo, cursors = set_memory(foo, cursors.alloc, DRAM_STATIC)

Advantages:

1. Explicit API
2. Uniform return values convention (Procedure, dataclass of relevant cursors)
3. It is easy to be forward compatiable and add more cursors in the future

Disatvantages:

1. Have to define a dataclass per scheduling op
2. Might be too verbose: cursors.load_stage
3. Users may need to name for each returned dataclass to save them as temporary varaibles. Alternatively, overwrite previous one just like we do for proedures; you need to manually save which cursors you actually care about. So, the advantage your really save is not having to do a lot of navigation which I think is worth it.
4. Unclear what to return when there is nothing to return. Some "Null" dataclass.

Proposal 2:

Similar to proposal 1, but we keep the return value of primitives to be just a procedure. Then, we implement a standard library wrapper around each primitive that returns a procedure and a dataclass of relative cursors.

Advantages:

1. Keeping the primitives implementation as small as possible.

Disatvantages:

1. Duplicate shceduling operations that perform the same rewrite.
2. Naming problems:
   • Use same name
   • or add a suffix to the wrappers e.g. _c.

Proposal 3:
Instead of returning multiple relevant cursor, we can return exactly on relevant cursor along with the procedure. In the case of stage_mem, this could be the new block cursor including the load stage, the original block we staged around, and the store stage.

• There is again the question here of whether this would a part of the primitive or a standard library wrapper.

Advantages:

1. Might be easier to use since you don't have to think about the name bindings (in a dataclass or a dictionary) anymore.

Disatvantages:

1. Returning one curosor is not alawys a viable options when the relevant cursor are not co-located i.e. they cannot be all grouped in one block. It is esepcially difficult with operations that could have an aribitrary number relevant cursors in arbitrary places in the procedure.
2. We will still require a fair amount of navigation in the user-code.

[SamirDroubi]

@yamaguchi1024 @skeqiqevian. Can you add any other propsals you had in mind?

[gilbo]

Some thoughts:

1. I think there are a few semi-orthogonal sub-issues here that would be worth breaking out when thinking about the design. In particular, I am thinking of (a) what is the return value (just proc, proc + dataclass, proc + one cursor); (b) how should the "return cursors" functionality be exposed? (e.g. should this be an alternate version of a primitive or should the expected signature of a scheduling operation primitive change?); (c) slightly different is "how should this be implemented?" (e.g. purely in user-library vs. changing the primitives' implementation?)

2. Another option (for dimension (b)) is to support a flag (e.g. p, cs = sched_op(p, c, ..., return_cursors=True)) that changes the returned value. If you think of this functionality as "advanced" then a little bit of extra typing might be appropriate.

3. I like that the dataclass idea can be combined with lvalue pattern matching to unpack the different cursors "positionally" e.g. p, (alloc, load, store) = stage_mem(...) or unpack the cursors by name as you suggested @SamirDroubi. In the event that a scheduling op naturally wants to return only one or zero cursors, what would we do? Still return a dataclass for uniformity of the interface? (That doesn't seem like a terrible idea...?)

4. The "return one cursor" variant is probably important for a combinator user library along the lines of the one we sketched in our PLDI submission. However, I don't think that option is either maximally simple (no cursor returned) or maximally useful/informative (the multiple cursor return). One could imagine e.g. a new combinator that allows one to wrap an op and specify which of the dataclass entries should be returned as "the" cursor. Bottom line is we can implement better "one cursor" convention libraries on top of the other options.

5. The second foo code (post stage_mem) has a bunch of arrows pointing to parts of it. This reminds me of the "pattern matching multiple cursor groups" idea we floated as a more self-documenting introspection mechanism. Suppose we added some such feature. We would have to return a set of cursors in that case; one for each group. Shouldn't that ultimately use the same convention that we decide on here? If so, maybe thinking about that feature could help us feel more confident about a decision here (e.g. exactly what something like the dataclass idea should look like)

[SamirDroubi]

1. I think splitting the problem into thoss three dimensions makes sense and is helpful to think more clearly about it.
2. If you don't care about the returned cursors, you could simply ignore it by writing p, _ = sched(...). On the other hand, having to specify return_cursros=True is more verbose (and more characters to write). I guess ultimately it would depend on how frequently you would need to specify that argument. I would say the nice thing about having this option is someone who doesn't want to think about cursors, can completely operate as if they didn't exist.
3. Unpacking should be simple with dataclasses, we just need to make sure to make the dataclass iterable by implementing __iter__. I am on the fence of the case where there is exactly one cursor because just returning the cursor wouldn't change the number of returned values, just the type of the second one. When there is no cursors to return, I think it would be nice to have uniformity at least in terms of number of returned arguments. Although it is a bit weird that every use of some op is of the form: p, _ = sched(...).
4. I thought about the combinators and my thinking is what you likely do is build lambdas on the fly. The one returned cursor by the lambda might likely be one of the many returned cursors by the scheduling operation. The other way is what you described which is to have a wrapping combinator: we can make the dataclasses indexable by the order in which the fields are defined, and then we can build the combinator you described to generate a version that returns exactly one cursor given an index.
5. I think we should be able to also define dataclasses dynamically for this feature as well and build the same functionality: https://docs.python.org/3/library/dataclasses.html#dataclasses.make_dataclass.

[gilbo]

Continuing the conversation where we're not converged:

2. yes, p,_ = op(p, ...) isn't unreasonable. I think the argument from the other side is that it's conceptually another spinning wheel. There was an old proposal which would simplify this to op(...) at the cost of not having as much explicit control. Absent other considerations, we should try to avoid the situation where there are a bunch of extra parameters/return-values that always have to be handled because such designs are more complicated to learn and tend to produce more syntactic boilerplate. If the return_cursors is an "advanced" feature, you can imagine it being as simple as adding op(..., rc=1) (using "truthiness" instead of True) or something similarly brief.
3. (a) You can get dataclass to autogenerate the __iter__ implementation.
4. (b) One proposal going with rc=1 style is that if it returns 0 cursors, then p,_ = op(..., rc=1) makes sense even if it is unnecessary. I think the awkward case is p,(c,) = op(..., rc=1) but maybe that's OK given that this is invoking the advanced feature?
5. You can use a string with the dataclasses or an index. Effectively, we would have a system with a few different standard signature conventions for scheduling operators. e.g. Basic is p = op(p, ...); Basic at location is p = op(p, c, ...); Combinator is p,c = op(p,c, ...); Full Navigation is p,(c1,...) = op(p,c, ...)
6. Yes.

An additional thought (not to roll into this proposal but thinking further ahead) is that it would be great if there was a "standard" way to decorate some user-defined scheduling ops in order to encourage adhering to a convention and maybe generate some boilerplate functionality as a consequence. There could be different conventions (as suggested above). You could imagine such a decorator allowing us to install different kinds of cursor coercion functionality (already exists for primitives, but not user ops), maybe help with documentation, and maybe package rc=1 optional functionality etc. You could also imagine auto-coercion between some different operator conventions.

[SamirDroubi]

2. Yeah, I agree with you that there should be a supported scheduling mode that allows that user to ignore anything cursors related. I think the consise arguemnt you described makes sense, it is not clear what it does from looking at it, but it is in every operation and will be common in code that users will just learn.
3. Agreed.
4. I agree it should be okay to expect the user to unpack the dataclass in that one off case.
5. Agreed.
6. I am not sure exactly what you mean by auto-coersion, it would be great if you could elaborate more on this.
7. I was generally thinking that we would extend the @sched_op decorator so that it automatically adds the new default argument rc=False.
8. We can potentially make the dataclass be automatically generated and associated with a scheduling op via the same decorator or a different one: @returned_cursors("alloc", "load", "block", "store") which then would serve as a docuemntation. If we go with the two decorators apporach, then the sched_op decorator should likely enforce the returned_cursors decorator.
9. How existing and new decorators can be exported to the user world is really important discussion as well. The current decorator and argument processors can greatly simplify the user ops implementation and provide easy documentation. I think simply exporting them as is may not be a terrible idea since the almost all return API cursors. There are some exceptions (NewExprA) where the return value is a non-API object. One option is to not export any thing that could return a non-API object for now and make the user-level argument signature be a string instead which is what you would do now anyways.
10. Another remaining point of the discussion is (1.c). How should this be implemented?

[gilbo]

Re auto-coercion, note that the Argument Processors do this. That is one concept of coercion. Another concept concerns totally standard liftings between different SchedOp signatures. eg an op with signature (p,c,…) -> p can be lifted to an op with signature (p,c,…) -> (p,c) implicitly (a coercion) whenever supplied as input to `repeat` or a similar combinator.The argument processors are currently a bit tricky to disentangle from the compiler implementation into a user-level concept, which is why I’m suggesting turning that (and more general coercions, op decorators, etc) into one or more future proposals.— GilbertOn Dec 26, 2023, at 5:36 PM, Samir Droubi ***@***.***> wrote: Yeah, I agree with you that there should be a supported scheduling mode that allows that user to ignore anything cursors related. I think the consise arguemnt you described makes sense, it is not clear what it does from looking at it, but it is in every operation and will be common in code that users will just learn. Agreed. I agree it should be okay to expect the user to unpack the dataclass in that one off case. Agreed. I am not sure exactly what you mean by auto-coersion, it would be great if you could elaborate more on this. I was generally thinking that we would extend the @sched_op decorator so that it automatically adds the new default argument rc=1. We can potentially make the dataclass be automatically generated and associated with a scheduling op via the same decorator or a different one: @returned_cursors("alloc", "load", "block", "store") which then would serve as a docuemntation. If we go with the two decorators apporach, then the sched_op decorator should likely enforce the returned_cursors decorator. How existing and new decorators can be exported to the user world is really important discussion as well. The current decorator and argument processors can greatly simplify the user ops implementation and provide easy documentation. I think simply exporting them as is may not be a terrible idea since the almost all return API cursors. There are some exceptions (NewExprA) where the return value is a non-API object. One option is to not export any thing that could return a non-API object for now and make the user-level argument signature be a string instead which is what you would do now anyways. Another remaining point of the discussion is (1.c). How should this be implemented? —Reply to this email directly, view it on GitHub, or unsubscribe.You are receiving this because you commented.Message ID: ***@***.***>

[SamirDroubi]

To summarize, the design we arrived at above:
a. Relevant cursors will be returned as a dataclass.
b. Using this functionality will require opting-in by specifcying a default parameter rc=True (return cursors) which will be the last parameter of a scheduling op.
c. This functionlity will be added to the primitives themsevles.
d. Cases where the are 0 or 1 relevant cursors will still return a dataclass when the the functionality is inovked for consisntency.
e. The per-op dataclass will be automatically defined by a decorator or explicitly.
f. The default argument will either be enforced or automatically added by a decorator.

I think (e) and (f) present a new design decision:
Option 1:

1. Require each scheduling operation to have an extra final argument **kwargs
2. We modify the argument processor to enforce what goes in that. Currently, we only allow rc with boolean values.
3. We can add a decorator that dynamically defines the dataclass e.g. (@returned_cursors("alloc", "load", "block", "store")). This decorator binds the type of the dataclass in **kwargs

Option 2:

1. Require each operation to have the last argument be rc similarly to how we enforce the first argument to be a procedure.
2. To define the dataclass:
   a) Statically (explicitly) define a dataclass in the globalscope with some naming convention e.g. "{op_name}_rc"
   b) Use a decorator @returned_cursors("alloc", "load", "block", "store") to the function that binds the dynamically-created definition of the dataclass to the globalscope under some naming convention.

I think Option 2 with (2.b) is the cleanest and most explicit. The rc argument will be explicit in the operation defintion. It exposes the dataclass definition to the users, so they can introspect it with help(). The information about what cursors are returned is clearly coupled with the scheduling operation via the decorator and we can add them to the information displayed when a user calls help() on a scheduling operation.

[gilbo]

Can we just chop off e & f from this proposal and fold those discussions into a subsequent proposal? (focused on how to better support user-defined scheduling operations using decorators)

…

On Dec 28, 2023, at 9:18 AM, Samir Droubi ***@***.***> wrote: To summarize, the design we arrived at above: a. Relevant cursors will be returned as a dataclass. b. Using this functionality will require opting-in by specifcying a default parameter rc=True (return cursors) which will be the last parameter of a scheduling op. c. This functionlity will be added to the primitives themsevles. d. Cases where the are 0 or 1 relevant cursors will still return a dataclass when the the functionality is inovked for consisntency. e. The per-op dataclass will be automatically defined by a decorator or explicitly. f. The default argument will either be enforced or automatically added by a decorator. I think (e) and (f) present a new design decision: Option 1: Require each scheduling operation to have an extra final argument **kwargs We modify the argument processor to enforce what goes in that. Currently, we only allow rc with boolean values. We can add a decorator that dynamically defines the dataclass e.g. ***@***.***_cursors("alloc", "load", "block", "store")). This decorator binds the type of the dataclass in **kwargs Option 2: Require each operation to have the last argument be rc similarly to how we enforce the first argument to be a procedure. To define the dataclass: a) Statically (explicitly) define a dataclass in the globalscope with some naming convention e.g. "{op_name}_rc" b) Use a decorator @returned_cursors("alloc", "load", "block", "store") to the function that binds the dynamically-created definition of the dataclass to the globalscope under some naming convention. I think Option 2 with (2.b) is the cleanest and most explicit. The rc argument will be explicit in the operation defintion. It exposes the dataclass definition to the users, so they can introspect it with help(). The information about what cursors are returned is clearly coupled with the scheduling operation via the decorator and we can add them to the information displayed when a user calls help() on a scheduling operation. — Reply to this email directly, view it on GitHub <#541 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAKZHLBUJTZ2SD2DYSFBC6TYLWSYHAVCNFSM6AAAAABA2NWFC2VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTQNZRGM2TQNBXGE>. You are receiving this because you commented.

[SamirDroubi]

(e) and (f) were referring to the primitives here. I think we need to have something if we are to implement this feature for the primitives. The other alternative is to hold off on implementing this feature until we have the second discussion about how decorators are exported by the API.

[gilbo]

Primitives are going to have to modify their implementation (in the `API_Scheduling.py` file) to specify _what_ the cursor values are. Note that all primitive scheduling operations are instances of the class `AtomicSchedulingOp`, so that (and the ***@***.***_op` decorator) provides an existing centralized location to change the behavior of all primitives. The `rc` kwarg can be intercepted in `AtomicSchedulingOp.__call__`. Presumably, we would then modify the functions later in that file to return a data class of cursors by default and optionally suppress that return value in `AtomicSchedulingOp.__call__`. Then, we need to figure out how to specify the dataclasses in question. In general, these dataclasses could be defined above the corresponding scheduling op. They then need to be user-exposed in `stdlib/scheduling.py` alongside the corresponding scheduling op. Another option could be to abbreviate these dataclass definitions as (e.g.) `Op1Cursors = MakeCursorSet(‘Op1Cursors’, ‘alloc’, ‘load’, ’store’)`. This will work if we don’t want to do anything clever with the type annotations for the cursors (i.e. they can all just be generic cursors). I think that accounts for almost all the issues. Have a look at the code and see if what I’m saying makes sense.

…

On Dec 28, 2023, at 1:59 PM, Samir Droubi ***@***.***> wrote: (e) and (f) were referring to the primitives here. I think we need to have something if we are to implement this feature for the primitives. The other alternative is to hold off on implementing this feature until we have the second discussion about how decorators are exported by the API. — Reply to this email directly, view it on GitHub <#541 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAKZHLCXL56RYXVQIMOC66TYLXTTXAVCNFSM6AAAAABA2NWFC2VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTQNZRGUYTKMBRGA>. You are receiving this because you commented.

[SamirDroubi]

Yeah that makes sense. That's a cool idea, I didn't think about making the implementation always return the value by default, then optionally supress it in the object. Thanks!