update tracing codegen to use redispatch API

[bdhirsh]

Taking advantage of the new redispatch API to clean up the codegen'd tracing kernels. Instead of directly interacting with the Dispatcher, the tracing kernels now just call the redispatch API directly.

One small benefit to this: hopefully the compiler is more likely to inline Dispatcher::redispatch(), since it's now used in fewer call-sites. After this change, the only places it's used are:

• the redispatch API (RedispatchFunctions.cpp)
• BackendSelect kernels.

One small complication: the redispatch API doesn't interact too well with manual_cpp_binding ops currently. I put a note with some thoughts in the comments.

Example tracing kernel before:

Tensor add_Tensor(c10::DispatchKeySet ks, const Tensor & self, const
  torch::jit::Node* node = nullptr;
  std::shared_ptr<jit::tracer::TracingState> tracer_state;
  if (jit::tracer::isTracing()) {
    tracer_state = jit::tracer::getTracingState();
    at::Symbol op_name;
    op_name = jit::Symbol::fromQualString("aten::add");
    node = tracer_state->graph->create(op_name, /*num_outputs=*/0);
    jit::tracer::recordSourceLocation(node);
    jit::tracer::addInputs(node, "self", self);
    jit::tracer::addInputs(node, "other", other);
    jit::tracer::addInputs(node, "alpha", alpha);
    tracer_state->graph->insertNode(node);

    jit::tracer::setTracingState(nullptr);
  }
  static auto op = c10::Dispatcher::singleton()
      .findSchemaOrThrow("aten::add", "Tensor")
      .typed<Tensor (const Tensor &, const Tensor &, Scalar)>();
  auto result =c10::Dispatcher::singleton()
      .redispatch<Tensor, const Tensor &, const Tensor &, Scalar>(op,
  if (tracer_state) {
    jit::tracer::setTracingState(std::move(tracer_state));
    jit::tracer::addOutput(node, result);
  }
  return result;
}

after: (note the lack of Dispatcher:: calls)

Tensor add_Tensor(c10::DispatchKeySet ks, const Tensor & self, const Tensor & other, Scalar alpha)
  torch::jit::Node* node = nullptr;
  std::shared_ptr<jit::tracer::TracingState> tracer_state;
  if (jit::tracer::isTracing()) {
    tracer_state = jit::tracer::getTracingState();
    at::Symbol op_name;
    op_name = jit::Symbol::fromQualString("aten::add");
    node = tracer_state->graph->create(op_name, /*num_outputs=*/0);
    jit::tracer::recordSourceLocation(node);
    jit::tracer::addInputs(node, "self", self);
    jit::tracer::addInputs(node, "other", other);
    jit::tracer::addInputs(node, "alpha", alpha);
    tracer_state->graph->insertNode(node);

    jit::tracer::setTracingState(nullptr);
  }
  auto result =at::redispatch::add(ks & c10::DispatchKeySet(c10::DispatchKeySet::FULL_AFTER, c10::D
  if (tracer_state) {
    jit::tracer::setTracingState(std::move(tracer_state));
    jit::tracer::addOutput(node, result);
  }
  return result;
}

Stack from ghstack:

• update tracing codegen to use redispatch API #52009 update tracing codegen to use redispatch API
• update autograd kernels to use redispatch #51363 update autograd kernels to use redispatch
• codegen'd API for redispatching #52008 codegen'd API for redispatching

Differential Revision: D26356078

[bdhirsh]

manual_cpp_binding ops are hardcoded directly into Functions.h, so they aren't currently part of the redispatch API :( I opted instead to call their codegen'd __dispatch_*, which ARE available.

It would probably be nice if I could just call the fast-path variants instead. I opted not to for two reasons:

• I can't think of an elegant way to do it. One option is to additionally tell the codegen to generate variants in the redispatch API that directly call the hardcoded fast-path variant for ops that have manual_cpp_binding set, appropriately calling the function/method variant depending on the op.
• Another would just be to duplicate each hardcoded function from Functions.h into both namespaces (bad).
• The perf benefit is probably minimal, since tracing kernels aren't called that frequently.

If anyone thinks that tradeoff is worth it though, lmk and I'm happy to add it in!

[ezyang]

Reasoning seems fine.

But consider constructing a full CppSignatureGroup and then extracting name from the faithful signature (if it exists). Helper method in CppSignatureGroup for returning faithful if it exists and regular if not will be helpful. Then you can axe the duplication of __dispatch logic.

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit 07f759f (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions to the (internal) Dr. CI Users group.

[ezyang]

the copy paste here got botched by screen

  auto result =at::redispatch::add(ks & c10::DispatchKeySet(c10::DispatchKeySet::FULL_AFTER, c10::D

[codecov]

Codecov Report

Merging #52009 (07f759f) into gh/bdhirsh/77/base (ad8e906) will decrease coverage by 0.00%.
The diff coverage is n/a.

@@                  Coverage Diff                   @@
##           gh/bdhirsh/77/base   #52009      +/-   ##
======================================================
- Coverage               80.76%   80.76%   -0.01%     
======================================================
  Files                    1969     1969              
  Lines                  216037   216037              
======================================================
- Hits                   174480   174478       -2     
- Misses                  41557    41559       +2     

[facebook-github-bot]

@bdhirsh merged this pull request in 947225c.