#include "atlas_misc.h"
#define ATL_LAUNCHORDER         /* we want static ATL_launchorder array */
#include "atlas_threads.h"
#include "atlas_tlvl3.h"
/*
 * =========================================================================
 * This file contains support routines for TGEMM that are not type-dependent
 * =========================================================================
 */
static int ATL_thrdecompMM_rec
   (ATL_TMMNODE_t *ptmms, const enum ATLAS_TRANS TA, const enum ATLAS_TRANS TB,
    ATL_CINT Mblks, const int mr, ATL_CINT Nblks, const int nr, ATL_CINT Kblks,
    const int kr, const void *A, ATL_INT lda, const void *B, const ATL_INT ldb,
    const void *C, ATL_CINT ldc, const int P, const int indx, const int COPYC)
{
   int pR, pL, rL, rR, np, eltsh;
   ATL_INT nblksL, nblksR, i, j;
   double d;

#ifdef DEBUG
   ATL_assert(P > 0);
#endif
/*
 * End recursion if we are down to 1 processor, or if we are out of blocks
 */
   if ( P <= 1 || (Mblks <= 1 && Nblks <= 1 && Kblks <= 2) )
   {
      ptmms[indx].A = A;
      ptmms[indx].B = B;
      ptmms[indx].C = (void*)C;
      ptmms[indx].lda = lda;
      ptmms[indx].ldb = ldb;
      ptmms[indx].ldc = ldc;
      ptmms[indx].M = ptmms[indx].mb*Mblks + mr;
      ptmms[indx].N = ptmms[indx].nb*Nblks + nr;
      ptmms[indx].K = ptmms[indx].kb*Kblks + kr;
      if (COPYC)
      {
         ptmms[indx].nCw = 1;
         ptmms[indx].nCp = ptmms[indx].ownC = 0;
         ptmms[indx].Cinfp[0] = ptmms+indx;
/*
 *       Make ldcw a multiple of 4 that is not a power of 2
 */
         i = ((ptmms[indx].M + 3)>>2)<<2;
         if (!(i & (i-1)))
            i += 4;
         ptmms[indx].ldcw = i;
      }
      else
      {
         ptmms[indx].ldcw = ptmms[indx].nCw = 0;
         ptmms[indx].nCp = ptmms[indx].ownC = 1;
         ptmms[indx].Cinfp[ATL_NTHREADS-1] = ptmms+indx;
      }
      ptmms[indx].Cw = NULL;
#ifdef DEBUG
fprintf(stderr, "%d: M=%d, N=%d, K=%d, ownC=%d, nCp=%d, nCw=%d\n",
        indx, ptmms[indx].M, ptmms[indx].N, ptmms[indx].K,
        ptmms[indx].ownC, ptmms[indx].nCp, ptmms[indx].nCw);
#endif
      return(1);
   }

   eltsh = ptmms[indx].eltsh;
   pR = P>>1;    /* on right, take P/2 threads */
   pL = P - pR;  /* on left, take remaining threads */
   d = (pR == pL) ? 0.5 : ((double)pL)/((double)P);    /* percent on left */

/*
 * Only cut K if it dominates M & N (here we say K must be 4 time larger)
 * and M&N are small enough that we can afford to malloc C
 * (here we say C workspace must be 16MB or less)
 */
   if ( ( ((Mblks < 2) && Nblks < 2) ||
          (((Kblks>>2) > Mblks) && ((Kblks>>2) > Nblks)) )
       && (Mblks*((Nblks*ptmms[indx].mb*(ptmms[indx].nb<<ptmms[indx].eltsh)
           +1023)>>10) <= 16*1024))
   {
#ifdef DEBUG
   fprintf(stderr, "Cut K\n");
#endif
      nblksL = (d * Kblks);
      nblksR = Kblks - nblksL;
      if (nblksR < nblksL)
      {
         rL = 0;
         rR = kr;
      }
      else
      {
         rL = kr;
         rR = 0;
      }
      i = (nblksL*ptmms[indx].kb + rL)<<eltsh;
      np = ATL_thrdecompMM_rec(ptmms, TA, TB, Mblks, mr, Nblks, nr, nblksL, rL,
                               A, lda, B, ldb, C, ldc, pL, indx, COPYC);
      np += ATL_thrdecompMM_rec(ptmms, TA, TB, Mblks, mr, Nblks, nr, nblksR, rR,
                                (TA==AtlasNoTrans)?MindxT(A,lda*i):MindxT(A,i),
                                lda,
                               (TB == AtlasNoTrans)?MindxT(B,i):MindxT(B,i*ldb),
                               ldb, C, ldc, pR, indx+pL, 1);
      return(np);
   }
   else if (Mblks >= Nblks)  /* split M */
   {
#ifdef DEBUG
   fprintf(stderr, "Cut M\n");
#endif
      nblksL = (d * Mblks);
      nblksR = Mblks - nblksL;
      if (nblksR < nblksL)
      {
         rL = 0;
         rR = mr;
      }
      else
      {
         rL = mr;
         rR = 0;
      }
      i = (nblksL*ptmms[indx].mb+rL) << eltsh;
      np = ATL_thrdecompMM_rec(ptmms, TA, TB, nblksL, rL, Nblks, nr, Kblks, kr,
                               A, lda, B, ldb, C, ldc, pL, indx, COPYC);
      np += ATL_thrdecompMM_rec(ptmms, TA, TB, nblksR, rR, Nblks, nr, Kblks, kr,
                                (TA==AtlasNoTrans)?MindxT(A,i):MindxT(A,i*lda),
                                lda, B, ldb, MindxT(C,i), ldc, pR, indx+pL,
                                COPYC);
      return(np);
   }
   else /* split N */
   {
#ifdef DEBUG
   fprintf(stderr, "Cut N\n");
#endif
      nblksL = (d * Nblks);
      nblksR = Nblks - nblksL;
      if (nblksR < nblksL)
      {
         rL = 0;
         rR = nr;
      }
      else
      {
         rL = nr;
         rR = 0;
      }
      i = (nblksL*ptmms[indx].nb+rL) << eltsh;
      np = ATL_thrdecompMM_rec(ptmms, TA, TB, Mblks, mr, nblksL, rL, Kblks, kr,
                               A, lda, B, ldb, C, ldc, pL, indx, COPYC);
      np += ATL_thrdecompMM_rec(ptmms, TA, TB, Mblks, mr, nblksR, rR, Kblks, kr,
                                A,  lda, (TB == AtlasNoTrans) ?
                                MindxT(B,i*ldb) : MindxT(B,i), ldb,
                                MindxT(C,i*ldc), ldc, pR, indx+pL, COPYC);
      return(np);
   }
}

static int ATL_thrdecompMM_M
   (ATL_TMMNODE_t *ptmms, const enum ATLAS_TRANS TA, const enum ATLAS_TRANS TB,
    ATL_CINT Mblks, const int mr, ATL_CINT Nblks, const int nr, ATL_CINT Kblks,
    const int kr, const void *A, ATL_INT lda, const void *B, const ATL_INT ldb,
    const void *C, ATL_CINT ldc, const int P, const int indx, const int COPYC)
{
   int j, i, m;
   const char *a=A, *c=C;
   const int eltsh = ptmms[0].eltsh, mb = ptmms[0].mb, n = ptmms[0].nb*Nblks+nr,
             k = ptmms[0].kb*Kblks+kr, minblks = Mblks / ATL_NTHREADS,
             extrablks = Mblks - minblks*ATL_NTHREADS;

   for (i=0; i < ATL_NTHREADS; i++)
   {
      m = minblks * mb;
      if (i < extrablks)
         m = (minblks + 1)*mb;
      else if (i == extrablks)
         m = minblks*mb + mr;
      else
         m = minblks*mb;

      j = ATL_launchorder[i];  /* use log2-launch order */
      ptmms[j].A = a;
      ptmms[j].B = B;
      ptmms[j].C = (void*)c;
      ptmms[j].lda = lda;
      ptmms[j].ldb = ldb;
      ptmms[j].ldc = ldc;
      ptmms[j].M = m;
      ptmms[j].N = n;
      ptmms[j].K = (m) ? k : 0;
      ptmms[j].ownC = 1;
      ptmms[j].nCp = ptmms[j].nCw = 0;
      ptmms[j].Cw = NULL;
      ptmms[j].ldcw = 0;
      m <<= eltsh;
      a = (TA == AtlasNoTrans) ? MindxT(a,m) : MindxT(a,m*lda);
      c = MindxT(c,m);
   }
   return(i);
}

#if ATL_NTHREADS != (1<<ATL_NTHRPOW2)
void EnforceNonPwr2LO(ATL_TMMNODE_t *ptmms, const int P)
/*
 * If threads aren't a power of 2, then the recursive decomposition will
 * fill in the ptmms array in a different order than the log2 spawn.
 * As long as P >= NTHREADS, all entries are filled in, so the worst that
 * happens is that thread 3 does the work that we expect to be done by 4.
 * However, if P < NTHREADS, then threads that launch starts won't have
 * work, so we must make sure that the 1st P elts in launchorder have
 * work to do.
 */
{
   int i, j, k, kk, h;

   if (P >= ATL_NTHREADS)
      return;
   for (i=0; i < P; i++)
   {
      j = ATL_launchorder[i];
      if (!ptmms[j].K)  /* we have found an empty required entry! */
      {
/*
 *      Search for a filled in entry that will not be used in launch
 */
        for (kk=ATL_NTHREADS-1; kk >= P; kk--)
        {
           k = ATL_launchorder[kk];
           if (ptmms[k].K)  /* found a non-empty entry */
           {
              ptmms[j].A = ptmms[k].A;
              ptmms[j].B = ptmms[k].B;
              ptmms[j].C = ptmms[k].C;
              ptmms[j].lda = ptmms[k].lda;
              ptmms[j].ldb = ptmms[k].ldb;
              ptmms[j].ldc = ptmms[k].ldc;
              ptmms[j].M = ptmms[k].M;
              ptmms[j].N = ptmms[k].N;
              ptmms[j].K = ptmms[k].K;
              ptmms[j].ownC = ptmms[k].ownC;
              ptmms[j].nCp = ptmms[k].nCp;
              ptmms[j].nCw = ptmms[k].nCw;
              ptmms[j].Cw = ptmms[k].Cw;
              ptmms[j].ldcw = ptmms[k].ldcw;
              for (h=0; h < ptmms[j].nCw; h++)
                 ptmms[j].Cinfp[h] = (ptmms[k].Cinfp[h] != ptmms+k) ?
                                     ptmms[k].Cinfp[h] : ptmms+j;
              for (h=ATL_NTHREADS-1; h >= ATL_NTHREADS-ptmms[j].nCp; h--)
                 ptmms[j].Cinfp[h] = (ptmms[k].Cinfp[h] != ptmms+k) ?
                                     ptmms[k].Cinfp[h] : ptmms+j;
              ptmms[k].K = 0;                /* entry k now empty */
              break;
           }
        }
        ATL_assert(kk >= P);
      }
   }
}
#endif

int ATL_thrdecompMM
   (ATL_TMMNODE_t *ptmms, const enum ATLAS_TRANS TA, const enum ATLAS_TRANS TB,
    ATL_CINT Mblks, const int mr, ATL_CINT Nblks, const int nr, ATL_CINT Kblks,
    const int kr, const void *A, ATL_INT lda, const void *B, const ATL_INT ldb,
    const void *C, ATL_CINT ldc, const int P, const int indx, const int COPYC)
{
/*
 * Divide only the M-dimension to cut down on JIK workspace & improve CE
 * efficiency if we have enough M blocks to make it worthwile
 */
#if ATL_NTHREADS != (1<<ATL_NTHRPOW2)
   int np;
#endif
   if ((Mblks >> ATL_NTHRPOW2) >= ATL_TMMMINMBLKS)
      return(ATL_thrdecompMM_M(ptmms, TA, TB, Mblks, mr, Nblks, nr, Kblks,
                               kr, A, lda, B, ldb, C, ldc, P, indx, COPYC));
   else
#if ATL_NTHREADS == (1<<ATL_NTHRPOW2)
      return(ATL_thrdecompMM_rec(ptmms, TA, TB, Mblks, mr, Nblks, nr, Kblks,
                                 kr, A, lda, B, ldb, C, ldc, P, indx, COPYC));
#else
   {
      np = ATL_thrdecompMM_rec(ptmms, TA, TB, Mblks, mr, Nblks, nr, Kblks,
                               kr, A, lda, B, ldb, C, ldc, P, indx, COPYC);
      if (np < ATL_NTHREADS)
         EnforceNonPwr2LO(ptmms, np);
      return(np);
   }
#endif
}

int ATL_StructIsInitMM(void *vp)
{
   return(((ATL_TMMNODE_t*)vp)->K);
}

void ATL_DoWorkMM(ATL_LAUNCHSTRUCT_t *lp, void *vp)
/*
 * Current implementation doesn't need lp, but if we had an error queue or
 * something similar we would need it, so keep it around
 */
{
   ATL_TMMNODE_t *mmp = vp;
/*
 * Allocate space if needed, do operation
 */
   if (mmp->nCw)
   {
/*
 *    If malloc fails, we'll do the operation during the combine
 */
      mmp->Cw = malloc(((mmp->ldcw)<<mmp->eltsh)*mmp->N+ATL_Cachelen);
      if (mmp->Cw)
      {
         mmp->gemmK(mmp->M, mmp->N, mmp->K, mmp->alpha, mmp->A, mmp->lda,
                    mmp->B, mmp->ldb, mmp->zero,
                    ATL_AlignPtr(mmp->Cw), mmp->ldcw);
      }
#ifdef DEBUG
      else
         fprintf(stderr, "%d: unable to allocate C(%dx%d)!!\n",
                 mmp->rank, mmp->M, mmp->N);
#endif
   }
   else  /* do GEMM directly into original C; no possibility of failure! */
      mmp->gemmK(mmp->M, mmp->N, mmp->K, mmp->alpha, mmp->A, mmp->lda,
                 mmp->B, mmp->ldb, mmp->beta, mmp->C, mmp->ldc);
}
