freqencoder/src/freqencoder.cu

#include <stdint.h>

#include <cuda.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>

#include <ATen/cuda/CUDAContext.h>
#include <torch/torch.h>

#include <algorithm>
#include <stdexcept>

#include <cstdio>


#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor")
#define CHECK_IS_INT(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, #x " must be an int tensor")
#define CHECK_IS_FLOATING(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Float || x.scalar_type() == at::ScalarType::Half || x.scalar_type() == at::ScalarType::Double, #x " must be a floating tensor")

inline constexpr __device__ float PI() { return 3.141592653589793f; }

template <typename T>
__host__ __device__ T div_round_up(T val, T divisor) {
	return (val + divisor - 1) / divisor;
}

// inputs: [B, D]
// outputs: [B, C], C = D + D * deg * 2
__global__ void kernel_freq(
    const float * __restrict__ inputs, 
    uint32_t B, uint32_t D, uint32_t deg, uint32_t C,
    float * outputs
) {
	// parallel on per-element
	const uint32_t t = threadIdx.x + blockIdx.x * blockDim.x;
	if (t >= B * C) return;

	// get index
	const uint32_t b = t / C;
	const uint32_t c = t - b * C; // t % C;

	// locate
	inputs += b * D;
	outputs += t;

	// write self
	if (c < D) {
		outputs[0] = inputs[c];
	// write freq
	} else {
		const uint32_t col = c / D - 1;
		const uint32_t d = c % D;
		const uint32_t freq = col / 2;
		const float phase_shift = (col % 2) * (PI() / 2);
		outputs[0] = __sinf(scalbnf(inputs[d], freq) + phase_shift);
	}
}

// grad: [B, C], C = D + D * deg * 2
// outputs: [B, C]
// grad_inputs: [B, D]
__global__ void kernel_freq_backward(
    const float * __restrict__ grad,
	const float * __restrict__ outputs,
    uint32_t B, uint32_t D, uint32_t deg, uint32_t C,
    float * grad_inputs
) {
	// parallel on per-element
	const uint32_t t = threadIdx.x + blockIdx.x * blockDim.x;
	if (t >= B * D) return;

	const uint32_t b = t / D;
	const uint32_t d = t - b * D; // t % D;

	// locate
	grad += b * C;
	outputs += b * C;
	grad_inputs += t;

	// register 
	float result = grad[d];
	grad += D;
	outputs += D;

	for (uint32_t f = 0; f < deg; f++) {
		result += scalbnf(1.0f, f) * (grad[d] * outputs[D + d] - grad[D + d] * outputs[d]);
		grad += 2 * D;
		outputs += 2 * D;
	}

	// write
	grad_inputs[0] = result;
}


void freq_encode_forward(at::Tensor inputs, const uint32_t B, const uint32_t D, const uint32_t deg, const uint32_t C, at::Tensor outputs) {
    CHECK_CUDA(inputs);
    CHECK_CUDA(outputs);
    
    CHECK_CONTIGUOUS(inputs);
    CHECK_CONTIGUOUS(outputs);

    CHECK_IS_FLOATING(inputs);
    CHECK_IS_FLOATING(outputs);

	static constexpr uint32_t N_THREADS = 128;

	kernel_freq<<<div_round_up(B * C, N_THREADS), N_THREADS>>>(inputs.data_ptr<float>(), B, D, deg, C, outputs.data_ptr<float>());
}


void freq_encode_backward(at::Tensor grad, at::Tensor outputs, const uint32_t B, const uint32_t D, const uint32_t deg, const uint32_t C, at::Tensor grad_inputs) {
    CHECK_CUDA(grad);
    CHECK_CUDA(outputs);
    CHECK_CUDA(grad_inputs);
    
    CHECK_CONTIGUOUS(grad);
    CHECK_CONTIGUOUS(outputs);
    CHECK_CONTIGUOUS(grad_inputs);

    CHECK_IS_FLOATING(grad);
    CHECK_IS_FLOATING(outputs);
    CHECK_IS_FLOATING(grad_inputs);

	static constexpr uint32_t N_THREADS = 128;

	kernel_freq_backward<<<div_round_up(B * D, N_THREADS), N_THREADS>>>(grad.data_ptr<float>(), outputs.data_ptr<float>(), B, D, deg, C, grad_inputs.data_ptr<float>());
}