#include <amp.h>
using namespace concurrency;

#define TS 16

int do_it(tiled_index<TS,TS> t_idx,array_view<int,2> a,array_view<int,2> b) restrict(direct3d)
{
  int row = t_idx.local[0]; 
  int col = t_idx.local[1];
  int sum = 0;
  for (int i = 0; i < a.extent[1]; i += TS) {
    tile_static int locA[TS][TS], locB[TS][TS];
    locA[row][col] = a(t_idx.global[0], col + i);
    locB[row][col] = b(row + i, t_idx.global[1]);
    t_idx.barrier.wait();

    for (int k = 0; k < TS; k++)
      sum += locA[row][k] * locB[k][col];            
    t_idx.barrier.wait();
  }
  return sum;
}

void MatMul_AMP(std::vector<int>& vC, std::vector<int>& vA, std::vector<int>& vB, int M, int N, int W )
{
  array_view<int,2> a(M, W, vA), b(W, N, vB), c(M,N,vC);   

  parallel_for_each(
    c.grid.tile<TS, TS>(), 
    [=] (tiled_index<TS, TS> t_idx) restrict(direct3d)  
    {
      int sum = do_it(t_idx, a, b);
      index<2> glob_idx = t_idx.global;
      c[glob_idx] = sum;
    });
}

void main()
{
  const int M = 32;
  const int N = 32;
  const int W = 64;

  // 3 vectors
  std::vector<int> A(M*W);
  std::vector<int> B(W*N);
  std::vector<int> C(M*N);

  // populate vector A
  for (unsigned int i = 0; i < A.size(); i++)
    A[i] = i;

  // populate vector B
  for (unsigned int i = 0; i < B.size(); i++)
    B[i] = i;

  // do the matrix multiplication using C++ AMP
  MatMul_AMP(C, A, B, M, N, W);

  //bye
  printf("all good, hit enter, ");
  getchar();
}