Lightning-AI/lightning-thunder

Better name for elements of list in `prologue_trace` and `computation_trace`

开放

#152 创建于 2024年4月10日

 (0 条评论) (0 个反应) (0 位负责人)Python (114 个派生)github user discovery
enhancementhelp wantedjit

仓库指标

星标
 (1,460 个星标)
PR 合并指标
 (30 天内没有已合并 PR)

描述

import thunder
import torch

def foo(xs):
    result = []
    for x in xs:
        result.append(x + x)
    return result

jfoo = thunder.jit(foo)

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

Names for the arguments to the computation trace are : res, x, a, b, t_0_4, t_0_5. It would be nice if there was a consistent pattern.

Prologue Trace

# Constructed by Transform for execution (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast()
def prologue(*args, **kwargs):
  # args: "Any"
  check_len(args, 1)
    # prims.check_len(args, 1)
  # kwargs: "Any"
  check_len(kwargs, 0)
    # prims.check_len(kwargs, 0)
  subscr: "Any" = args[0]
  res: "cpu f32[3]" = subscr[0]
  x: "cpu f32[3]" = subscr[1]
  a: "cpu f32[3]" = subscr[2]
  b: "cpu f32[3]" = subscr[3]
  t_0_4: "cpu f32[3]" = subscr[4]
  t_0_5: "cpu f32[3]" = subscr[5]
  ...
  return (res, x, a, b, t_0_4, t_0_5)

Computation Trace

# Constructed by Delete Last Used (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast()
def computation(res, x, a, b, t_0_4, t_0_5):
  # res: "cpu f32[3]"
  # x: "cpu f32[3]"
  # a: "cpu f32[3]"
  # b: "cpu f32[3]"
  # t_0_4: "cpu f32[3]"
  # t_0_5: "cpu f32[3]"
  result = torch.add(res, res)  # result: "cpu f32[3]"
    # result = ltorch.add(res, res, alpha=None)  # result: "cpu f32[3]"
      # result = prims.add(res, res)  # result: "cpu f32[3]"
  del res
  ...
  return [result, t1, t2, t3, t4, t5]

贡献者指南