본문 바로가기

AI Compiler framework

CudaExecutor.run - IR 을 실제 텐서 계산으로 바꾸는 과정

aicf.CudaExecutor().run(m, feed) 가 하는 일 : vid - torch.Tensor 슬롯 테이블을 완성하고, op 를 순서대로 실행

핵심 자료구조 : slots

slots: List[Optional[torch.Tensor]] = [None] * len(b.values)
  • 길이는 b.values 개수와 같은
  • 인덱스 = vid
  • slots[vid] 에는 그 Value 에 해당하는 실제 텐서가 들어감

run 이 끝나면 builder 의 모든 value 가 실제 torch.Tensor 로 구체화됨

 

단계별 생성 텐서

1) feed 로 들어온 값이 slots 에 꽂힘 ( inputs + params)

for vid in b.input_vids:
    slots[vid] = feed[name]
  • x
  • fc1.W
  • fc1.b

같은 것들이 그대로 slots 에 참조로 들어감, slots[vid] 는 모델이 쓰는 실제 입력 / 파라미터 텐서 핸들이 된다.

 

2) feed 에 없는 Value 들은 전부 allocate 된다

for v in b.values:
    if slots[v.vid] is None:
        slots[v.vid] = torch.empty(spec.shape, dtype=..., device=...)

여기가 run 에서 실제로 새 텐서가 생성되는 지점

  • fc1.out
  • fc1.out_bias
  • relu1.out
  • res.out

같은 중간 / 출력 Value 들이 torch.empty 로 생성됨

 

3) lower(m) 호출 : LoweredOp 리스트가 생성됨

lowered = self.lower(m)

텐서가 아닌 메타가 만들어진다.

  • kind_id
  • attr_schema
  • attr_blob
  • in_vids/out_vids

실행 지시서 생성 단계

 

4) 실행 루프 : _C.op_call(...) 가 slots 텐서를 읽고 / 쓰면서 값을 채움

for lop in lowered:
    ins  = [slots[v] for v in lop.in_vids]
    outs = [slots[v] for v in lop.out_vids]
    op_call(kind_id, ins, outs, schema, blob)
  • outs 텐서는 torch.empty 로 만들어져 있음
  • _C.op_call 은 그 outs 텐서 메모리에 값을 write 한다
  • 새 텐서를 만들지 않고 기존 버퍼에 쓰는 느낌

 

5) outputs 반환

out[name] = slots[vid]

출력은 그냥 slots 의 일부를 딕셔너리로 포장해서 반환