c10::KernelFunction by smessmer · Pull Request #26337 · pytorch/pytorch

smessmer · 2019-09-17T05:25:26Z

Stack from ghstack:

Move Generator ops to c10 #26434 Move Generator ops to c10
Enable creation of boxing wrappers for some aten operators #26273 Enable creation of boxing wrappers for some aten operators
Add two levels to use_c10_dispatcher #26272 Add two levels to use_c10_dispatcher
c10::KernelFunction #26337 c10::KernelFunction

Factor out boxing and unboxing functionality from the c10 dispatcher into a c10::KernelFunction class
Move that class and everything else it depends on into ATen/core/boxing
This also allows us to get rid of c10::KernelCache. Instead, we now store a pointer to the unboxed functor in c10::KernelFunction.
We're also getting rid of the DispatchTableEntry struct and instead store KernelFunction directly.
To make this work, we need to change the dispatcher calling API from Dispatcher::lookup().callBoxed/callUnboxed and OperatorEntry::lookup().callBoxed/callUnboxed to Dispatcher::callBoxed/callUnboxed and OperatorEntry::callBoxed/callUnboxed.

Differential Revision: D17416607

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

Pull Request resolved: #26337 ghstack-source-id: 90228146 Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/)

aten/src/ATen/core/dispatch/DispatchTable.h

aten/src/ATen/core/dispatch/KernelFunction.h

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

aten/src/ATen/core/dispatch/KernelFunction.h

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

dzhulgakov · 2019-09-19T06:18:31Z

aten/src/ATen/core/boxing/KernelFunction.h

+namespace detail {
+template<class Type>
+constexpr uint64_t hashType() {
+  return typeid(Type).hash_code()


did you verify (diasm in gdb or something) that it indeed becomes a constexpr value?

aten/src/ATen/core/boxing/KernelFunction.h

dzhulgakov · 2019-09-19T06:30:55Z

aten/src/ATen/core/boxing/KernelFunction.h

+  // We use this indirection because many KernelFunctions are created
+  // at static initialization time but are created with functors that
+  // store Tensor and we can't call the Tensor() constructor at static
+  // initialization time yet (SIOF). So these register with a


Even in C2 world it's not true - we wouldn't create an instance of operator at registration time and lifetime of operator would be managed by the invocation context, not centralized dispatcher. I honestly don't think there's a usecase for kernelFactory_ today (and in the earlier design it was not fully figured out either). Can't we just nuke the functorFactory and just declare that it's always functor?

- Factor out boxing and unboxing functionality from the c10 dispatcher into a c10::KernelFunction class - Move that class and everything else it depends on into ATen/core/boxing - This also allows us to get rid of c10::KernelCache. Instead, we now store a pointer to the unboxed functor in c10::KernelFunction. - We're also getting rid of the DispatchTableEntry struct and instead store KernelFunction directly. - To make this work, we need to change the dispatcher calling API from Dispatcher::lookup().callBoxed/callUnboxed and OperatorEntry::lookup().callBoxed/callUnboxed to Dispatcher::callBoxed/callUnboxed and OperatorEntry::callBoxed/callUnboxed. Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

aten/src/ATen/core/boxing/KernelFunction.h

- Factor out boxing and unboxing functionality from the c10 dispatcher into a c10::KernelFunction class - Move that class and everything else it depends on into ATen/core/boxing - This also allows us to get rid of c10::KernelCache. Instead, we now store a pointer to the unboxed functor in c10::KernelFunction. - We're also getting rid of the DispatchTableEntry struct and instead store KernelFunction directly. - To make this work, we need to change the dispatcher calling API from Dispatcher::lookup().callBoxed/callUnboxed and OperatorEntry::lookup().callBoxed/callUnboxed to Dispatcher::callBoxed/callUnboxed and OperatorEntry::callBoxed/callUnboxed. Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

ezyang · 2019-09-20T14:21:09Z

aten/src/ATen/core/boxing/KernelFunction.h

+        TORCH_INTERNAL_ASSERT(false, "Tried to call KernelFunction::callBoxed() on an uninitialized KernelFunction.");
+      } else {
+        // TODO We want to introduce the invariant that all kernels must be callable in a boxed way, then this case should be impossible.
+        TORCH_INTERNAL_ASSERT(false, "Tried to call KernelFunction::callBoxed() on a KernelFunction that can only be called with KernelFunction::callUnboxed().");


So, in the meantime, this should be a TORCH_CHECK, right?

I don't think so. This doesn't fail due to user error but if we mucked up codegen. KernelFunction is not called directly by users and I don't think they can trigger this by using public APIs. Can you give an explanation why it should be TORCH_CHECK?

I thought the TODO implied this was triggerable from userspace.

ezyang · 2019-09-20T14:30:33Z

aten/src/ATen/core/boxing/KernelFunction.h

+   */
+  void callBoxed(Stack* stack) const {
+    if (C10_UNLIKELY(boxed_kernel_func_ == nullptr)) {
+      if (unboxed_kernel_func_ == nullptr) {


As an organizational point, it would be nice for this error handling code to live in the cpp file rather than the header. There's no reason for it to live inline, and it would make certain types of refactoring/debugging work easier (as it would avoid a recompilation of the world)

Calling operators is a hot path that I want to be fully inlineable. I could move it to a KernelFunction_inl.h header if you think that's better? Let me know.

I'm talking about the error reporting code, not the calling code.

ezyang · 2019-09-20T15:35:07Z

aten/src/ATen/core/boxing/KernelFunction.h

+   *
+   * KernelFunction::callUnboxed() is generally better, since it will allow
+   * calling KernelFunctions that are backed by either boxed or unboxed
+   * kernels, but that one will not work for all types.


I don't understand, is this a compilation-time error? I would prefer if we have a single entry point callUnboxed, that just skips the fallback to boxed types if the types in question are not supported. This is how I implemented it in #26368

Yes, I chose a bit a different approach here. I want the compiler to show an error if you have an op with use_c10_dispatcher: full but the boxed fallback doesn't work. That means the use_c10_dispatcher flag is the only configuration setting you have to move to fully migrate ops over to the c10 dispatcher. If you have some SFINAE config to remove ops from creating boxed wrappers, that would add an additional configuration setting and make the system more complex.

I don't buy the argument that it's more complex. The reason I don't buy it is because #26368 is able to avoid a bunch of pieces which you have to do in this PR:

There's no per-operator setting saying whether or not it's full or unboxedOnly. The templates figure it out automatically.

There's no change to code generation, as is here. The templates figure it out automatically.

The actual difference in the dispatcher is a single if statement https://github.com/pytorch/pytorch/pull/26368/files#diff-4f2e119446bc2edf856334581e5c70c9R217 testing if the boxed version is supported, and error if it is not

The only downside is that you don't get compile time errors in the way you've done here. But it's easy to turn all of the runtime errors into compile time errors, to see if you have 100% coverage: change supports_boxed_fallback to return unconditionally true.

aten/src/ATen/core/boxing/KernelFunction.h

ezyang · 2019-09-20T15:37:52Z

aten/src/ATen/core/boxing/KernelFunction.h

+    // TODO Remove this function once all kernels support a boxed variant
+
+    if (unboxed_kernel_func_ != nullptr) {
+      using ActualSignature = Return (OperatorKernel*, Args...);


This signature looks weird to me. Why is the unboxed kernel taking an OperatorKernel* as an argument? (I.e., what does this enable you to implement?) For reference, unboxed FuncType in #26368 is just Return(Args...) which provides evidence that you don't need the OperatorKernel*

The autogenerated boxed wrappers need to call into the unboxed functor so they need a pointer to it. This replaces what c10::KernelCache did before, which was also passed in to the boxed kernel. In #26368, you don't call unboxed from boxed but only the other way round, which is much easier and which is why you don't need this.

You're right!

ezyang · 2019-09-20T15:53:53Z

aten/src/ATen/core/boxing/KernelFunction.h

+    if (unboxed_kernel_func_ != nullptr) {
+      using ActualSignature = Return (OperatorKernel*, Args...);
+      ActualSignature* func = reinterpret_cast<ActualSignature*>(unboxed_kernel_func_);
+      return (*func)(getFunctor_(), std::forward<Args>(args)...);


Are you using the term functor here in an untraditional way? In C++, I expect functor to be a "function like object". But you don't call the functor, you call a function pointer, passing the functor as an argument. That's weird?

this is calling into the boxed kernel, passing in the functor so the boxing wrapper can call it if the boxed kernel is just a wrapper that calls into the unboxed functor.

ezyang · 2019-09-20T16:03:34Z

aten/src/ATen/core/boxing/KernelFunction.h

+  , unboxed_kernel_func_(unboxed_kernel_func)
+  {}
+
+  OperatorKernel* getFunctor_() const {


Why does this method have a trailing underscore? Why is it called "get functor" but returns an "operator kernel*"?

because c10::OperatorKernel is the public API all kernel functors have to inherit from in the registration API.

ezyang · 2019-09-20T17:08:20Z

aten/src/ATen/core/boxing/KernelFunction.h

+    (*boxed_kernel_func)(functor, &stack);
+
+    TORCH_INTERNAL_ASSERT(stack.size() == 1, "A boxed kernel should only push one return to the stack");
+    return std::move(stack[0]).to<Return>();


Isn't stack[0] already a temporary, so the std::move is not necessary?

no, std::vector<T>::operator[] returns a T&, not a T.

ezyang · 2019-09-20T17:09:23Z

aten/src/ATen/core/boxing/KernelFunction.h

+  }
+};
+template<class... Args>
+struct boxAndCallBoxedFunc<void, Args...> final {


Man, I hate that void needs special casing here XD

yes, me too. if constexpr would fix this but that's C++17 only.

ezyang · 2019-09-20T17:18:40Z

aten/src/ATen/core/dispatch/OperatorEntry.h

-      return OpKernel(kernel.kernel_func, kernel.cache_creator_func, kernel.unboxed_kernel_func);
+  template<class Return, class... Args>
+  Return callUnboxed(TensorTypeId dispatchKey, Args... args) const {
+    // TODO Remove dispatchKey argument and instead infer dispatchKey from args...


As I think you are aware, multi_dispatch_tensor_type_set should work well for this.

yes, that's a perfect match. I'll probably move the logic into the dispatch table so it's next to the boxed logic calling multi_dispatch_tensor_type_set.

ezyang

This seems like an improvement. But I still think this code can be simplified more.

smessmer · 2019-09-20T17:36:57Z

Going to land this. As summary, I think there's three main points we need to address in future:

It doesn't do error checking for callUnboxed() yet. I'm currently working on moving our code base to GCC 5, which would allow me to add error checking using compile time type ids.
The kernelFactory_ is only needed if we want to allow kernel functors to store Tensors (or other things that can't be created at static initialization time).
The API is quite wide because our custom op API offers a lot of ways of creating kernels. This could be narrowed if we're willing to break the custom op API.

facebook-github-bot · 2019-09-21T02:38:03Z

This pull request has been merged in ed207b5.

Summary: Pull Request resolved: pytorch/pytorch#26337 - Factor out boxing and unboxing functionality from the c10 dispatcher into a c10::KernelFunction class - Move that class and everything else it depends on into ATen/core/boxing - This also allows us to get rid of c10::KernelCache. Instead, we now store a pointer to the unboxed functor in c10::KernelFunction. - We're also getting rid of the DispatchTableEntry struct and instead store KernelFunction directly. - To make this work, we need to change the dispatcher calling API from Dispatcher::lookup().callBoxed/callUnboxed and OperatorEntry::lookup().callBoxed/callUnboxed to Dispatcher::callBoxed/callUnboxed and OperatorEntry::callBoxed/callUnboxed. ghstack-source-id: 90459911 Test Plan: unit tests Differential Revision: D17416607 fbshipit-source-id: fd221f1d70eb3f1b4d33092eaa7e37d25684c934

Summary: Pull Request resolved: pytorch#26337 - Factor out boxing and unboxing functionality from the c10 dispatcher into a c10::KernelFunction class - Move that class and everything else it depends on into ATen/core/boxing - This also allows us to get rid of c10::KernelCache. Instead, we now store a pointer to the unboxed functor in c10::KernelFunction. - We're also getting rid of the DispatchTableEntry struct and instead store KernelFunction directly. - To make this work, we need to change the dispatcher calling API from Dispatcher::lookup().callBoxed/callUnboxed and OperatorEntry::lookup().callBoxed/callUnboxed to Dispatcher::callBoxed/callUnboxed and OperatorEntry::callBoxed/callUnboxed. ghstack-source-id: 90459911 Test Plan: unit tests Differential Revision: D17416607 fbshipit-source-id: fd221f1d70eb3f1b4d33092eaa7e37d25684c934

[wip] c10::KernelFunction

46b3266

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

smessmer requested a review from apaszke as a code owner September 17, 2019 05:25

pytorchbot added caffe2 oncall: jit Add this issue/PR to JIT oncall triage queue module: internals Related to internal abstractions in c10 and ATen module: operators labels Sep 17, 2019

Update on "[wip] c10::KernelFunction"

cdb60ce

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

Update on "[wip] c10::KernelFunction"

466587c

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

smessmer added a commit that referenced this pull request Sep 17, 2019

[wip] c10::KernelFunction

a913153

Pull Request resolved: #26337 ghstack-source-id: 90228146 Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/)

ezyang reviewed Sep 17, 2019

View reviewed changes

aten/src/ATen/core/dispatch/DispatchTable.h Outdated Show resolved Hide resolved

ezyang reviewed Sep 17, 2019

View reviewed changes

aten/src/ATen/core/dispatch/KernelFunction.h Outdated Show resolved Hide resolved

ezyang reviewed Sep 17, 2019

View reviewed changes

aten/src/ATen/core/dispatch/KernelFunction.h Outdated Show resolved Hide resolved

Update on "[wip] c10::KernelFunction"

7095126

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

ezyang reviewed Sep 17, 2019

View reviewed changes

aten/src/ATen/core/dispatch/KernelFunction.h Outdated Show resolved Hide resolved

Update on "[wip] c10::KernelFunction"

c08c02f

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

Update on "[wip] c10::KernelFunction"

f2227fe

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

pytorchbot added the module: docs Related to our documentation, both in docs/ and docblocks label Sep 17, 2019

Update on "[wip] c10::KernelFunction"

ab7c73b

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

This was referenced Sep 18, 2019

Add two levels to use_c10_dispatcher #26272

Closed

Enable creation of boxing wrappers for some aten operators #26273

Closed

smessmer added 4 commits September 17, 2019 19:11

Update on "[wip] c10::KernelFunction"

a851344

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

Update on "[wip] c10::KernelFunction"

b0acbdb

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

Update on "[wip] c10::KernelFunction"

a888248

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

Update on "[wip] c10::KernelFunction"

d3d611e

Differential Revision: [D17416607](https://our.internmc.facebook.com/intern/diff/D17416607/) [ghstack-poisoned]

dzhulgakov reviewed Sep 19, 2019

View reviewed changes

smessmer requested a review from dzhulgakov September 19, 2019 18:41

smessmer mentioned this pull request Sep 19, 2019

[wip] Error checking for KernelFunction::callUnboxed() #26482

Closed

ezyang reviewed Sep 19, 2019

View reviewed changes

aten/src/ATen/core/boxing/KernelFunction.h Outdated Show resolved Hide resolved

smessmer requested a review from ezyang September 19, 2019 23:02

smessmer added 2 commits September 19, 2019 16:04

ezyang reviewed Sep 20, 2019

View reviewed changes

aten/src/ATen/core/boxing/KernelFunction.h Show resolved Hide resolved

ezyang reviewed Sep 20, 2019

View reviewed changes

ezyang approved these changes Sep 20, 2019

View reviewed changes

facebook-github-bot closed this in ed207b5 Sep 21, 2019

facebook-github-bot added the merged label Sep 21, 2019

smessmer mentioned this pull request Oct 18, 2019

c10 dispatcher data composition diagram is hard to understand #26171

Closed

facebook-github-bot deleted the gh/smessmer/49/head branch October 28, 2019 22:19

mruberry added the Merged label Oct 28, 2020

Conversation

smessmer commented Sep 17, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

smessmer Sep 20, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

smessmer Sep 20, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

smessmer Sep 20, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ezyang left a comment

Choose a reason for hiding this comment

Uh oh!

smessmer commented Sep 20, 2019

Uh oh!

facebook-github-bot commented Sep 21, 2019

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

smessmer commented Sep 17, 2019 •

edited

Loading

smessmer Sep 20, 2019 •

edited

Loading

smessmer Sep 20, 2019 •

edited

Loading

smessmer Sep 20, 2019 •

edited

Loading