3D Gaussian Splatting의 submodule - simple-knn 수정하기

knn 함수를 수정할 필요가 있어서 simple-knn 모듈 내의 spatial.h, spatial.cu에 knn2 함수를 추가해줬다.

 

// spatial.h
#include <torch/extension.h>

torch::Tensor distCUDA2(const torch::Tensor& points);
torch::Tensor knn2(const torch::Tensor& points);

 

그리고 이를 빌드하기 위해 ext.cpp를 아래처럼 수정했다

#include <torch/extension.h>
#include "spatial.h"

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("distCUDA2", &distCUDA2);
  m.def("knn2", &knn2); // 추가
}

 

이후 gaussian-splatting 루트 디렉토리에서 submodules/simple-knn으로 이동한 후, pip install -e를 통해 setup.py를 실행하면 된다

python3 -m pip install -e .

 

그냥 pip install -e .를 하니 설치가 제대로 진행되지 않았는데, 저번에 겪었던 문제랑 비슷하여 python3 -m "명령"을 이용하니 해결.