CUDA build

CUDA extension은 GPU 가속을 위해 사용자 정의 CUDA 커널을 작성하고 이를 Python 코드에서 사용할 수 있도록 만든 확장 모듈

torch.utils.cpp_extension을 사용하여 CUDA 커널을 Python 코드에서 호출할 수 있도록 빌드함

 

• CUDA extension은 사용자 정의 CUDA 커널을 PyTorch와 연결하기 위한 모듈
• 일반적인 Python/CPU 코드로는 수백만 개의 점에 대해 이웃을 찾는 연산이 너무 느리기 때문에, CUDA 커널을 직접 작성하여 속도를 높임
• 병렬 연산을 통해 GPU 가속이 가능해지며, PyTorch의 autograd 기능과도 통합할 수 있음

 

• C++ Code (ball_query.cpp): This file contains the interface between Python and CUDA. It defines how Python calls the underlying CUDA kernel.
• CUDA Code (ball_query_cuda.cu): This file contains the actual GPU kernel implementation, where computationally intensive tasks are performed.

 

 

코드 주요 부분 설명

🔍 1. 전역 변수 extern THCState *state;

무엇인가?

• THCState는 구형 PyTorch의 내부 상태를 나타내는 구조체입니다.
• CUDA 메모리 할당, CUDA 스트림 관리 및 장치 메모리 풀링 등을 관리합니다.
• **구형 PyTorch (1.0 이전)**에서는 이 상태를 명시적으로 참조해야 했으나, 현재는 필요하지 않습니다.
• 현대 PyTorch에서는 at::cuda::getCurrentCUDAStream() 같은 인터페이스로 대체되었습니다.\

왜 선언하는가?

• 과거의 잔재입니다. => 그렇기 때문에 CUDA 11.3 이상부터 얘가 있으면 아예 빌드가 되지 않음
• 예전의 PyTorch 코드는 THCState를 통해 CUDA 장치 정보를 공유했지만, 최신 PyTorch에서는 더 이상 필요하지 않습니다.
• 따라서, 이 코드는 불필요한 부분일 가능성이 높으며, 삭제해도 잘 작동할 수 있습니다.
• 결론: extern THCState *state;는 CUDA 상태 관리에 사용되었으나, 최신 PyTorch에서는 at::cuda::getCurrentCUDAStream()으로 대체되었습니다. 현대 PyTorch 코드에서는 사용하지 않습니다.

• .

---

🔍 5. PyBind11 모듈 연결

이 부분은 C++과 Python을 연결하기 위해 PYBIND11_MODULE을 사용합니다.

cpp

코드 복사

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("ball_query_wrapper", &ball_query_wrapper, "ball_query_wrapper"); }

 

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("ball_query_wrapper", &ball_query_wrapper, "ball_query_wrapper");
}

 

• TORCH_EXTENSION_NAME: setup.py의 name과 동일한 이름으로 정의됩니다.
• m.def(...): Python 함수로 연결됩니다.
  • Python에서는 import ball_query 후 ball_query.ball_query_wrapper(...)로 함수를 호출할 수 있습니다.