

#pragma once

#include <blam/detail/config.h>
#include <blam/adl/detail/customization_point.h>

BLAM_CUSTOMIZATION_POINT(gemm_batch)

#if defined(BLAM_USE_DECAY)
# include <blam/blas/level3/gemm.h>
#endif

namespace blam
{

// Backend entry point
template <typename ExecutionPolicy,
          typename Alpha, typename MA, typename MB,
          typename Beta, typename MC>
void
generic(blam::gemm_batch_t, const ExecutionPolicy& exec,
        Layout order, Op transA, Op transB,
        int m, int n, int k,
        const Alpha& alpha,
        const MA* A, int ldA, int loA,
        const MB* B, int ldB, int loB,
        const Beta& beta,
        MC* C, int ldC, int loC,
        int batch_size)
#if defined(BLAM_USE_DECAY)               // XXX: Remove
{
  for (int i = 0; i < batch_size; ++i) {
    blam::gemm(exec,
               order, transA, transB,
               m, n, k,
               alpha,
               A + i*loA, ldA,
               B + i*loB, ldB,
               beta,
               C + i*loC, ldC);
  }
}
#else
= delete;
#endif

// Default to ColMajor
template <typename ExecutionPolicy,
          typename Alpha, typename MA, typename MB,
          typename Beta, typename MC>
auto
generic(blam::gemm_batch_t, const ExecutionPolicy& exec,
        Op transA, Op transB,
        int m, int n, int k,
        const Alpha& alpha,
        const MA* A, int ldA, int loA,
        const MB* B, int ldB, int loB,
        const Beta& beta,
        MC* C, int ldC, int loC,
        int batch_size)
BLAM_DECLTYPE_AUTO_RETURN
(
  blam::gemm_batch(exec,
                   ColMajor, transA, transB,
                   m, n, k,
                   alpha,
                   A, ldA, loA,
                   B, ldB, loB,
                   beta,
                   C, ldC, loC,
                   batch_size)
)

// Default to NoTrans
template <typename ExecutionPolicy,
          typename Alpha, typename MA, typename MB,
          typename Beta, typename MC>
auto
generic(blam::gemm_batch_t, const ExecutionPolicy& exec,
        int m, int n, int k,
        const Alpha& alpha,
        const MA* A, int ldA, int loA,
        const MB* B, int ldB, int loB,
        const Beta& beta,
        MC* C, int ldC, int loC,
        int batch_size)
BLAM_DECLTYPE_AUTO_RETURN
(
  blam::gemm_batch(exec,
                   NoTrans, NoTrans,
                   m, n, k,
                   alpha,
                   A, ldA, loA,
                   B, ldB, loB,
                   beta,
                   C, ldC, loC,
                   batch_size)
)

} // end namespace blam
