tc.sac

/*
name:	 tc.sac
author:	 Kevin Klein
compile: sac2c -o tc tc.sac
run:	 ./tc
descr.:	 SaC implementation of Tensor Comprehension kernels
*/

use Array: all;
use StdIO: all;
use CommandLine: {argc, argv};
use String: {string};


//---------------------------------------------------------
// 1) MATRIX VECTOR PRODUCT
//---------------------------------------------------------
/*
// Matrix Vector Product in TC
def mv(float(M,K) A, float(K) x) → (C) {
	C(i) = 0
	C(i) += A(i,k) * x(k)
}
// or as a one liner...
def mv(float(M,K) A, float(K) x) → (C) {
	C(i) +=! A(i,k) * x(k)
}
*/
// Matrix Vector Product in SaC
//
noinline
float[.] mv(float[.,.] A, float[.] x)
{
   return{[i] -> sum({[k] -> A[[i,k]] * x[[k]]})};
}

//---------------------------------------------------------
// 2) SGEEM BLAS
//---------------------------------------------------------
/*
// SGEEM BLAS in TC
def sgemm(float a, float b,
float(N,M) A, float(M,K) B) → (C) {
	# initialization
	C(i,j) = b * C(i,j)
	# accumulation
	C(i,j) += a * A(i,k) * B(k,j)
}
*/
// sgemm BLAS in SaC
noinline
float[.,.] sgemm(float a, float b, float[.,.] A, float[.,.] B, float[.,.] C)
{
	// initialisation
  C = b * C;
	// accumulation
	C += {[i,j] -> sum({[k] -> a * (A[[i,k]] * B[[k,j]])})};

	return(C);
}

//-------------------------------------------------------------
// 3) FULLY CONNECTED LAYER FOLLOWED BY A RECTIFIED LINEAR UNIT
//-------------------------------------------------------------
/*
// CFRELU in TC
def fcrelu(float(B,I) in, float(O,I) weight, float(I) bias) → (out) {
	out(i,j) = bias(j)
	out(b,o) += in(b,i) * weight(o,i)
	out(i,j) = fmaxf(out(i,j), 0)
}
*/
// CFRELU in SaC
noinline
float[.,.] fcrelu(float[.,.] in, float[.,.] weight, float[.] bias)
{
	B = shape(in)[[0]];
	out = genarray([B], bias);
	out += {[b,o] -> sum({[i] -> in[[b,i]] * weight[[o,i]]})};
	out = max(out, 0f);
	return(out);
}

//---------------------------------------------------------
// 4) 2-D CONVOLUTION
//---------------------------------------------------------
/*
// 2-D Convolution in TC
def conv2d(float(B,IP,H,W) in, float(OP,IP,KH,KW) weight) → (out) {
	out(b,op,h,w) +=! in(b,ip,h+kh,w+kw) * weight(op,ip,kh,kw)
}
*/
// 2-D Convolution in SaC
noinline
float[.,.,.,.] conv2d(float[.,.,.,.] in, float[.,.,.,.] weight)
{
  lb = 0*shape(in);
  ub = shape(in) - [0,0,shape(weight)[[2]], shape(weight)[[3]]];
  return with{
    (lb <= [b,op,h,w] < ub) : sum({[ip,kh,kw] -> in[[b,ip,h+kh,w+kw]] * weight[[op,ip,kh,kw]]});
  } : modarray(in);
}

//---------------------------------------------------------
// 5) MAX POOLING LAYER
//---------------------------------------------------------
/*
// Max Pooling Layer in TC
def maxpool2x2(float(B,C,H,W) in) → (out) {
	out(b,c,i,j) max=! in(b,c, 2 * i + kw, 2 * j + kh)
	where kw in 0:2, kh in 0:2
}
*/
// maxpool2x2 in SaC
noinline
float[.,.,.,.] maxpool2x2(float[.,.,.,.] in)
{
  lb = 0*shape(in);
  ub = [shape(in)[0], shape(in)[1], (shape(in)[2]-3)/2, (shape(in)[3]-3)/2];
  return with {
    (lb <= [b,c,i,j] < ub) : maxval(with{
      ([0,0] <= [kw,kh] <= [2,2]) : in[[b,c,2*i+kw,2*j+kh]];
    } : genarray([3,3], 0f));
  } : modarray(in);
}

//---------------------------------------------------------
// 6) GATHER OPERATION
//---------------------------------------------------------
/*
// gather operation in TC
def gather(float(N) X, int(A,B) I) → (Z) {
	Z(i,j) = X(I(i,j))
}
*/
// gather operation in SaC
// this is a  comnination of iterators and subscript expressions, how do I do this in SaC?
noinline
float[.,.] gather(float[.] X, int[.,.] I)
{
  return {[i,j] -> X[I[[i,j]]]};
}
//---------------------------------------------------------
// 7) STRIDED CONVOLUTION
//---------------------------------------------------------
/*
// strided convolution in TC
def sconv2d(int sh, int sw, float(N,C,H,W) I, float(F,C,KH,KW) W, float(F) B) → (O) {
	O(n,f,h,w) = B(f)
	O(n,f,h,w) += I(n,c, sh * h + kh, sw * w + kw) * W(f,c,kh,kw)
}
*/
// strided convolution in SaC
noinline
float[.,.,.,.] sconv2d(int sh, int sw, float[.,.,.,.] I, float[.,.,.,.] W, float[.] B)
{
	f = shape(I)[[0]];
	O = genarray([f], B);
	lb = 0*shape(I);
	ub = shape(I) - [0,0,shape(W)[[2]], shape(W)[[3]]];
	return with{
    (lb <= [n,f,h,w] < ub) : sum({[c,kh,kw] -> I[[n,c,sh*h+kh,sw*w+kw]] * W[[f,c,kh,kw]]});
  } : modarray(I);
}

//---------------------------------------------------------
// 8) 1-DIMENSIONAL CONVOLUTION
//---------------------------------------------------------
/*
// 1-D Convolution in TCshape(K)[[0]]
def conv1d(float(M) I, float(N) K) → (O) {
	O(i) = K(x) * I(i + x)
}
*/
// 1-D Convolution in SaC
noinline
float[.] conv1d(float[.] I, float[.] K)
{
  lb = 0*shape(I);
  ub = shape(I) - shape(K);
  return with{
    (lb <= [i] < ub) : sum({[x] -> K[[x]] * I[[i + x]]});
  } : modarray(I); // don't understand this, whether you choose K or I affects the result...
}

//---------------------------------------------------------
// 9) OUTER PRODUCT MATRIX MULTIPLICATION
//---------------------------------------------------------
/*
// outer product matrix multiply in TC
def outerProductMM(float(P,Q,R) A, float(S,R,T) B) → (O) {
	O(p,s,q,t) +=! A(p,q,r) * B(s,r,t)
}
*/
// outer product matrix multiply in SaC
noinline
float[.,.,.,.] outerProductMM(float[.,.,.] A, float[.,.,.] B)
{
	O = {[p,s,q,t] -> sum(A[[p,q,.]] * B[[s,.,t]])};
	return(O);
}

//---------------------------------------------------------
// 10) TRANSPOSED MATRIX MULTIPLICATION
//---------------------------------------------------------
/*
// transposed matrix multiply in TC
def tmm(float(M,K) A, float(N,K) B) → (C) {
	C(m,n) +=! A(m,kk) * B(n,kk)
}
*/
// transposed matrix multiply in SaC
noinline
float[.,.] tmm(float[.,.] A, float[.,.] B)
{
	C = {[m,n] -> sum(A[[m,.]] * B[[n,.]])};
	return(C);
}

//---------------------------------------------------------
// 11) TRANSPOSED BATCH MATRIX MULTIPLICATION
//---------------------------------------------------------
/*
// transposed batch matrix multiply in TC
def tbmm(float(B,N,M) X, float(B,K,M) Y) → (Z) {
	Z(b,n,k) +=! X(b,n,m) * Y(b,k,m)
}
*/
// transposed batch matrix multiply in SaC
noinline
float[.,.,.] tbmm(float[.,.,.] X, float[.,.,.] Y)
{
	Z = {[b,n,k] -> sum(X[[b,n,.]] * Y[[b,k,.]])};
	return(Z);
}

//---------------------------------------------------------
// 12) GROUPED CONVOLUTIONS
//---------------------------------------------------------
/*
// grouped convolutions in TC
def gconv(float(N,G,C,H,W) I,float(G,F,C,KH,KW) W1, float(M) B) → (O) {
	O(n,g,o,h,w) +=! I(n,g,i, h + kh, w + kw) * W1(g,o,i,kh,kw)
	O(n,g,o,h,w) = O(n,g,o,h,w) + B(m)
}
*/
// grouped convolutions in SaC
noinline
float[.,.,.,.,.] gconv(float[.,.,.,.,.] I, float[.,.,.,.,.] W1, float[.] B)
{
	lb = 0*shape(I);
	ub = shape(I) - [0,0,0,shape(W1)[3], shape(W1)[4]];

	O = with{
		(lb <= [n,g,o,h,w] < ub) : sum({[i,kh,kw] -> I[n,g,i,h+kh,w+kw] * W1[g,o,i,kh,kw]});
	} : modarray(I);
  O = {[n,g,o,h,w] -> sum({[m] -> O[n,g,o,h,w] * B[[m]]})};
  return(O);
}

//---------------------------------------------------------
// 13) 2 PARALLEL LOOKUP TABLES
//---------------------------------------------------------
/*
// parallel lookup tables in TC
def 2LUT(float(E1,D) LUT1, int(B,L1) I1, float(E2,D) LUT2, int(B,L2) I2) → (O1, O2) {
	O1(i,j) +=! LUT1(I1(i,k),j)
	O2(i,j) +=! LUT2(I2(i,k),j)
}
*/
// parallel lookup tables in SaC
noinline
float[.,.], float[.,.] LUT2(float[.,.]LUT1, int[.,.] I1, float[.,.] LUT2, int[.,.] I2)
{
  // No shape information found for index scalar 'j'.
  O1 = {[i,j] -> sum({[k] -> LUT1[ I1[i, k], j] })};
  O2 = {[i,j] -> sum({[k] -> LUT2[ I2[i, k], j] })};
  return(O1, O2);
}


//---------------------------------------------------------
// MAIN FUNCTION
//---------------------------------------------------------
int main()
{
#define M 20
#define N 22
#define K 12
#define B 20
#define IP 22
#define H 24
#define W 26
#define OP 28
#define KH 3
#define KW 3
#define G 25
#define C 20
#define F 23

/*
	// 1) call MATRIX VECTOR PRODUCT
  A = genarray([M,K], 2f);
  x = genarray([K], 3f);
	print(mv(A, x));
*/
//---------------------------------------------------------
	// 2) call SGEEM BLAS
/*
  A = genarray([N,M], 2f);
  B = genarray([M,K], 3f);
  C = genarray([N,K], 4f);
	print(sgemm(2f, 3f, A, B, C));
*/
//---------------------------------------------------------
/*
	// 3) call FCRELU
  A = genarray([N,M], 2f);
  B = genarray([M,K], 3f);
	print(fcrelu(B, C, x));
*/
//---------------------------------------------------------
/*
	// 4) call 2-D CONVOLUTION
  in = genarray([B,IP,H,W], 2f);
  weight = genarray([OP,IP,KH,KW], 3f);
	print(conv2d(in,weight));
*/
//---------------------------------------------------------
/*
	// 5) call Max Pooling
  in = genarray([B,IP,H,W], 2f);
  print(maxpool2x2(in));
*/
//---------------------------------------------------------
/*
	// 6) call GATHER OPERATION
  I = reshape([M, N], iota(M * N));
  v = genarray([M * N], 2f);
	print(gather(v, I));
*/
//---------------------------------------------------------
/*
	// 7) call STRIDED CONVOLUTION
	I = genarray([B,IP,H,W], 2f);
	V = genarray([B,IP,H,W], 2f);
	print(sconv2d(3,4,I,V,[3f]));
*/
//---------------------------------------------------------
/*
	// 8) call 1D CONVOLUTION
	print(conv1d([2f,3f],[4f,5f]));
*/
//---------------------------------------------------------
/*
	// 9) call OUTER PRODUCT MATRIX MULTIPLY
	cube1 = [[[1f,2f,3f]],[[3f,4f,5f]]];
	call9 = outerProductMM(cube1, cube1);
	print(call9);
*/
//---------------------------------------------------------
/*
	// 10) call TRANSPOSED MATRIX MULTIPLY
	mat1 = [[1f,2f],[3f,4f]];
	//call10 = tmm(mat1, mat1);
	//print(call10);
*/
//---------------------------------------------------------
/*
	// 11) call TRANSPOSED BATCH MATRIX MULTIPLICATION
	cube1 = [[[1f,2f,3f]],[[3f,4f,5f]]];
	call11 = tbmm(cube1, cube1);
	print(call11);
*/
//---------------------------------------------------------
/*
	// 12) call gconv
	I = genarray([N,G,C,H,W], 2f);
	W1 = genarray([G,F,C,KH,KW], 3f);
	print(gconv(I, W1,[3f]));
*/
//---------------------------------------------------------
/*
  // 13) call 2LUT
  mat1 = [[1f,2f],[3f,4f]];
  mat2 = [[1,2],[3,4]];
  mat3 = [[5f,6f],[7f,8f]];
  mat4 = [[5,6],[7,8]];
  res1, res2 = LUT2(mat1,mat2,mat3,mat4);
  print(res1);
  print(res2);
*/

	return 0;
}