RFC-0005: Structured kernel definitions RFC

[ezyang]

This is a proposal for a new code generation facility for writing kernels in PyTorch, where we will automatically generate easy-to-get-wrong boilerplate for functional (add), inplace (add_) and out (add_out) variants of functions, as well as common code (device guards, version counter tracking). The net result is you only need to write a shape checking function and an out-kernel when writing a function.

Rendered

[bhosmer]

maybe worth noting somewhere in here that you're limiting the example to CPU and CUDA for expository purposes? the use of explicit dispatch keys above should probably imply to the careful reader that this setup is parameterized over all backend dispatch keys in the usual way (IIUC), but if so probably useful to make it explicit

[bhosmer]

I think it's worth thinking carefully about whether this is the best way to make structured kernels available out of tree, given the modality it introduces, risk of drift, perf handicap (albeit diminishing over time) it saddles out of tree backends with, etc.

E.g. once the rest of the codegen has been ported, it's not absurd to imagine fitting it with a frontend that can take inputs other than the in-tree yaml...

[ezyang]

The common dispatch key is primarily oriented towards backend implementers.

Let's suppose for a moment, that an overriding design goal is safety/correctness first, with the ability to opt into performance. Then it feels feel like I am forced to introduce the Common dispatch key, because without it, an average backend overrider has to faithfully replicate all of the functionality that we have otherwise fused into CPU/CUDA operators (shape checking, device guards, and with some refactors coming soon, it will also include version counter bumps). It's unrealistic to expect a backend implementer to actually manage all of this without code generation.

That being said, there is a certain optionality to the Common dispatch key. We don't have to implement it (and indeed, in the current posted PR, it is not implemented), and if it is not implemented, the burden is simply on backend implementers to implement all of the necessary scaffolding (which is the de facto situation today). If, for example, we pivoted to publishing code generation for backend implementers, that would alleviate or perhaps eliminate the need for a common dispatch key.

The common dispatch key is mostly irrelevant for custom operators, since average use of common operators is via catch all registration and there will not be any common dispatch key in any case.

[bhosmer]

Yeah, agree on all the motivations. My pitch would be to (wait and) audition a lightly parameterized codegen pipeline as an alternative, once the rewrite is complete. (Waiting would serve not just to let the codegen settle down, but also give us more time to see how quickly dispatcher overhead was coming down.)

[smessmer]

I think some of the weirdness like making native_functions.yaml entries depend on each other or having two ways of doing things (dispatch key and fusing into kernels) is because you designed with the constraint of not changing the fact that the three operator overloads (regular, inplace, out) are actually registered as separate operators to the dispatcher.

Have we considered an alternative solution where the dispatcher only knows about the out overload and the regular and inplace variants are generated in the frontend before the call to the dispatcher? That would also make it extensible for out-of-tree ops and take the burden away from backend implementors.

[smessmer]

This means we are using this mechanism for all ops, even ops that don't have out or inplace variants, right? Otherwise those ops wouldn't get device guards or version counter bumps.

[ezyang]

Yeah, this is an orthogonal change that we can do with or without structured kernels.

[smessmer]

As long as we use this mechanism internally for all ops, we have meta functions for them defined and backends adding new kernels don't need to care about it. Only backends that add new operators would have to use it. Did I understand this correctly?

[ezyang]

That's right. So if you don't care about performance, you don't even have to lift a finger.

[smessmer]

wait I think I misunderstood this. To get device guards, you'll still have to go through the dispatch key for those backends, but not if calling into internal backends like CPU. How do you selectively make that dispatch key fallthrough for some backends but not for others?

[ezyang]

I claimed this would work in our meeting, but rereading your comment, I take it back, you really do need a per backend key here :(

[smessmer]

This dual approach makes the system a bit complex, it can be hard to figure out where version counter bumps, device guards or shape checking happen for a corresponding op, especially when you're not familiar with this system (yet).

[smessmer]

This implicit connection makes reasoning about entries in native_functions.yaml harder since they're not independent entries anymore.

[smessmer]

hm maybe we should make this explicit by also adding structured: True here, to make the information more local and easier to read. Otherwise you have to look at other entries in native_functions.yaml to know what this entry is actually doing. Omitting dispatch keys is already a valid kernel definition even if it's not structured, so by just looking at this one it would be ambigious.

[smessmer]

these should be easy to generate with templates instead and it should be possible to do that in a readable way

[ezyang]

The current expressivity problems:

1. Guard logic is weird and special-casey in codegen right now. We probably should be able to simplify it a bit, but right now it would be quite difficult to faithfully replicate the logic in a template
2. The upcoming version counter bumps will be difficult to do without out-of-band information about what arguments are mutable or not (right now you can check this using Tensor& but when we fix everything to uniformly be const Tensor& you'll lose this type info)

[smessmer]

We could offer a simple metaprogram for them so they don't have to manually write this for each op. But if I understand correctly, they only have to write this for ops they introduce right? Ops from native_functions.yaml that the backend only extends will already have shape checking etc. through our codegen right?

[ezyang]

Yes

[smessmer]

I don't understand the connection to mobile selective build yet

[ezyang]

We can talk about this in the meeting. It's constraint solving from the problem "mobile requires registrations to be a separate compilation unit from kernels"

[salexspb]

This is very exciting! Looking forward having this structure.
Ability to infer shapes statically and pre-allocate all the memory to be passed on later to _out kernel versions is going to bring PyTorch runtime to the next level, I think.

Do you guys plan to support all GPU standard kernels using this new framework? (so they have shape inference and _out versions).
And what about custom ops / classes, do you plan to support them in the kernel definition framework as well? Ideally I would like to be able to have 100% static models via using a mix of built-in _out ops versions and custom ops (also accepting external memory and providing shape checking / inference functionality) one may implement.

[ezyang]

Do you guys plan to support all GPU standard kernels using this new framework? (so they have shape inference and _out versions).

Yeah. The aspiration is every kernel in PyTorch is in this framework. That's gonna be a lot of work, but hopefully we can lay the groundwork and then roll it out as we go.

And what about custom ops / classes, do you plan to support them in the kernel definition framework as well?

This is not entirely settled yet, but the intention for the section at #9 (comment) was to make this possible. So if you are doing a custom op, you now define (e.g.) three parts: a CPU part, a CUDA part, and the static shape checking part, and the framework would put it all together for you.

One extra thing that I'd add, though, is when I've been chatting with other static runtime people at FB, they seem to want strange things like being able to run the static inference really fast at the beginning of each run to work out the preallocation. At least the first iteration, the API for actually running these shape computations won't be particular fast, and will be mostly useful for offline use cases.

[ezyang]

The *this is deeply weird until you know the context this code is being defined in... it makes me think we might want the macro names to avoid overselling that it's a freestanding function. It wouldn't need to literally say "you're looking at a method" but like TORCH_STRUCTURED_IMPL or something.

Fair enough. We can use your naming!

[bhosmer]

I think you mean the const modifier can be added, not eliminated, right? Like, out arguments can now be typed as const Tensor& in TORCH_IMPL_FUNC declarations, as in

TORCH_IMPL_FUNC(upsample_nearest1d_out_cpu) (
    const Tensor& input,
    IntArrayRef output_size,
    c10::optional<double> scales,
    const Tensor& output             // currently just Tensor& output
) {
  upsample_nearest1d_kernel(kCPU, output, input, scales);
}

Aside: it happens that upsample_nearest1d_kernel itself takes Tensor& output, which might hint at additional porting work for at least some kernels to adopt this change.

[bhosmer]

upcast

[bhosmer]

at::cpu seems much less risky around the edges than CPUTensor[1] - if they serve near equivalent purposes, what's an example that motivates adding CPUTensor either instead of or in addition to at::cpu?

[1] stemming from CPUTensor not reeeally being a subtype of Tensor. A stray call to something like set_storage_and_dtype would damage a CPUTensor that's been upcast to a Tensor (or provoke a runtime error, if we've put safety measures in place).

[ezyang]

It is an ergonomics thing (but important ergonomics): with at::cpu, you have to remember to use the CPU only function (and you have to make sure not to accidentally call it with a non-CPU tensor). CPUTensor ensures that the CPU tensor is used, and you only need to prove you have CPU tensor when you initially construct the CPUTensor.

[bhosmer]

A counterargument to consider here I think is that making structured kernels ubiquitous would allow other system capabilities to be defined in terms of them - e.g., shape analysis, memory planning.

The answer wouldn't be to torture everything into having in/func/out variants, but probably to add a handful of other little composite kernel system variations, each with its own codegen and triggering annotation in native_functions.yaml. Factory functions are an obvious fit for this kind of approach, but maybe they're the only good fit? Haven't thought it through any further.

[ezyang]

Well, this certainly would work. What is less certain is if the number of operators that would be covered by the system variation warrant the variations. Well, at least we can delay designing this for later...