rank_k.hpp

#ifndef RANK_K_HPP
#define RANK_K_HPP

#include <assert.h>
#include <typeinfo>
#include <algorithm>

#include <hmlp.h>
#include <hmlp_internal.hpp>
#include <hmlp_base.hpp>

#include <packing.hpp>
#include <semiring_mrxnr.hpp>
#include <fused_mrxnr.hpp>

using namespace std;
using namespace hmlp;

namespace hmlp
{
template<int KC, typename SEMIRINGKERNEL, typename TA, typename TB, typename TV>
void rank_k_macro_kernel
(
  tci::Comm &Comm3rd,
  int ic, int jc, int pc,
  int  m, int  n, int  k,
  TA *packA,
  TB *packB,
  TV *V, int rs_v, int cs_v,
  SEMIRINGKERNEL semiringkernel
)
{
  const static int MR         = SEMIRINGKERNEL::mr;
  const static int NR         = SEMIRINGKERNEL::nr;
  const static int PACK_MR    = SEMIRINGKERNEL::pack_mr;
  const static int PACK_NR    = SEMIRINGKERNEL::pack_nr;
  auto Comm2nd = Comm3rd.Split( hmlp_read_nway_from_env( "KS_JR_NT" ) );
  auto Loop3rd = Comm3rd.DistributeOver1DGangs(        0, n,      NR );
  auto Pack3rd = Comm3rd.DistributeOver1DGangs(        0, n, PACK_NR );
  auto Loop2nd = Comm2nd.DistributeOver1DThreads(      0, m,      MR );
  auto Pack2nd = Comm2nd.DistributeOver1DThreads(      0, m, PACK_MR );
  for ( int j  = Loop3rd.beg(), jp  = Pack3rd.beg();
            j  < Loop3rd.end();
            j += Loop3rd.inc(), jp += Pack3rd.inc() )
  {
    struct aux_s<TA, TB, TV, TV> aux;
    aux.pc       = pc;
    aux.b_next   = packB;
    aux.do_packC = 0;
    aux.jb       = std::min( n - j, NR );
    for ( int i  = Loop2nd.beg(), ip  = Pack2nd.beg();
              i  < Loop2nd.end();
              i += Loop2nd.inc(), ip += Pack2nd.inc() )
    {
      aux.ib = std::min( m - i, MR );
      if ( i + MR >= m ) aux.b_next += Pack3rd.inc() * k;
           
      if ( aux.jb == NR && aux.ib == MR )                 
      {
        semiringkernel( k, &packA[ ip * k ], &packB[ jp * k ],
          &V[ i * rs_v + j * cs_v ], rs_v, cs_v, &aux );
      }
      else
      {
        TV vtmp[ MR * NR ];

        if ( pc ) // initilize ctmp
        {
          for ( auto jj = 0; jj < aux.jb; jj ++ )
            for ( auto ii = 0; ii < aux.ib; ii ++ )
              vtmp[ jj * MR + ii ] = 
                V[ ( j + jj ) * cs_v + ( i + ii ) * rs_v ];
        }

        semiringkernel( k, &packA[ ip * k ], &packB[ jp * k ],
          vtmp, 1, MR, &aux );

        for ( auto jj = 0; jj < aux.jb; jj ++ )
          for ( auto ii = 0; ii < aux.ib; ii ++ )
            V[ ( j + jj ) * cs_v + ( i + ii ) * rs_v ] 
              = vtmp[ jj * MR + ii ];
      }
    } 
  } 
}; 
template<
  int MC, int NC, int KC,
  typename TPACKA, typename TPACKB, typename TV,
  typename     TA, typename     TB, typename TC,
  typename SEMIRINGKERNEL>
void rank_k_internal
(
  tci::Comm &Comm6th,
  int batchId, int m, int n, int k, int k_stra,
  TA& A, 
  TB& B, 
  TV* V, int rs_v, int cs_v,
  SEMIRINGKERNEL semiringkernel
)
{
  const static int MR         = SEMIRINGKERNEL::mr;
  const static int NR         = SEMIRINGKERNEL::nr;
  const static int PACK_MR    = SEMIRINGKERNEL::pack_mr;
  const static int PACK_NR    = SEMIRINGKERNEL::pack_nr;
  const static int ALIGN_SIZE = SEMIRINGKERNEL::align_size;
  const static int PACK_MC    = ( MC / MR ) * PACK_MR;
  const static int PACK_NC    = ( NC / NR ) * PACK_NR;
  auto Comm5th = Comm6th.Split( hmlp_read_nway_from_env( "KS_JC_NT" ) );
  auto Comm4th = Comm5th.Split( 1 );
  auto Comm3th = Comm4th.Split( hmlp_read_nway_from_env( "KS_IC_NT" ) );
  int nc = Comm6th.BalanceOver1DGangs( n, NC, NR );
  int pack_nc = ( nc / NR ) * PACK_NR;
  auto *packB = Comm4th.AllocateSharedMemory<ALIGN_SIZE, TPACKB>( KC * ( pack_nc + 1 ) );
  auto *packA = Comm3th.AllocateSharedMemory<ALIGN_SIZE, TPACKA>( KC * ( PACK_MC + 1 ) );
  auto Loop6th = Comm6th.DistributeOver1DGangs(      0, n, nc );
  auto Loop5th = Comm5th.DistributeOver1DGangs( k_stra, k, KC );
  auto Loop4th = Comm4th.DistributeOver1DGangs(      0, m, MC );
  for ( int jc  = Loop6th.beg(); 
            jc  < Loop6th.end(); 
            jc += Loop6th.inc() )
  {
    auto jb = std::min( n - jc, nc );
    for ( int pc  = Loop5th.beg(); 
              pc  < Loop5th.end(); 
              pc += Loop5th.inc() )
    {
      auto pb = std::min( k - pc, KC );
      auto LooppkB = Comm4th.DistributeOver1DThreads( 0, jb,      NR );
      auto PackpkB = Comm4th.DistributeOver1DThreads( 0, jb, PACK_NR );
      for ( int j  = LooppkB.beg(), jp  = PackpkB.beg();
                j  < LooppkB.end();
                j += LooppkB.inc(), jp += PackpkB.inc() )
      {
        B.Pack( k, pc, pb, n, jc + j, std::min( jb - j, NR ), 
            &packB[ jp * pb ] );
      }
      Comm4th.Barrier();
      for ( int ic  = Loop4th.beg();
                ic  < Loop4th.end();
                ic += Loop4th.inc() )
      {
        auto ib = std::min( m - ic, MC );
        auto LooppkA = Comm3th.DistributeOver1DThreads( 0, ib, MR );
        auto PackpkA = Comm3th.DistributeOver1DThreads( 0, ib, PACK_MR );
        for ( int i  = LooppkA.beg(), ip  = PackpkA.beg();
                  i  < LooppkA.end();
                  i += LooppkA.inc(), ip += PackpkA.inc() )
        {
          A.Pack( m, ic + i, std::min( ib - i, MR ), 
              k, pc, pb, &packA[ ip * pb ] );
        }
        Comm3th.Barrier();
        rank_k_macro_kernel<KC>( Comm3th, 
          ic, jc, pc, ib, jb, pb, packA, packB,
          V + ic * rs_v + jc * cs_v, rs_v, cs_v,
          semiringkernel );
        Comm3th.Barrier();
      } 
      Comm4th.Barrier();
    } 
    Comm5th.Barrier();
  } 
  Comm6th.Barrier();
  Comm3th.FreeSharedMemory( packA );
  Comm4th.FreeSharedMemory( packB );
}; 
template<
  int MC, int NC, int KC,
  typename TPACKA, typename TPACKB, typename TV,
  typename     TA, typename     TB, typename TC,
  typename SEMIRINGKERNEL>
void rank_k
(
  int batchId, int m, int n, int k,
  TA& A, 
  TB& B, 
  TC& C,
  SEMIRINGKERNEL semiringkernel
)
{
  const static int MR         = SEMIRINGKERNEL::mr;
  const static int NR         = SEMIRINGKERNEL::nr;
  const static int PACK_MR    = SEMIRINGKERNEL::pack_mr;
  const static int PACK_NR    = SEMIRINGKERNEL::pack_nr;
  const static int ALIGN_SIZE = SEMIRINGKERNEL::align_size;
  const static int PACK_MC    = ( MC / MR ) * PACK_MR;
  const static int PACK_NC    = ( NC / NR ) * PACK_NR;
  const static bool USE_STRASSEN = false; 

  if ( m == 0 || n == 0 || k == 0 ) return;
  if ( !is_same<TC, MatrixLike<PACK_MR, TV, TV>>::value )
  {
    exit( 1 );
  }
  auto *V = reinterpret_cast<TV*>( C.X );
  auto rs_v = C.rs;
  auto cs_v = C.cs;


  int k_stra = 0; 
  if ( USE_STRASSEN )
  {
    assert( typeid(TPACKA) == typeid(TPACKB) );
    assert( typeid(TC) == typeid(TV) );
    k_stra = k - k % KC;

    if ( k_stra == k ) k_stra -= KC;
  }

  tci::Parallelize( NULL, rank_k_internal<MC, NC, KC, TPACKA, TPACKB, TV, 
      TA, TB, TC, SEMIRINGKERNEL>, 
      batchId, m, n, k, k_stra, A, B, C, V, rs_v, cs_v, 
      semiringkernel );
}; 
}; 
#endif