Lookaside for torch.ops.higher_order.autograd_function_apply

[crcrpar]

What does this PR do?

As per #1248, the support of torch.ops.higher_order.autograd_function_apply would be a bit more flexible by tracing into both fwd and bwd.

cc @apaz-cli

[IvanYashchuk]

I want the lookasides' scope to be limited only to the preprocessing of PyTorch code. If the removed code is reused in the updated lookaside we'll achieve that.

[t-vi]

initial_computation trace is not a valid Python function

import thunder.torch as ltorch
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(x):
  # x: "cpu f32[2, 2]"

  # /home/tv/data/firma/grid/thunder/lightning-thunder/thunder/tests/test_jit_general.py:1216: 	        return grad_output * torch.cos(x)
  t6 = ltorch.autograd_function_apply(_function_0, _function_1, x, args_tensor_mask=[True], non_differentiable_idx=[])  # t6: "cpu f32[2, 2]"
    # t6 = ltorch.sin(x)  # t6: "cpu f32[2, 2]"
      # t6 = prims.sin(x)  # t6: "cpu f32[2, 2]"
  return t6
.

[crcrpar]

Is it a contract that any of the traces generated from a callable is executable?

[IvanYashchuk]

What part of the provided code snippet makes it an invalid Python function?

[t-vi]

Is it a contract that any of the traces generated from a callable is executable?

This has been the case up to now and it is what I have repeatedly said about why I want the things properly inlined.

[IvanYashchuk]

Could you please elaborate on what part of the example trace makes it improperly executable?

A valid Python program from a trace is generated using its string representation (trace.python()) and its "context" (trace.python_ctx). The context is used to pass as the globals= argument to the built-in exec function (https://docs.python.org/3/library/functions.html#exec).

[crcrpar]

I brought the two

lightning-thunder/thunder/tests/test_jit_general.py

Line 1233 in 0205c73

initial_computation_trace = thunder.last_traces(jitted)[0]

with this PR and main branch.
The top is this PR, the bottom, main.

I'm not seeing the difference.

import thunder
import thunder.torch as ltorch
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(x):
  # x: "cpu f32[2, 2]"

  # /home/mkozuki/ghq/github.com/crcrpar/lightning-thunder/thunder/tests/test_jit_general.py:1218:              return grad_output * x.cos()
  t6 = ltorch.autograd_function_apply(_function_0, _function_1, x, args_tensor_mask=[True], non_differentiable_idx=[])  # t6: "cpu f32[2, 2]"
    # t6 = ltorch.sin(x)  # t6: "cpu f32[2, 2]"
      # t6 = prims.sin(x)  # t6: "cpu f32[2, 2]"
  return t6

import thunder
import thunder.torch as ltorch
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(x):
  # x: "cpu f32[2, 2]"

  # /home/mkozuki/ghq/github.com/crcrpar/lightning-thunder/thunder/tests/test_jit_general.py:1217:              return torch.ops.higher_order.autograd_function_apply(
  t0 = ltorch.autograd_function_apply(_function_0, _function_1, x, args_tensor_mask=[True], non_differentiable_idx=[])  # t0: "cpu f32[2, 2]"
    # t0 = ltorch.sin(x)  # t0: "cpu f32[2, 2]"
      # t0 = prims.sin(x)  # t0: "cpu f32[2, 2]"
  return t0

[t-vi]

Masaki offered #1526 as an alternative without higher order functions. I merged it to unblock uses of autograd_function_apply.