如何在自定义的TensorFlow C++ op中调用sgemm函数

Question

我正在遵循如何在C++中定义自己的TensorFlow op的教程。

我想在我的自定义TensorFlow C++ op中调用sgemm。我正在编写两个内核，一个用于CUDA，另一个用于CPU。每种情况下sgemm的调用会是什么样子？或者有没有一种通用的方法适用于两种情况？

我尝试了这段代码，但由于缺少包含文件（请参见此处），我无法使其正常工作：

auto dev_ctx = context->op_device_context();
auto* dev_stream = dev_ctx->stream();
OP_REQUIRES(context, dev_stream, errors::Internal("No stream available."));

bool blas_launch_status =
    dev_stream
         ->ThenBlasGemm(...

此外，我不确定这是通用的还是仅适用于CUDA。

这是否有文档记录？

我如何在我的GPU/CUDA实现中调用？更准确地说，如何获取？

我在TF代码中搜索了一下，发现class CUDABlas提供了cuBLAS函数的包装器。我需要使用它吗？还是可以直接使用？我想我需要使用包装器，因为这将确保CUDA流执行器处于正常状态。我如何使用包装器？

我还发现了contrib/rnn/kernels/blas_gemm.cc和core/kernels/matmul_op.cc，它们似乎可以做我想要的事情。代码看起来像这样：

#define EIGEN_USE_THREADS

#if GOOGLE_CUDA
#include "tensorflow/core/platform/stream_executor.h"
#endif  // GOOGLE_CUDA

#include "tensorflow/contrib/rnn/kernels/blas_gemm.h"
#include "tensorflow/core/framework/op_kernel.h"
namespace tensorflow {

#if GOOGLE_CUDA
namespace {
template <typename T>
perftools::gputools::DeviceMemory<T> AsDeviceMemory(const T* cuda_memory) {
  perftools::gputools::DeviceMemoryBase wrapped(const_cast<T*>(cuda_memory));
  perftools::gputools::DeviceMemory<T> typed(wrapped);
  return typed;
}
}  // namespace
#endif  // GOOGLE_CUDA

namespace functor {
template <typename T>
void TensorCuBlasGemm<T>::operator()(OpKernelContext* ctx,
                                     bool transa, bool transb, uint64 m,
                                     uint64 n, uint64 k, T alpha, const T* a,
                                     int lda, const T* b, int ldb, T beta, T* c,
                                     int ldc) {
#if GOOGLE_CUDA
  perftools::gputools::blas::Transpose trans[] = {
      perftools::gputools::blas::Transpose::kNoTranspose,
      perftools::gputools::blas::Transpose::kTranspose};

  auto a_ptr = AsDeviceMemory(a);
  auto b_ptr = AsDeviceMemory(b);
  auto c_ptr = AsDeviceMemory(c);

  bool blas_launch_status =
      ctx->op_device_context()
          ->stream()
          ->ThenBlasGemm(trans[transa], trans[transb], m, n, k, alpha, a_ptr,
                         lda, b_ptr, ldb, beta, &c_ptr, ldc)
          .ok();
  OP_REQUIRES(ctx, blas_launch_status, errors::Aborted("CuBlasGemm failed!"));
#else
  ctx->SetStatus(errors::InvalidArgument("CuBlasGemm needs CUDA."));
#endif
}

例如，在我的Compute(OpKernelContext* ctx)中，我会调用。

ctx->op_device_context()
      ->stream()
      ->ThenBlasGemm(...)

我尝试了一下，但似乎对我来说缺少一些包含头文件（TensorFlow 0.12.0 with GPU for Linux）。我收到了“fatal error: tensorflow/stream_executor/lib/status.h: No such file or directory”的错误消息。我已在此处上报告了这个问题。

是否有关于所有这些的文档，例如如何处理cuBLAS，或者这个“DeviceStream”接口，流执行逻辑等等？

我的当前解决方案有点hacky。对于CPU，我尝试链接系统上可用的一些Blas库，并使用其中的sgemm。对于CUDA，我链接到tensorflow/contrib/rnn/python/ops/_lstm_ops.so，因为在那里我找到了TensorCuBlasGemm，可以使用。请参见这里。基本上，在该contrib op中，他们面临着相同的问题，并提出了这个。但这部分取决于通常不可用的包含文件，请参见上面的问题。

- Albert

---

你是在询问是否要在你将要编写的CUDA内核中使用CUBLAS，还是从主机端操作符中调用CUBLAS？ - talonmies

@talonmies：这会有什么区别吗？我对两种情况都很感兴趣。在我的情况下，我编写了自己的GPU操作，应该调用一些自己的CUDA内核以及cublasSgemm。 - Albert

是的，这很重要。而且是两个不同的问题。一个与CUDA有关，另一个与tensorflow有关。我建议选择其中一个。 - talonmies

@talonmies：我扩展了我的问题。 - Albert

好的，这并不是一个关于CUDA编程的问题，更多地涉及到TF内部机制。 - talonmies

@Albert，你解决了如何在TensorFlow中使用流的问题吗？ - warunapww

Answer 1

你可以尝试以下方法，这对我今天有效： 在你的*.cu.cc文件开头：

#include <cublas_v2.h>
cublasHandle_t cublas_handle = NULL;

在同一个 *.cu.cc 文件中的 functor 实现中：

if (cublas_handle == NULL)
{
    assert(cublasCreate(&cublas_handle) == CUBLAS_STATUS_SUCCESS);
    assert(cublasSetStream(cublas_handle, d.stream()) == CUBLAS_STATUS_SUCCESS);
}

在*.cc文件中，将d作为参数从函数对象传入，其值为ctx->eigen_device<Eigen::GpuDevice>()

希望这可以帮到你，干杯！

- whywhywhy

---

如果现在dev_ctx->stream()可以工作，那么使用自定义cublas句柄的整个hack就不再需要了，对吗？因此，我的第一个代码片段应该已经可以正常工作了。 - Albert