Python <-> CUDA Ops Binding 설계 및 pyd 모듈 생성 과정

1. 목표와 설계 의도

1.1 문제 의식

초기에는 모든 CUDA ops 를 하나의 pybind 모듈에 바인딩하는 방식이 가장 단순

하지만 ops 수가 증가할수록 다음의 문제 발생

• 하나의 바인딩 파일이 비대
• 특정 op 수정 시 전체 바인딩 모듈 재빌드 필요
• ops 단위 디버깅/프로파일링이 어려움

1.2 해결 전략

CUDA op 구조와 동일하게 Python 바인딩도 op 단위로 분리한다

op 하나 = CUDA launcher 하나 = python 모듈 하나

 

2. 전체 구조 요약 (Binding 관련)

src/
 ┣ python_bindings/
 ┃ ┣ add.cpp
 ┃ ┣ relu.cpp
 ┃ ┣ gemm.cpp
 ┃ ┗ common.hpp
build/
 ┗ python/
   ┗ aicf_cuda/
     ┣ __init__.py
     ┣ add.pyd
     ┣ relu.pyd
     ┗ gemm.pyd

 

 

3. 공통 Binding 유틸 ( common.hpp )

3.1 역할

• torch Tensor 검사 로직 중복 제거
• aicf::Stream 기본값 통일
• aicf::Status 처리 공통화

3.2 핵심 내용

inline bool is_cuda_f32_contig(const torch::Tensor& t) {
    return t.is_cuda() && t.scalar_type() == at::kFloat && t.is_contiguous();
}

inline aicf::Stream default_stream() {
    aicf::Stream s{};
    s.handle = nullptr; // default CUDA stream
    return s;
}

• 바인딩 레벨에서는 Torch semantics 를 흉내내지 않는다
• contiguous + f32 + cuda 만 지원
• 그 외는 Pythno backend 만 torch fallback 처리

 

4. Op 별 pybind 바인딩 코드 구조

4.1 기본 패턴 ( add.cpp )

PYBIND11_MODULE(add, m) {
    m.def("add_f32", [](const torch::Tensor& a,
                        const torch::Tensor& b,
                        torch::Tensor& out) -> bool {

        if (!is_cuda_f32_contig(a) || !is_cuda_f32_contig(b)) return false;

        int N = a.numel();
        auto st = aicf::cuda::add_f32(
            a.data_ptr<float>(),
            b.data_ptr<float>(),
            out.data_ptr<float>(),
            N,
            default_stream());

        return st == aicf::Status::Ok;
    });
}

4.2 설계 원칙

• pybind 모듈 이름 = op 이름
• C++ public API 만 호출
• Python 에서는 반환값으로 성공 여부 판단

 

5. CMake 를 통한 op 별 pyd 모듈 생성

5.1 핵심 아이디어

CMake 함수 하나로 op 별 pybind MODULE 타겟을 반복 생성한다

function(aicf_add_op_pybind target_name source_file output_name)
  add_library(${target_name} MODULE ${source_file})
  ...
endfunction()

 

5.2 Torch + Python 연동 준비

find_package(Python REQUIRED COMPONENTS Interpreter Development)
find_package(Torch REQUIRED)

Windows 에서는 추가 필요

find_library(TORCH_PYTHON_LIBRARY
  NAMES torch_python
  PATHS "<python>/Lib/site-packages/torch/lib"
)

pybind11 에서 torch::Tensor caster 사용 시 필수

 

5.3 op 별 pybind 타겟 정의

aicf_add_op_pybind(aicf_cuda_add  src/python_bindings/add.cpp  add)
aicf_add_op_pybind(aicf_cuda_relu src/python_bindings/relu.cpp relu)
aicf_add_op_pybind(aicf_cuda_gemm src/python_bindings/gemm.cpp gemm)

각 op 별 .pyd 파일 생성

 

5.4 출력 디렉토리 설정

set_target_properties(${target_name} PROPERTIES
  PREFIX ""
  OUTPUT_NAME "${output_name}"
  LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/python/aicf_cuda"
  RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/python/aicf_cuda"
)

 

 

6. Python 패키지화 ( aicf_cuda / __init__.py )

from . import add, relu, gemm
__all__ = ["add", "relu", "gemm"]

이 파일을 build / python / aicf_cuda 에 복사함으로써

import aicf_cuda
aicf_cuda.add.add_f32(...)

형태도 지원

 

7. 빌드 및 테스트 흐름 요약

7.1 빌드

cmake .. -G Ninja
ninja aicf_cuda_add aicf_cuda_relu aicf_cuda_gemm

7.2 Python 연동 테스트

set PYTHONPATH=%CD%\build\python
python examples/python/test_aicf_pyd.py

7.3 확인 포인트

• import 성공 여부
• torch 결과와 수치 일치
• ncu 에서 실제 CUDA kernel 실행 확인

 

8. 이 구조의 장점

• ops 증가 시 바인딩 파일 선형 증가
• 특정 op 만 재빌드 가능
• CUDA / Python / IR 구조가 동일한 축으로 정렬도미
• 향후
  • kernel variant 추가
  • op capability table 기반 자동 바인딩 생성
  • graph capture 시 op 단위 replay

로 확장 가능

 

9. 다음 단계

• AicfBackend 에서 op 별 pyd 호출 + fallback 정책 고정
• registry 기반 dispatcher 와 pybind 경로 연결
• NVTX 계층화
• kernel variant 선택 로직과  Python attrs 연동
•  
•