DebugTransform - new interface, fixes and doc update.

[kshitij12345]

Repro

import torch
import thunder

from thunder.dev_utils.debug_transform import DebugTransform

def pre_callback(bsym, *args, **kwargs):
    return f"Pre - {bsym.sym.name}"

def foo(x):
    y = x + 1
    z = y * 2
    return z

jfoo = thunder.jit(foo, transforms=[DebugTransform(pre_callback=pre_callback)])

x = torch.randn(3, 3)
o = jfoo(x)
print(thunder.last_traces(jfoo)[-1])

Output

# Constructed by Debug trace (took 0.03 milliseconds)
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(x):
  # x: "cpu f32[3, 3]"
  debug_pre_add1(x, 1)
  y = torch.add(x, 1)  # y: "cpu f32[3, 3]"
    # y = ltorch.add(x, 1, alpha=None)  # y: "cpu f32[3, 3]"
      # _ = prims.convert_element_type(1, float)
      # y = prims.add(x, 1.0)  # y: "cpu f32[3, 3]"
  del x
  # Pre - mul
  debug_pre_mul1(y, 2)
  z = torch.mul(y, 2)  # z: "cpu f32[3, 3]"
    # z = ltorch.mul(y, 2)  # z: "cpu f32[3, 3]"
      # _ = prims.convert_element_type(2, float)
      # z = prims.mul(y, 2.0)  # z: "cpu f32[3, 3]"
  del y
  return z

Problems:

1. In the output trace - note the name of debug symbols debug_pre_add1 and debug_pre_mul1. It should be debug_pre_add1 and debug_pre_mul2 otherwise we may have name collision if we had multiple add.
2. We only updated the header for debug_pre_mul1 but debug_pre_add1 didn't update the captured information.
   This happens because we capture pre_debug_bsym which at the end of the loop refers to the last pre_debug_bsym (so essentially all _pre_call_ctx update the header of the last pre_debug_bsym.
   

   lightning-thunder/thunder/dev_utils/debug_transform.py

   Lines 42 to 49 in 815bff0

   	def _pre_call_ctx(bsym, *args, **kwargs):
   	out = self.pre_callback(bsym, *args, **kwargs)
   	thunder.core.utils.check_type(out, str)
   	pre_debug_bsym.header = out
   	pre_debug_name = f"debug_pre_{sym_name}{debug_counter}"
   	pre_debug_bsym = create_debug_boundsymbol(pre_debug_name, bsym, _pre_call_ctx)
   	new_bsyms.append(pre_debug_bsym)

Changes:
This PR fixes the above 2 problem and adds a test.

Besides that PR adds a new interface debug_execution_trace (open to suggestion for name) function which takes care of adding the transform (this functional interface should be more intuitive as discussed with @riccardofelluga). This PR also adds documentation for debug_execution_trace.

[riccardofelluga]

LGTM, stamping!