AICF CUDA Plan A : Op 선택과 Kernel(Variant) 선택 메커니즘

목표

• Python 레벨에서는 커널 이름을 알 필요 없이 OpKind 와 텐서 / 속성만 넘긴다
• C++ 런타임(Register / Dispatch) 이 입력 텐서의 메타데이터와 attrs 를 기반으로 가장 적합한 커널 구현 (variant) 을 선택한다
• 동일 조건 입력은 동일 선택을 유지하고, CUDA Graph Capture / Replay 에서도 동일하게 동작한다.

 

1. 전체 실행 흐름

1.1 Python 호출 ( 사용자 코드 )

테스트 스크립트에서 핵심 호출 코드

aicf.op_call(kind, inputs, outputs, attrs)

Python 이 제공하는 정보

• kind : 어떤 op 인지 지정하는 enum ( aicf.OpKind)
• inputs : torch.Tensor list
• outputs : torch.Tensor list
• attrs : dict

Python 은 커널 이름을 전달할지 않음, 커널 선택은 전적으로 C++ 쪽에서 수행

 

1.2 PyBind 바인딩 (bindings.cpp)

• TensorDesc 생성
  • dtype / shape / stride / device pointer
• AttrPack packing
  • { key : bool / int / float } 를 key sort 후 AttrKV[] 로 구성
• CUDA Stream 결정
  • PyTorch current stream 사용 ( current_cuda_stream())

그 다음 이 정보를 dispatch entry 에 전달한다.

aicf::cuda::dispatch_v0(kind, in_descs, out_descs, attr_ptr, stream);

 

1.3 Dispatch entry ( dispatch.hpp)

dispatch_v0 은 선택 로직 없는 단순 wrapper

• OpCall 구조체로 포장
• Dispatch(call) 호출

OpCall call{};
call.kind = kind;
call.inputs = inputs;
call.outputs = outputs;
call.attrs = attrs;
call.stream = stream;
return Dispatch(call);

 

1.4 Registry Dispatch (registry.cpp) - 핵심 선택 로직

Dispatch(const OpCall& call) 은 두 가지를 한다

(A) OpKind 로 op 후보군 조회

vars = KernelRegistry::variants(call.kin)

• OpKind::EltwiseAdd - Add 용 variants
• OpKind::Gemm - Gemm 용 variants

 

(B) op 내부에서 kernel (variant) 선택 

선택 정책 존재

• priority 내림차순 정렬
• supported(...) == ture 인 첫 variant 선택
• workspace 필요하면 현재는 거부 상태
• 최종 chosen -> launch 호출

stable_sort(priority desc)
for v in order:
  if v.supported(...) -> choose
return chosen->launch(..., stream)

 

2. Op 선택 : Python 에서 이미 결정되었음

2.1 OpKind 가 연산을 결정

op_call(aicf.OpKind.EltwiseAdd, [a, b], [out], {})

이 순간, 이미 op 는 Add 로 고정, C++ dispatch 는 OpKind::EltwiseAdd 에 등록된 variants 만 검색

 

3. Op 내부 Kernel 선택 : TensorDesc / AttrPack 기반

3.1 KernelVariant 계약 (kernel_variant.hpp)

각 variant는 3개의 함수 포인터로 구성

• supported(...) : 이 커널이 이 입력에서 가능한가?
• query_workspace(...) : workspace 필요량 ( 현재는 0 만 허용 )
• launch(...) : 실제 커널 런치 (launcher entry)

struct KernelVariant {
  const char* name = nullptr;

  // higher = earlier selection
  int priority = 0;

  // reserved for future policies (arch/capture_safe/etc.)
  uint32_t flags = 0;

  aicf::Status (*launch)(
      const TensorDesc* inputs, int num_inputs,
      TensorDesc* outputs, int num_outputs,
      const void* attr,
      void* workspace, size_t workspace_bytes,
      cudaStream_t stream) = nullptr;

  size_t (*query_workspace)(
      const TensorDesc* inputs, int num_inputs,
      const void* attr) = nullptr;

  bool (*supported)(
      const TensorDesc* inputs, int num_inputs,
      const TensorDesc* outputs, int num_outputs,
      const void* attr) = nullptr;
}

커널 이름이 아닌

• 입력 dtype / shape / stride
• 포인터 alignment
• attrs 값
• arch, flags, worksapce 등...

입력 조건 자체가 선택 신호가 된다.

 

3.2 add op 예시 : f16 half2 variant 가 선택되는 조건

half2 선택 조건 ( add_f16_vec2_check )

• 2 inputs / 1 output
• f16 contiguous 1d
• same shape
• length even ( N % 2 == 0 )
• data pointer 4B aligned (half2 alignment)

if (!is_f16_contig_1d) return false;
if (!same_shape_1d) return false;
if (!even_len) return false;
if (!aligned(4)) return false;

선택 결과

• 조건 만족 - add_f16_vec2_variant_supported == true
• priority = 10 이므로 먼저 선택
• 실행은 add_f16_vec2_variant_launch 로 점프
• 내부에서 add_f16x2_kernel<<<...>>>(__half2*) 런치

 

 

4. register_all : 선택 가능한 variants 를 registry 에 등록

4.1 등록이 곧 연결선

aicf_cuda_register_all_kernels() 에서 OpKind 별 variants 가 등록된다.

R.register_kernel(OpKind::EltwiseAdd, make_add_f16_vec2_variant());

이 등록이 의미하는 것

• OpKind:EltwiseAdd 일 때 
• 후보군 vars 에 kernelVariant { launch = add_f16_vec2_variant_launch } 가 포함된다. 

즉, 런처 연결은 kernelVariant.launch 에 저장된 함수 포인터로 만들어진다.

 

5. 확장 방식 : 새 op 추가와 op 내부 variant 추가

5.1 새 op 추가 

• OpKind 확장
  • op_kind.hpp 에 enum 추가
  • _Count 업데이트
• TensorDesc 계약 정리
  • inputs / outputs 개수
  • shape rank / contig 요구
  • dtype 허용 범위
• launcher + kernels 구현
  • 공용 shim validate / util 활용
• KernelVariant factories + register_all 등록
  • KernelVariant make_layernorm_f16_variant();
  • register_all.cpp 에 등록

 

5.2 같은 op 에서 variant 추가 

• variant 하나 = supported + launch + query_worksapce
• priority 는 빠른 구현이 이길 가능성의 표현
• supported 는 정확한 조건을 표현

 

 

6. 개선 포인트

6.1 선택 정책을 priority + supported 에서 score 기반으로 확장

현재 priority 가 정적 우선순위, 작은 N 에 대해서도 vectorized 가 항상 이길 수 있음

• int ( *score)(...) 도입
• score 가 음수면 unsupporte, 0 이상이면 후보
• 가장 높은 score 선택

이 경우

• N 이 작을 때는 naive 가 더 빠를 수 있는 케이스도 자동 처리 가능

 

6.2 선택 결과 로깅 지원

NVTX 나 로그를 남기기

• op kind
• chosen variant name
• input dtype / shape
• stream id

이를 통해 ncu 없이도 fallback 이유 탐색 가능