text: #include "nervana.h"
  
  #include "caffe2/core/context_gpu.h"
  #include "caffe2/utils/math.h"
  
  namespace caffe2 {
  
  namespace math {
  
  // Caffe2 gemm provides a simpler interface to the gemm functions, with the
  // limitation that the data has to be contiguous in memory.
  template <>
  void Gemm<float, CUDAContext, NervanaEngine>(
      const CBLAS_TRANSPOSE TransA, const CBLAS_TRANSPOSE TransB,
      const int M, const int N, const int K, const float alpha, const float* A,
      const float* B, const float beta, float* C, CUDAContext* context) {
  
    // Note that cublas follows fortran order, so the order is different from
    // the cblas convention.
    int lda = (TransA == CblasNoTrans) ? K : M;
    int ldb = (TransB == CblasNoTrans) ? N : K;
    bool a_t = (TransA == CblasTrans);
    bool b_t = (TransB == CblasTrans);
    CHECK(nervana_sgemm(
        const_cast<float*>(A), const_cast<float*>(B), C, a_t, b_t, M, N, K,
        lda, ldb, N, alpha, beta, nullptr, false, false, context->cuda_stream()));
  }
  
  }  // namespace math
  }  // namespace caffe2
  