gemm.hpp

#ifndef GEMM_HPP
#define GEMM_HPP

#include <hmlp.h>
#include <hmlp_base.hpp>
//#include <View.hpp>

using namespace std;
using namespace hmlp;

namespace hmlp
{
namespace gemm
{

template<typename T>
class xgemmTask : public Task
{
  public:

    T alpha = 0.0;

    View<T> A;

    View<T> B;

    T beta = 0.0;

    View<T> C;

    void Set( 
        T alpha, View<T> &A,
                 View<T> &B,
        T beta,  View<T> &C )
    {
      this->alpha = alpha;
      this->A = A;
      this->B = B;
      this->beta = beta;
      this->C = C;

      ostringstream ss;
      name = string( "gemm" );

      double flops = 0.0, mops = 0.0;
      cost  = 2.0 * C.row() * C.col();
      event.Set( name + label, flops, mops );
    };

    void DependencyAnalysis()
    {
      A.DependencyAnalysis( R,  this );
      B.DependencyAnalysis( R,  this );
      C.DependencyAnalysis( RW, this );
      assert( !this->TryEnqueue() );
    };

    void Execute( Worker* user_worker )
    {
      string transA, transB;
      if ( A.IsTransposed() ) transA = "Transpose";
      else                    transA = "No transpose";
      if ( B.IsTransposed() ) transB = "Transpose";
      else                    transB = "No transpose";

      size_t m = C.row();
      size_t n = C.col();
      size_t k = A.col();

      assert( A.row() == m );
      assert( B.row() == k );
      assert( B.col() == n );

      //int rand_id = rand();

      //printf( "%d GEMM task %s %s %lu %lu %lu, %E, %E\n", 
      //    rand_id, transA.data(), transB.data(), m, n, k, alpha, beta ); fflush( stdout );
      //printf( "%d lda %lu ldb %lu ldc %lu\n", rand_id, A.ld(), B.ld(), C.ld() ); fflush( stdout );

      xgemm( transA.data(), transB.data(), m, n, k,
        alpha, A.data(), A.ld(),
               B.data(), B.ld(),
        beta,  C.data(), C.ld() );

      //printf( "%d end GEMM task %s %s %lu %lu %lu, %E, %E\n", 
      //    rand_id, transA.data(), transB.data(), m, n, k, alpha, beta ); fflush( stdout );
    };

}; 
template<typename T>
class xgemmBarrierTask : public Task
{
  public:

    View<T> C;

    void Set( 
        T alpha, View<T> &A,
                 View<T> &B,
        T beta,  View<T> &C )
    {
      this->C = C;
      this->stealable = false;

      name = string( "gemmBarrier" );

      double flops, mops;
      flops = 0.0;
      mops  = 0.0;
      cost  = 1.0;
      event.Set( name + label, flops, mops );
    };

    void DependencyAnalysis() { C.DependencyAnalysis( RW, this ); };

    void Execute( Worker* user_worker ) {};

}; 
template<typename T>
void CreatexgemmTask( T alpha, View<T> &A, View<T> &B, T beta,  View<T> &C )
{
  auto *task = new xgemmTask<T>();
  task->Set( alpha, A, B, beta, C );
  task->Submit();
  task->DependencyAnalysis();
}; 
template<size_t NB = 512, typename T>
void xgemm_var1( T alpha, View<T> &A, View<T> &B, T beta,  View<T> &C )
{
  View<T> AL, AR, 
          A0, A1, A2;
  View<T> BT, BB, 
          B0, B1, B2;
  A.Partition1x2( AL, AR, 0, LEFT );
  B.Partition2x1( BT,
                  BB,     0, TOP  ); 

  //printf( "AL.col() %lu AR.col() %lu A.col() %lu\n", AL.col(), AR.col(), A.col() );

  while ( AL.col() < A.col() )
  {
    //printf( "AL.col() %lu AR.col() %lu A.col() %lu\n", AL.col(), AR.col(), A.col() );
    size_t b = std::min( AR.col(), NB );

    Repartition1x2To1x3( AL,      AR,
                         A0,  A1, A2, b, RIGHT );
    Repartition2x1To3x1( BT,  B0,
                              B1,
                         BB,  B2, b, BOTTOM );

    CreatexgemmTask( alpha, A1, B1, beta, C );
    beta = 1.0;
    ContinueWith1x3To1x2( AL,      AR,
                          A0,  A1, A2, LEFT );
    ContinueWith3x1To2x1( BT,  B0,
                               B1,
                          BB,  B2,  TOP );

  } 
}; 
template<size_t NB = 512, typename T>
void xgemm_var2( T alpha, View<T> &A, View<T> &B, T beta,  View<T> &C )
{
  View<T> CL, CR, 
          C0, C1, C2;
  View<T> BL, BR, 
          B0, B1, B2;
  
  C.Partition1x2( CL, CR, 0, LEFT );
  B.Partition1x2( BL, BR, 0, LEFT );

  while ( BL.col() < B.col() )
  {
    size_t b = std::min( BR.col(), NB );

    Repartition1x2To1x3( CL,      CR,
                         C0,  C1, C2, b, RIGHT );
    Repartition1x2To1x3( BL,      BR,
                         B0,  B1, B2, b, RIGHT );

    xgemm_var1( alpha, A, B1, beta, C1 );
    ContinueWith1x3To1x2( CL,      CR,
                          C0,  C1, C2, LEFT );
    ContinueWith1x3To1x2( BL,      BR,
                          B0,  B1, B2, LEFT );

  } 
}; 
template<size_t NB = 512, typename T>
void xgemm_var3( T alpha, View<T> &A, View<T> &B, T beta,  View<T> &C )
{
  View<T> AT, A0, CT, C0, 
          AB, A1, CB, C1,
              A2,     C2;

  A.Partition2x1( AT,
                  AB,     0, TOP  ); 
  C.Partition2x1( CT,
                  CB,     0, TOP  ); 

  while ( AT.row() < A.row() )
  {
    size_t b = std::min( AB.row(), NB );

    Repartition2x1To3x1( AT,  A0,
                              A1,
                         AB,  A2, b, BOTTOM );
    Repartition2x1To3x1( CT,  C0,
                              C1,
                         CB,  C2, b, BOTTOM );

    xgemm_var2( alpha, A1, B, beta, C1 );
    ContinueWith3x1To2x1( AT,  A0,
                               A1,
                          AB,  A2,  TOP );
    ContinueWith3x1To2x1( CT,  C0,
                               C1,
                          CB,  C2,  TOP );
  }; 
}; 
template<size_t NB = 512, typename T>
void xgemm( T alpha, View<T> &A, View<T> &B, T beta,  View<T> &C )
{
  //string transA, transB;
  //if ( A.IsTransposed() ) transA = "Transpose";
  //else                    transA = "No transpose";
  //if ( B.IsTransposed() ) transB = "Transpose";
  //else                    transB = "No transpose";
  //size_t m = C.row();
  //size_t n = C.col();
  //size_t k = A.col();
  //assert( A.row() == m );
  //assert( B.row() == k );
  //assert( B.col() == n );
  //hmlp::xgemm( transA.data(), transB.data(), m, n, k,
  //  alpha, A.data(), A.ld(),
  //         B.data(), B.ld(),
  //  beta,  C.data(), C.ld() );
  //return;





  A.CreateLeafMatrixBlocks( NB, NB );
  B.CreateLeafMatrixBlocks( NB, NB );
  C.CreateLeafMatrixBlocks( NB, NB );

  if ( hmlp_is_in_epoch_session() )
  {
    auto *begXGEMMtask = new xgemmBarrierTask<T>();
    auto *endXGEMMtask = new xgemmBarrierTask<T>();
    begXGEMMtask->Set( alpha, A, B, beta, C );
    begXGEMMtask->Submit();
    begXGEMMtask->DependencyAnalysis();

    xgemm_var3( alpha, A, B, beta, C );

    endXGEMMtask->Set( alpha, A, B, beta, C );
    endXGEMMtask->Submit();
    endXGEMMtask->DependencyAnalysis();

    begXGEMMtask->TryEnqueue();
    endXGEMMtask->CallBackWhileWaiting();
  }
  else
  {
    xgemm_var3( alpha, A, B, beta, C );
  }
}; 
template<typename T>
void xgemm( hmlpOperation_t transA, hmlpOperation_t transB,
    T alpha, Data<T> &A, Data<T> &B, T beta,  Data<T> &C ) 
{
  const bool TRANS = true;
  const bool NOTRANS = true;

  View<T> Aview, Bview, Cview;

  if ( transA == HMLP_OP_T ) Aview.Set(  true, A );
  else                       Aview.Set( false, A );
  if ( transB == HMLP_OP_T ) Bview.Set(  true, B );
  else                       Bview.Set( false, B );
  Cview.Set( C );

  xgemm( alpha, Aview, Bview, beta, Cview );
  
}; 
template<typename T>
void xgemm( T alpha, Data<T> &A, Data<T> &B, T beta,  Data<T> &C ) 
{
  xgemm( HMLP_OP_N, HMLP_OP_N, alpha, A, B, beta, C );
}; 
}; 
}; 
#endif