<commit_before>/*Copyright (c) 2011, <NAME>, all rights reserved.*/

#include "dist_tensor_internal.h"
#include "../shared/offload.h"

#if (defined BGP || defined BGQ)
#define BLACS_GRIDINFO blacs_gridinfo
#else
#define BLACS_GRIDINFO blacs_gridinfo_
#endif
#ifdef USE_SCALAPACK
extern "C" {
  void BLACS_GRIDINFO(int *, int *, int *, int *, int *);
}
#else
inline
void BLACS_GRIDINFO(int *, int *, int *, int *, int *) { assert(0); }
#endif

template<typename dtype, int is_herm_A, int is_herm_B>
int gemm_ctr(  dtype  const     alpha,
               dtype  const *   A,
               int const        ndim_A,
               int const *      edge_len_A,
               int const *      lda_A,
               int const *      sym_A,
               int const *      idx_map_A,
               dtype  const *   B,
               int const        ndim_B,
               int const *      edge_len_B,
               int const *      lda_B,
               int const *      sym_B,
               int const *      idx_map_B,
               dtype  const     beta,
               dtype  *         C,
               int const        ndim_C,
               int const *      edge_len_C,
               int const *      lda_C,
               int const *      sym_C,
               int const *      idx_map_C){
  char ta, tb;
  int n, m, k;
  int la_A, la_B, la_C;
  LIBT_ASSERT(ndim_A == 2);
  LIBT_ASSERT(ndim_B == 2);
  LIBT_ASSERT(ndim_C == 2);

  if (idx_map_A[0] == 0){
    k = edge_len_A[0];
    m = edge_len_A[1];
    la_A = k;
    if (is_herm_A)
      ta = 'C';
    else
      ta = 'T';
  } else {
    k = edge_len_A[1];
    m = edge_len_A[0];
    la_A = m;
    ta = 'N';
  }
  if (idx_map_B[0] == 0) {
    LIBT_ASSERT(k==edge_len_B[0]);
    n = edge_len_B[1];
    la_B = k;
    tb = 'N';
  } else {
    LIBT_ASSERT(k==edge_len_B[1]);
    n = edge_len_B[0];
    la_B = n;
    if (is_herm_B)
      tb = 'C';
    else
      tb = 'T';
  }
  LIBT_ASSERT(m==edge_len_C[0]);
  LIBT_ASSERT(n==edge_len_C[1]);
  la_C = m;


#ifdef OFFLOAD
  TAU_FSTART(offload_alloc);
  offload_ptr<dtype> ptr_A(m*k);
  offload_ptr<dtype> ptr_B(k*n);
  offload_ptr<dtype> ptr_C(m*n);
  TAU_FSTOP(offload_alloc);
  TAU_FSTART(offload_upload);
  ptr_A.upload(A);
  ptr_B.upload(B);
  ptr_C.upload(C);
  TAU_FSTOP(offload_upload);
  TAU_FSTART(offload_gemm);
  TAU_FSTART(dgemm);
  offload_gemm<dtype>(ta, tb, m, n, k, alpha, 
                      ptr_A, la_A,
                      ptr_B, la_B, beta,
                      ptr_C, la_C);
  TAU_FSTOP(dgemm);
  TAU_FSTART(offload_download);
  ptr_C.download(C);
  TAU_FSTOP(offload_download);
#else
  TAU_FSTART(dgemm);<commit_msg>upload/download now no longer necessary for scala_backednd. Getting decent single-GPU perf for the benchmark now.<commit_after>/*Copyright (c) 2011, <NAME>, all rights reserved.*/

#include "dist_tensor_internal.h"
#include "../shared/offload.h"

#if (defined BGP || defined BGQ)
#define BLACS_GRIDINFO blacs_gridinfo
#else
#define BLACS_GRIDINFO blacs_gridinfo_
#endif
#ifdef USE_SCALAPACK
extern "C" {
  void BLACS_GRIDINFO(int *, int *, int *, int *, int *);
}
#else
inline
void BLACS_GRIDINFO(int *, int *, int *, int *, int *) { assert(0); }
#endif

template<typename dtype, int is_herm_A, int is_herm_B>
int gemm_ctr(  dtype  const     alpha,
               dtype  const *   A,
               int const        ndim_A,
               int const *      edge_len_A,
               int const *      lda_A,
               int const *      sym_A,
               int const *      idx_map_A,
               dtype  const *   B,
               int const        ndim_B,
               int const *      edge_len_B,
               int const *      lda_B,
               int const *      sym_B,
               int const *      idx_map_B,
               dtype  const     beta,
               dtype  *         C,
               int const        ndim_C,
               int const *      edge_len_C,
               int const *      lda_C,
               int const *      sym_C,
               int const *      idx_map_C){
  char ta, tb;
  int n, m, k;
  int la_A, la_B, la_C;
  LIBT_ASSERT(ndim_A == 2);
  LIBT_ASSERT(ndim_B == 2);
  LIBT_ASSERT(ndim_C == 2);

  if (idx_map_A[0] == 0){
    k = edge_len_A[0];
    m = edge_len_A[1];
    la_A = k;
    if (is_herm_A)
      ta = 'C';
    else
      ta = 'T';
  } else {
    k = edge_len_A[1];
    m = edge_len_A[0];
    la_A = m;
    ta = 'N';
  }
  if (idx_map_B[0] == 0) {
    LIBT_ASSERT(k==edge_len_B[0]);
    n = edge_len_B[1];
    la_B = k;
    tb = 'N';
  } else {
    LIBT_ASSERT(k==edge_len_B[1]);
    n = edge_len_B[0];
    la_B = n;
    if (is_herm_B)
      tb = 'C';
    else
      tb = 'T';
  }
  LIBT_ASSERT(m==edge_len_C[0]);
  LIBT_ASSERT(n==edge_len_C[1]);
  la_C = m;


#ifdef OFFLOAD
  /*TAU_FSTART(offload_alloc);
  offload_ptr<dtype> ptr_A(m*k);
  offload_ptr<dtype> ptr_B(k*n);
  offload_ptr<dtype> ptr_C(m*n);
  TAU_FSTOP(offload_alloc);
  TAU_FSTART(offload_upload);
  ptr_A.upload(A);
  ptr_B.upload(B);
  ptr_C.upload(C);
  TAU_FSTOP(offload_upload);
  TAU_FSTART(offload_gemm);*/
  TAU_FSTART(offload_dgemm);
  offload_gemm<dtype>(ta, tb, m, n, k, alpha, 
                      A, la_A,
                      B, la_B, beta,
                      C, la_C);
  TAU_FSTOP(offload_dgemm);
/*  TAU_FSTART(offload_download);
  ptr_C.download(C);
  TAU_FSTOP(offload_download);*/
#else
  TAU_FSTART(dgemm);
{'fraction_non_alphanumeric': 0.12319644839067703, 'fraction_numerical': 0.008139104698483167, 'mean_word_length': 1.9148247978436657, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6867522', 'n_tokens_mistral': 2194, 'n_tokens_neox': 2153, 'n_words': 445}