#pragma kernel Relu
#pragma kernel Relu_CNyx
#pragma kernel Relu_Nyxc
#pragma kernel Relu6
#pragma kernel Relu6_CNyx
#pragma kernel Relu6_Nyxc
#pragma kernel Tanh
#pragma kernel Tanh_CNyx
#pragma kernel Tanh_Nyxc
#pragma kernel Swish
#pragma kernel Swish_CNyx
#pragma kernel Swish_Nyxc
#pragma kernel Sigmoid
#pragma kernel Sigmoid_CNyx
#pragma kernel Sigmoid_Nyxc
#pragma kernel Elu
#pragma kernel Elu_CNyx
#pragma kernel Elu_Nyxc
#pragma kernel LeakyRelu
#pragma kernel LeakyRelu_CNyx
#pragma kernel LeakyRelu_Nyxc
#pragma kernel Exp
#pragma kernel Exp_CNyx
#pragma kernel Exp_Nyxc
#pragma kernel Pow
#pragma kernel Pow_CNyx
#pragma kernel Pow_Nyxc
#pragma kernel Softmax

#include "Tensor.cginc"

TENSOR_DECL(X)
TENSOR_DECL_RW(O)

float _Alpha;

float relu(float v)
{
	return 0.5f * (v + abs(v));
}

float relu6(float v)
{
	return min(max(0, v), 6);
}

float swish(float v)
{
	return v / (1.f + exp(-v));
}

float sigmoid(float v)
{
	return 1.f / (1.f + exp(-v));
}

float elu(float v)
{
	if (v <= 0)
		v = _Alpha * (exp(v) - 1);
	return v;
}

float lrelu(float v)
{
	return max(v, _Alpha * v);	
}

float signed_pow(float f, float e)
{
	// handle negative f
	float v = pow(abs(f), e);
	float s = (e % 2 == 1) ?
		sign(f):	// exponent is odd  => sign(f) * pow(abs(f), e)
		1;			// exponent is even => pow(abs(f), e)
	return v * s;
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void Relu(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 v = X.Get(n, y, x, c);
		v = relu(v);
		O.Set(n, y, x, c, v);
	}
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void Relu6(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 v = X.Get(n, y, x, c);
		v = relu6(v);
		O.Set(n, y, x, c, v);
	}
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void Tanh(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 v = X.Get(n, y, x, c);
		v = tanh(v);
		O.Set(n, y, x, c, v);
	}
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
 void Sigmoid(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 v = X.Get(n, y, x, c);
 		v = sigmoid(v);
 		O.Set(n, y, x, c, v);
 	}
 }
 
 NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void Swish(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 v = X.Get(n, y, x, c);
		v = swish(v);
		O.Set(n, y, x, c, v);
	}
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void Elu(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 v = X.Get(n, y, x, c);
		v = elu(v);
		O.Set(n, y, x, c, v);
	}
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void LeakyRelu(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 v = X.Get(n, y, x, c);
		v = lrelu(v);
		O.Set(n, y, x, c, v);
	}
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void Exp(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 v = X.Get(n, y, x, c);
		v = exp(v);
		O.Set(n, y, x, c, v);
	}
}

NUMTHREADS((4,8,8), (4,8,4), (4,4,4))
void Pow(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 v = X.Get(n, y, x, c);
		v = signed_pow(v, _Alpha);
		O.Set(n, y, x, c, v);
	}
}

NUMTHREADS((16,16,1), (16,8,1), (16,4,1))
void Relu_CNyx(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.batch * O.height * O.width, 1);
	TENSOR_ARGS2(X, 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 v = X.Get(n, y, x, c);
	v = relu(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((512,1,1), (128,1,1), (64,1,1))
void Relu_Nyxc(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.batch * O.height * O.width * O.channels, 1, 1)
	TENSOR_ARGS2(X, O);

	uint nyxc = dispatchThreadID.x;

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

	if (n >= X.batch) return;

	float v = X.Get(n, y, x, c);
	v = relu(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((16,16,1), (16,8,1), (16,4,1))
void Relu6_CNyx(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.batch * O.height * O.width, 1);
	TENSOR_ARGS2(X, 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 v = X.Get(n, y, x, c);
	v = relu6(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((512,1,1), (128,1,1), (64,1,1))
void Relu6_Nyxc(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.batch * O.height * O.width * O.channels, 1, 1)
	TENSOR_ARGS2(X, O);

	uint nyxc = dispatchThreadID.x;

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

	if (n >= X.batch) return;

	float v = X.Get(n, y, x, c);
	v = relu6(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((16,16,1), (16,8,1), (16,4,1))
void Tanh_CNyx(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.batch * O.height * O.width, 1);
	TENSOR_ARGS2(X, 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 v = X.Get(n, y, x, c);
	v = tanh(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((512,1,1), (128,1,1), (64,1,1))
void Tanh_Nyxc(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.batch * O.height * O.width * O.channels, 1, 1)
	TENSOR_ARGS2(X, O);

	uint nyxc = dispatchThreadID.x;

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

	if (n >= X.batch) return;

	float v = X.Get(n, y, x, c);
	v = tanh(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((16,16,1), (16,8,1), (16,4,1))
void Sigmoid_CNyx(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.batch * O.height * O.width, 1);
	TENSOR_ARGS2(X, 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 v = X.Get(n, y, x, c);
	v = sigmoid(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((512,1,1), (128,1,1), (64,1,1))
void Sigmoid_Nyxc(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.batch * O.height * O.width * O.channels, 1, 1)
	TENSOR_ARGS2(X, O);

	uint nyxc = dispatchThreadID.x;

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

	if (n >= X.batch) return;

	float v = X.Get(n, y, x, c);
	v = sigmoid(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((16,16,1), (16,8,1), (16,4,1))
void Swish_CNyx(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.batch * O.height * O.width, 1);
	TENSOR_ARGS2(X, 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 v = X.Get(n, y, x, c);
	v = swish(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((512,1,1), (128,1,1), (64,1,1))
void Swish_Nyxc(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.batch * O.height * O.width * O.channels, 1, 1)
	TENSOR_ARGS2(X, O);

	uint nyxc = dispatchThreadID.x;

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

	if (n >= X.batch) return;

	float v = X.Get(n, y, x, c);
	v = swish(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((16,16,1), (16,8,1), (16,4,1))
void Elu_CNyx(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.batch * O.height * O.width, 1);
	TENSOR_ARGS2(X, 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 v = X.Get(n, y, x, c);
	v = elu(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((512,1,1), (128,1,1), (64,1,1))
void Elu_Nyxc(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.batch * O.height * O.width * O.channels, 1, 1)
	TENSOR_ARGS2(X, O);

	uint nyxc = dispatchThreadID.x;

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

	if (n >= X.batch) return;

	float v = X.Get(n, y, x, c);
	v = elu(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((16,16,1), (16,8,1), (16,4,1))
void LeakyRelu_CNyx(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.batch * O.height * O.width, 1);
	TENSOR_ARGS2(X, 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 v = X.Get(n, y, x, c);
	v = lrelu(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((512,1,1), (128,1,1), (64,1,1))
void LeakyRelu_Nyxc(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.batch * O.height * O.width * O.channels, 1, 1)
	TENSOR_ARGS2(X, O);

	uint nyxc = dispatchThreadID.x;

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

	if (n >= X.batch) return;

	float v = X.Get(n, y, x, c);
	v = lrelu(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((16,16,1), (16,8,1), (16,4,1))
void Exp_CNyx(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.batch * O.height * O.width, 1);
	TENSOR_ARGS2(X, 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 v = X.Get(n, y, x, c);
	v = exp(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((512,1,1), (128,1,1), (64,1,1))
void Exp_Nyxc(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.batch * O.height * O.width * O.channels, 1, 1)
	TENSOR_ARGS2(X, O);

	uint nyxc = dispatchThreadID.x;

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

	if (n >= X.batch) return;

	float v = X.Get(n, y, x, c);
	v = exp(v);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((16,16,1), (16,8,1), (16,4,1))
void Pow_CNyx(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.channels, O.batch * O.height * O.width, 1);
	TENSOR_ARGS2(X, 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 v = X.Get(n, y, x, c);
	v = signed_pow(v, _Alpha);
	O.Set(n, y, x, c, v);
}

NUMTHREADS((512,1,1), (128,1,1), (64,1,1))
void Pow_Nyxc(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.batch * O.height * O.width * O.channels, 1, 1)
	TENSOR_ARGS2(X, O);

	uint nyxc = dispatchThreadID.x;

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

	if (n >= X.batch) return;

	float v = X.Get(n, y, x, c);
	v = signed_pow(v, _Alpha);
	O.Set(n, y, x, c, v);
}


NUMTHREADS((64,4,1), (64,2,1), (64,1,1))
void Softmax(uint3 dispatchThreadID : SV_DispatchThreadID)
{
	DISPATCH_ARGS(O.flatWidth, O.flatHeight, 1);
	TENSOR_ARGS2(X, O);

	uint x = dispatchThreadID.x;
	uint y = dispatchThreadID.y;

	if (x >= O.GetFlatWidth()) return;
	if (y >= O.GetFlatHeight()) return;

	float maxV = -FLT_MAX;
	for (uint i = 0; i < X.GetFlatWidth(); ++i)
	{
		float v = X.Get(y, i);
		if (v > maxV)
			maxV = v;
	}

	float acc = 0.0f;
	for (i = 0; i < X.GetFlatWidth(); ++i)
	{
		float v = X.Get(y, i);
		acc += exp(v - maxV);
	}

	float v = X.Get(y, x);
	v = exp(v - maxV) / acc;
	O.Set(y, x, v);
}
