#include "conv.h"

constexpr int threads_per_bloc = 16;
constexpr int T = threads_per_bloc;

//
// example: CUDA_CHECK( cudaMalloc(dx, x, N*sizeof(int) );
//
#define CUDA_CHECK(code)                                                                                               \
	{ cuda_check((code), __FILE__, __LINE__); }
inline void cuda_check(cudaError_t code, const char* file, int line) {
	if (code != cudaSuccess) {
		std::cout << file << ':' << line << ": [CUDA ERROR] " << cudaGetErrorString(code) << std::endl;
		std::abort();
	}
}

#define RANGE(X, FROM, TO)                                                                                             \
	long X = FROM;                                                                                                     \
	X < TO;                                                                                                            \
	X += 1

#define THREAD_ID_X() (blockIdx.x * blockDim.x + threadIdx.x);

#define DBG(X, FMT) printf(#X ": %" FMT "\n", X);
#define DBG_S(X) #X << ": " << X << ", "

//
// 1D convolution
// - x: input array of size N
// - y: kernel of odd size M
//
// CPU
//
std::vector<int> conv1(const std::vector<int>& x, const std::vector<int>& y) {
	//
	// step 01
	//
	const int N = x.size();
	const int M = y.size();
	const int P = (M - 1) / 2;

	auto z = std::vector<int>(N);
	for (RANGE(result_index, 0, x.size())) {
		auto result = 0;
		for (RANGE(y_index, 0, y.size())) {
			auto x_index = result_index - P + y_index;
			if (x_index < 0 || x_index >= x.size()) continue;
			result += x.at(x_index) * y.at(y_index);
		}
		z.at(result_index) = result;
	}

	return z;
}

namespace kernel {

//
// step 02
//
__global__ void conv2(const int* dx, const int* dy, int x_length, int y_length, int* dz) {
	auto thread_id = (long)THREAD_ID_X();
	if (thread_id >= x_length) return;
	auto offset = (y_length - 1) / 2;

	auto result_index = thread_id;
	auto result = 0;
	for (RANGE(y_index, 0, y_length)) {
		auto x_index = result_index - offset + y_index;
		if (x_index < 0 || x_index >= x_length) continue;
		result += dx[x_index] * dy[y_index];
	}
	dz[result_index] = result;
}

} // namespace kernel

//
// 1D convolution
// - x: input array of size N
// - y: kernel of odd size M
//
// GPU (naive)
//
std::vector<int> conv2(const std::vector<int>& x, const std::vector<int>& y) {
	//
	// step 03
	//
	auto dx = (int*)nullptr;
	auto size_dx = x.size() * sizeof(int);
	cudaMalloc(&dx, size_dx);
	cudaMemcpy(dx, x.data(), size_dx, cudaMemcpyHostToDevice);

	auto dy = (int*)nullptr;
	auto size_dy = y.size() * sizeof(int);
	cudaMalloc(&dy, size_dy);
	cudaMemcpy(dy, y.data(), size_dy, cudaMemcpyHostToDevice);

	auto dz = (int*)nullptr;
	auto size_dz = x.size() * sizeof(int);
	cudaMalloc(&dz, size_dz);

	auto blocks = x.size() / threads_per_bloc + 1;
	kernel::conv2<<<blocks, threads_per_bloc>>>(dx, dy, x.size(), y.size(), dz);
	cudaFree(dx);
	cudaFree(dy);

	auto z = std::vector<int>(x.size());
	cudaMemcpy(z.data(), dz, size_dz, cudaMemcpyDeviceToHost);
	cudaFree(dz);

	return z;
}

namespace kernel {

//
// step 04
//
__global__ void conv3(const int* dx, const int* dy, int x_length, int y_length, int* dz) {
	__shared__ int buffer[T];
	auto thread_id = (long)THREAD_ID_X();
	if (thread_id >= x_length) return;

	buffer[thread_id % T] = dx[thread_id];
	__syncthreads();
	auto buffer_lower_x_index = (thread_id / T) * T;
	auto buffer_upper_x_index = buffer_lower_x_index + T;
	auto offset = (y_length - 1) / 2;

	auto result_index = thread_id;
	auto result = 0;
	for (RANGE(y_index, 0, y_length)) {
		auto x_index = result_index - offset + y_index;
		if (x_index < 0 || x_index >= x_length) continue;
		auto in_buffer = x_index >= buffer_lower_x_index && x_index < buffer_upper_x_index;
		if (in_buffer) {
			auto buff_index = x_index - buffer_lower_x_index;
			result += buffer[buff_index] * dy[y_index];
		} else result += dx[x_index] * dy[y_index];
	}
	dz[result_index] = result;
}

} // namespace kernel

//
// 1D convolution
// - x: input array of size N
// - y: kernel of odd size M
//
// GPU (optimized)
//
std::vector<int> conv3(const std::vector<int>& x, const std::vector<int>& y) {
	//
	// step 05
	//
	auto dx = (int*)nullptr;
	auto size_dx = x.size() * sizeof(int);
	cudaMalloc(&dx, size_dx);
	cudaMemcpy(dx, x.data(), size_dx, cudaMemcpyHostToDevice);

	auto dy = (int*)nullptr;
	auto size_dy = y.size() * sizeof(int);
	cudaMalloc(&dy, size_dy);
	cudaMemcpy(dy, y.data(), size_dy, cudaMemcpyHostToDevice);

	auto dz = (int*)nullptr;
	auto size_dz = x.size() * sizeof(int);
	cudaMalloc(&dz, size_dz);

	auto blocks = x.size() / threads_per_bloc + 1;
	kernel::conv3<<<blocks, threads_per_bloc>>>(dx, dy, x.size(), y.size(), dz);
	cudaFree(dx);
	cudaFree(dy);

	auto z = std::vector<int>(x.size());
	cudaMemcpy(z.data(), dz, size_dz, cudaMemcpyDeviceToHost);
	cudaFree(dz);

	return z;
}
