det-train / torch-ext /torch_binding.cpp

det-train v1: complete deterministic-training kernel

c8048c9 verified 26 days ago

496 Bytes

	#include <torch/library.h>

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

	TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
	ops.def("det_gemm(Tensor! out, Tensor x, Tensor w) -> ()");
	ops.def("det_gemm_partial(Tensor! E, Tensor! S, Tensor x, Tensor w) -> ()");

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

	REGISTER_EXTENSION(TORCH_EXTENSION_NAME)