Aluode/PerceptionLabPortable

1#include <torch/extension.h>2#include <ATen/ATen.h>3#include <vector>4 5#define min(a, b) ((a)<(b)?(a):(b))6#define max(a, b) ((a)>(b)?(a):(b))7 8std::vector<at::Tensor> index_max_kernel(9  at::Tensor index_vals,10  at::Tensor indices,11  int A_num_block,12  int B_num_block13);14 15at::Tensor mm_to_sparse_kernel(16  at::Tensor dense_A,17  at::Tensor dense_B,18  at::Tensor indices19);20 21at::Tensor sparse_dense_mm_kernel(22  at::Tensor sparse_A,23  at::Tensor indices,24  at::Tensor dense_B,25  int A_num_block26);27 28at::Tensor reduce_sum_kernel(29  at::Tensor sparse_A,30  at::Tensor indices,31  int A_num_block,32  int B_num_block33);34 35at::Tensor scatter_kernel(36  at::Tensor dense_A,37  at::Tensor indices,38  int B_num_block39);40