#pragma kernel ScaleBias 	
#pragma kernel ScaleBias_CNyx
#pragma kernel Upsample2D
#pragma kernel AvgPool2D
#pragma kernel MaxPool2D
#pragma kernel AvgPool2D_NoPads
#pragma kernel MaxPool2D_NoPads
//#pragma kernel MaxPool2D_Pool2x2_NoPads
#pragma kernel GlobalAvgPool2D
#pragma kernel InstanceNorm
#pragma kernel Copy

#include "Tensor.cginc"

TENSOR_DECL(X)
TENSOR_DECL(W)
TENSOR_DECL(B)
TENSOR_DECL(WBK)
TENSOR_DECL_RW(O)

uint4 _Pool;
uint4 _Stride;
uint4 _Pad;
float _Alpha;

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void ScaleBias(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.width, O.height);
	TENSOR_SHARED2_ARGS4(X, W, B, WBK, O);

	uint c = dispatchThreadID.x;
	uint x = dispatchThreadID.y;
	uint y = dispatchThreadID.z;

	if (c >= O.channels) return;
	if (x >= O.width) return;
	if (y >= O.height) return;

	float bias = B.Get(0, 0, 0, c);
	float scale = W.Get(0, 0, 0, c);

	for (uint n = 0; n < X.batch; ++n)
	{
		float v = X.Get(n, y, x, c);
		v = v * scale + bias;
		O.Set(n, y, x, c, v);
	}
}

NUMTHREADS((16,16,1), (16,8,1), (16,4,1))
void ScaleBias_CNyx(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.batch * O.height * O.width, 1);
	TENSOR_SHARED2_ARGS4(X, W, B, WBK, O);

	uint c = dispatchThreadID.x;
	uint nyx = dispatchThreadID.y;

	uint x = nyx % X.width;
	uint ny = nyx / X.width;
	uint y = ny % X.height;
	uint n = ny / X.height;

	if (c >= X.channels) return;
	if (n >= X.batch) return;

	float bias = B.Get(0, 0, 0, c);
	float scale = W.Get(0, 0, 0, c);

	float v = X.Get(n, y, x, c);
	v = v * scale + bias;
	O.Set(n, y, x, c, v);
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void Upsample2D(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	// NOTE: dispatched over X (not O)
	DISPATCH_ARGS(X.channels, X.width, X.height);
	TENSOR_ARGS2(X, O);

	uint c = dispatchThreadID.x;
	uint x = dispatchThreadID.y;
	uint y = dispatchThreadID.z;

	if (c >= X.channels) return;
	if (x >= X.width) return;
	if (y >= X.height) return;

	for (uint n = 0; n < O.batch; ++n)
	{
		float v = X.Get(n, y, x, c);

		for (uint dy = 0; dy < _Pool.y; ++dy)
			for (uint dx = 0; dx < _Pool.x; ++dx)
			{
				uint oy = y * _Pool.y + dy;
				uint ox = x * _Pool.x + dx;
				O.Set(n, oy, ox, c, v);
			}
	}
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void MaxPool2D(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.width, O.height);
	TENSOR_ARGS2(X, O);

	uint c = dispatchThreadID.x;
	uint x = dispatchThreadID.y;
	uint y = dispatchThreadID.z;

	if (c >= O.channels) return;
	if (x >= O.width) return;
	if (y >= O.height) return;

	for (uint n = 0; n < X.batch; ++n)
	{
		float maxV = -FLT_MAX;
		for (uint dy = 0; dy < _Pool.y; ++dy)
			for (uint dx = 0; dx < _Pool.x; ++dx)
			{
				uint oy = y * _Stride.y + dy;
				uint ox = x * _Stride.x + dx;

				bool mask = (oy >= _Pad.y) && (ox >= _Pad.x) && (oy - _Pad.w < X.height) && (ox - _Pad.z < X.width);
				float v = (mask)? X.Get(n, oy - _Pad.y, ox - _Pad.x, c): 0;
				maxV = max(v, maxV);
			}
		
		O.Set(n, y, x, c, maxV);
	}
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void AvgPool2D(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.width, O.height);
	TENSOR_ARGS2(X, O);

	uint c = dispatchThreadID.x;
	uint x = dispatchThreadID.y;
	uint y = dispatchThreadID.z;

	if (c >= O.channels) return;
	if (x >= O.width) return;
	if (y >= O.height) return;

	for (uint n = 0; n < X.batch; ++n)
	{
		float acc = 0;
		float counter = 0;
		for (uint dy = 0; dy < _Pool.y; ++dy)
			for (uint dx = 0; dx < _Pool.x; ++dx)
			{
				uint oy = y * _Stride.y + dy;
				uint ox = x * _Stride.x + dx;

				bool mask = (oy >= _Pad.y) && (ox >= _Pad.x) && (oy - _Pad.w < X.height) && (ox - _Pad.z < X.width);
				acc += (mask)? X.Get(n, oy - _Pad.y, ox - _Pad.x, c): 0;
				counter += (mask)? 1: 0;
			}
		
		acc /= counter;
		O.Set(n, y, x, c, acc);
	}
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void MaxPool2D_NoPads(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.width, O.height);
	TENSOR_ARGS2(X, O);

	uint c = dispatchThreadID.x;
	uint x = dispatchThreadID.y;
	uint y = dispatchThreadID.z;

	if (c >= O.channels) return;
	if (x >= O.width) return;
	if (y >= O.height) return;

	for (uint n = 0; n < X.batch; ++n)
	{
		float maxV = -FLT_MAX;
		for (uint dy = 0; dy < _Pool[1]; ++dy)
			for (uint dx = 0; dx < _Pool[0]; ++dx)
			{
				float v = X.Get(n, y * _Stride[1] + dy, x * _Stride[0] + dx, c);
				maxV = max(v, maxV);
			}
		
		O.Set(n, y, x, c, maxV);
	}
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void AvgPool2D_NoPads(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.width, O.height);
	TENSOR_ARGS2(X, O);

	uint c = dispatchThreadID.x;
	uint x = dispatchThreadID.y;
	uint y = dispatchThreadID.z;

	if (c >= O.channels) return;
	if (x >= O.width) return;
	if (y >= O.height) return;

	float invPoolSize = 1.0f / (_Pool[0] * _Pool[1]);
	for (uint n = 0; n < X.batch; ++n)
	{
		float v = 0;
		for (uint dy = 0; dy < _Pool[1]; ++dy)
			for (uint dx = 0; dx < _Pool[0]; ++dx)
				v += X.Get(n, y * _Stride[1] + dy, x * _Stride[0] + dx, c) * invPoolSize;

		O.Set(n, y, x, c, v);
	}
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
//NUMTHREADS((16,4,4), (16,4,2), (16,2,2))
void MaxPool2D_Pool2x2_NoPads(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.width, O.height);
	TENSOR_ARGS2(X, O);

	uint c = dispatchThreadID.x;
	uint x = dispatchThreadID.y;
	uint y = dispatchThreadID.z;

	if (c >= O.channels) return;
	if (x >= O.width) return;
	if (y >= O.height) return;

	for (uint n = 0; n < X.batch; ++n)
	{
		float v0 = X.Get(n, y*2,   x*2,   c);
		float v1 = X.Get(n, y*2+1, x*2,   c);
		float v2 = X.Get(n, y*2,   x*2+1, c);
		float v3 = X.Get(n, y*2+1, x*2+1, c);
		float v = max(v0, max(v1, max(v2, v3)));

		O.Set(n, y, x, c, v);
	}
}

[numthreads(32,1,1)]
void GlobalAvgPool2D(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, 1, 1);
	TENSOR_ARGS2(X, O);

	uint c = dispatchThreadID.x;
	if (c >= O.channels) return;
	//ASSERT(X.batch == O.batch)

	for (uint n = 0; n < X.batch; ++n)
	{
		float v = 0;
		for (uint y = 0; y < X.height; ++y)
			for (uint x = 0; x < X.width; ++x)
				v += X.Get(n, y, x, c);
		
		v /= (X.height * X.width);
		O.Set(n, 0, 0, c, v);
	}
}

[numthreads(64,1,1)]
void InstanceNorm(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, 1, 1);
	TENSOR_SHARED2_ARGS4(X, W, B, WBK, O);

	uint c = dispatchThreadID.x;
	if (c >= O.channels) return;
	//ASSERT(X.shape == O.shape)

	float gamma = W.Get(0, 0, 0, c);
	float beta = B.Get(0, 0, 0, c);

	for (uint n = 0; n < O.batch; ++n)
	{
		uint x, y;
		// calc mean
		float acc = 0;
		for (y = 0; y < O.height; ++y)
			for (x = 0; x < O.width; ++x)
				acc += X.Get(n, y, x, c);
		float mean = acc / (O.width * O.height);

		// calc variance
		acc = 0;
		for (y = 0; y < O.height; ++y)
			for (x = 0; x < O.width; ++x)
			{
				float delta = X.Get(n, y, x, c) - mean;
				acc += delta * delta;
			}
		float var = acc / (O.width * O.height);

		// normalization factor
		float invNormFactor = 1 / sqrt(var + FLT_EPSILON);

		float scale = gamma * invNormFactor;
		float bias = beta - gamma * mean * invNormFactor;

		// apply normalization
		for (y = 0; y < O.height; ++y)
			for (x = 0; x < O.width; ++x)
			{
				float v = X.Get(n, y, x, c);
				v = v * scale + bias;
				O.Set(n, y, x, c, v);
			}
	}
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void Copy(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	// NOTE: dispatched over X (not O)
	DISPATCH_ARGS(X.channels, X.width, X.height);
	TENSOR_ARGS2(X, O);

	uint c = dispatchThreadID.x;	uint x = dispatchThreadID.y;	uint y = dispatchThreadID.z;
	if (c >= X.channels) return;	if (x >= X.width) return;		if (y >= X.height) return;

	for (uint n = 0; n < X.batch; ++n)
	{
		float v = X.Get(n, y, x, c);
		O.Set(n + _Pad[0], y + _Pad[1], x + _Pad[2], c + _Pad[3], v);
	}
}
