graph_executor_v2 구성

핵심 4단계

1. 모델 레이어 -> 그래프 IR 로 정적 계획
2. 텐서 메모리 (디바이스) -> 아레나로 주소 확정/재사용
3. 노드 -> 각 독립 모듈의 호출 스텁 ( 속성 + 파라미터 + WS 포함 )
4. 런 시점 -> 노드를 순회하며 포인터와 스트림을 넘겨 제로-카피 호출

 

1) IR (GraphIR) : 레이어 시퀀스를 노드 + 텐서로 바꾸는 중간 표현

Node 가 무엇인가?

• Node(kind, inputs, outputs, attrs, layer_ref, ws_fwd, ws_bwd)
  • kind : 어떤 연산인지 식별 ( Conv2D, GEMM, ReLU, Flatten, ... )
  • inputs / outputs : 텐서 ID 의 리스트, 노드 간 데이터 흐름은 텐서 ID 로 연결됨
  • attrs : 이 노드가 필요한 하이퍼파라미터 ( Stride, Padding, Act  ... ) , C++/CUDA 런처에 그대로 매핑된 값
  • layer_ref : 파이썬 레이어 객체에 대한 참조
  • ws_fwd/ws_bwd : 빠른 경로용 워크스페이스 핸들 (디바이스 배열들)

Node 는 독립 모듈을 호출하기 위한 콜 시트

입력/출력은 텐서 ID 로 커널에 넘길 러닝타임 속성은 attrs 로, 가중치는 layer_ref 에서 포인터로 뽑아 쓴다.

 

TensorDesc 가 무엇인가?

• TensorDesc(shape, dtype, buffer_id, offset)
  • buffer_id + offset 조합은 아레나 내 1D 디바이스 버퍼의 부분 뷰
  • 같은 buffer_id 를 공유하면 인플레이스/앨리어싱이 가능
  • dtype 은 대부분 cp.float32 로 고정

GraphCompiler.compile 흐름

• 입력 텐서 하날르 먼저 등록 ( buffer_id = 0 ) 
• 모델 레이어를 순회하며
  • 레이어 종류별로 출력 shape 를 계산 -> 새 TensorDesc 를 만들고, ir.tensors 에 등록
  • 해당 연산에 맞는 Node 를 만들어 ir.nodes 에 추가
  • Conv(groups == 1) 인 경우 fast path WS 를 미리 할당해 node.ws_fwd 에 꽂아둠
  • ReLU 는 in-place alis 를 위해 입력과 동일한 buffer_id/offset 의 TensorDesc 를 하나 더 만들어 뷰만 변경
• 마지막으로 ir.output_ids 지정

즉, 컴파일 단계에서 이미 노드 그래프와 텐서-버퍼 매핑이 확정됨

 

2) MemoryArena : "한 장짜리 디바이스 메모리(아레나)" 에서 텐서를 뷰로 잘라쓰기

plan_memory

• 모든 TensorDesc 를 훑어 buffer_id 별 필요한 최대 길이 (=offset+numel) 를 계산
• 그 길이만큼 cp.empty 로 디바이스 1D 버퍼를 한 번만 할당
• 이후 실행 중 할당/해제 없음

arena.view(desc)

• 아레나의 1D 버퍼에서 offset : offset+numel 범위를 잘라 reshape 하여 디바이스 뷰를 돌려줌
• 이게 곧 텐서의 장소가 됨, 이 뷰의 data.ptr 이 커널에 들어감

효과 : 레이어 사이 텐서 복사 최소화, 인플레이스/앨리어시 제어

커널 런칭은 모든 메모리 준비가 끝난 상태

3) 독립 모듈바인딩과 주소 연결 : 포인터를 직접 넘기는 구조

포인터(주소) 전달

•  _run_impl 에서 노드를 순회하며, 입력/출력 텐서를 arena.view(...) 로 가져옴 - 이건 디바이스 배열 (cuda.ndarray)
• 저수준 런처를 쓸 때 (_ops_conv2d.forwad 같은 pybind 바인딩 )
  • int(X.data.ptr), int(W.data.ptr), int(Y.data.ptr) 식으로 디바이스 포인터를 넘김
  • 스트림도 cp.cuda.get_current_stream().ptr 로 그대로 전달
• 고수준 헬퍼를 쓸 때 (conv_ops.forward, gemm_ops.forward)
  • cupy 배열을 넘기면 내부에서 필요한 포인터를 꺼내 커널 호출 ( 편하지만 fast-path 최적화는 제한적 )
  • 출력 버퍼 out=Y 로 지정해 추가 할당/복사 방지.

워크스페이스(ws) 주입

• Conv fast path(groups==1) 에서 미리 확보한 ws_fwd 를 포인터로 주입
• 이렇게 하면 런타임 동적 할당 없이 곧바로 커널을 쓸 수 있어 CUDA Graph 캡처 안정성 + 성능을 동시에 확보

Z alias ( = 중간활성 저장/혹은 act 전의 pre-activation)

• fast path 에서 Z 가 필요한 구조라면 Z 를 별도의 버퍼 대신 Y 에 alias 로 박음 (추가 메모리 제로)

 

실행 경로 : 일반 실행 vs CUDA Graph 실행

일반 실행 ( _run_impl )

1. 입력을 아레나 입력 뷰에 써 넣고
2. for node in ir.nodes 로 순회하며 연산 호출
3. 최종 출력 텐서의 뷰를 리턴

CUDA Graph 캡처(compile(..., use_cuda_graph=True)

• 워밍업 1회로 모든 계획/WS 가 고정되게 함.
• 별도의 캡처 스트림에서 _run_imple(None) 을 수행하며 begin_capture ~ end_capture.
• 그래프 객체에 대해 3종 경로 지원
  • Graph 객체가 upload/launch 메서드를 제공 : graph.upload(stream), 이후 graph.launch(stream)
  • 모듈 함수 스타일 : cp.cuda.graph.upload(graph, stream) - cp.cuda.graph.launch(graph, stream)
  • 런타임 포인터 경로 : graph.graph 또는 graph.ptr 추출 - cudaGraphInstantiate / graphUpload / graphLaunch 직접 호출 ( self._use_runtime_graph_api=True 로 분기 )
• run(x) 에서는
  • 입력 뷰에 x 를 써 넣고
  • 위 3종 경로 중 해당 방식으로 그래프 런치
  • 스트림 sync 후 출력 뷰를 .copy() 해 반환

포인트 : 캡처 중에는 어떤 H2D/동적할당도 일어나면 안 됨, 

아레나/WS 선할당, 파이썬 측 수학은 CPU 에서만, 스트림은 항상 같은걸 넘기는 게 중요

 

레이어 연결이 "주소 차원" 에서 어떻게 이어지나? (데이터 흐름 예시)

예: Conv2D → ReLU → Flatten → Dense

1. Conv2D
   • 입력: tensor#0 (buffer_id=0, offset=0)
   • 출력: tensor#1 (buffer_id=1, offset=0) ← 새 버퍼
   • 호출 시 _gconv.forward(X.ptr, W.ptr, Y.ptr, ...)로 직접 포인터 전달.
2. ReLU (in-place alias)
   • 입력: tensor#1
   • 출력: tensor#2 (같은 buffer_id=1, 같은 offset) ← 같은 장소
   • cp.maximum(X,0,out=Y)지만 X,Y가 같은 버퍼를 가리키므로 사실상 in-place.
3. Flatten
   • 출력: tensor#3는 shape만 (N, C*H*W)로 바뀐 같은 버퍼 뷰(여전히 buffer_id=1).
   • 호출단에선 아무 것도 안 함(뷰 변경만으로 충분).
4. Dense/GEMM
   • 입력: tensor#3(buffer_id=1)
   • 출력: tensor#4(buffer_id=2) ← 새 버퍼
   • gemm_ops.forward(A, W, b, out=Y)로 출력 버퍼를 지정해 고정된 장소에 결과를 씀.

이렇게 버퍼 ID/offset이 실질적인 “주소계약”이 되고, 모든 노드는 그 계약에 따라 직접 포인터를 받아 자신의 커널을 실행해. 중간에 불필요한 할당/복사가 없다 보니 성능/캡처 안정성이 좋아짐.

 

각각의 독립 모듈과의 바인딩 관계

• Conv2D fast path: graph_executor_v2.ops._ops_conv2d (pybind 바인딩)
  • 완전 저수준 경로: 포인터/shape/attrs/스트림/WS를 정수로 직통 전달
  • C++/CUDA 런처(launcher.cu/cpp)가 그대로 받아 <<<grid,block,stream>>> 커널을 호출.
• Conv2D fallback / GEMM: graph_executor_v2.ops.conv2d, graph_executor_v2.ops.gemm
  • cupy ndarray를 받아 내부에서 포인터 꺼내고 필요한 임시 텐서를 관리.
  • out=를 주면 아레나 버퍼에 바로 써주도록 유도.
• ReLU/BN/Flatten:
  • ReLU는 cupy elementwise(단순)로 out=Y 지정.
  • BN은 간단한 학습모드 식으로 구현(cpu-side는 아님, cupy 연산으로 디바이스상에서 mean/var 계산).
  • Flatten은 IR단에서 “뷰 변경”만 하므로 런타임 계산 없음.

---

스트림과 캡처 안정성

• 항상 cp.cuda.get_current_stream().ptr를 그대로 바인딩에 전달.
  • 캡처 시점에는 “현재 스트림 == 캡처 스트림”이라 0(디폴트 스트림)로 던지면 안 됨.
  • 코드에 주석도 있음: # ★ 현재 스트림 사용 (0 금지: 캡처 스트림과 분리 위험)
• 워밍업 한 번으로 모든 메모리(아레나/WS)와 경로가 고정 → 캡처에서 동적 이벤트 없음.

 

확장/유지보수 팁

1. 새 연산 추가
   • compile에서 레이어 → shape 추론 → TensorDesc/Node 생성 규칙만 추가.
   • 가능한 한 out 버퍼를 미리 정해서 바인딩에 out=로 넘겨 추가 할당 금지.
   • fast path가 있다면 ws_fwd를 노드에 매달아 포인터 주입.
2. 메모리 최적화
   • 지금은 “buffer_id=증가하는 새 버퍼” 기본 + 일부 in-place alias(ReLU/Flatten).
   • 사용 끝난 텐서의 buffer_id를 재사용하는 라이프타임 기반 플래너를 넣으면 메모리 폭 감소.
   • dtype 풀 분리(float16/bfloat16 등)도 확장 가능(현재는 float32 단일 아레나).
3. 역전파(학습)
   • 지금 코드는 v0 forward 중심.
   • Node.ws_bwd에 bwd용 WS를 준비하고, Node에 grad_inputs, grad_outputs 정의 후
     bwd패스에서 동일한 아레나 정책(buffer_id 재사용/alias)을 적용하면 됨.
   • CUDA Graph로 학습 루프까지 캡처하려면(optimizer step 포함) 모든 버퍼/워크스페이스/랜덤상태가
     캡처 안정적으로 고정되도록 설계를 더 탄탄히 해야 함.
4. 안전장치
   • 캡처 중 H2D를 막기 위해 shape 계산/math.prod는 파이썬 CPU side에서만.
   • 헬퍼 내부에서 갑자기 새로운 cupy 배열을 할당하지 않도록 out=...를 일관되게 제공.
   • 스트림 전달은 항상 “현재 스트림 포인터”.