Kernels:

phanerozoic
/

fhe-ntt

Kernel card Files Files and versions

fhe-ntt / torch-ext /torch_binding.cpp

CharlesCNorton

Add kernel source and model card

37ecc0f 26 days ago

748 Bytes

	#include <torch/library.h>

	#include "registration.h"
	#include "torch_binding.h"

	TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
	ops.def("fhe_ntt_forward(Tensor(a!) a, Tensor tw, Tensor tw_shoup, Tensor primes) -> ()");
	ops.def("fhe_ntt_inverse(Tensor(a!) a, Tensor itw, Tensor itw_shoup, Tensor ninv, "
	"Tensor ninv_shoup, Tensor primes) -> ()");
	ops.def("fhe_pointwise_mul(Tensor(a!) out, Tensor a, Tensor b, Tensor primes) -> ()");

	#if defined(CUDA_KERNEL) \|\| defined(ROCM_KERNEL)
	ops.impl("fhe_ntt_forward", torch::kCUDA, &fhe_ntt_forward);
	ops.impl("fhe_ntt_inverse", torch::kCUDA, &fhe_ntt_inverse);
	ops.impl("fhe_pointwise_mul", torch::kCUDA, &fhe_pointwise_mul);
	#endif
	}

	REGISTER_EXTENSION(TORCH_EXTENSION_NAME)