c++ - TensorFlow，自定义 C++ 运算，使用现有的 MatMul

Question

我想为 CPU 和 GPU（在 CUDA 中）编写一个自定义操作（在 C++ 中）。这很简单，除了：我如何调用 BLAS (s)gemm （在这两种情况下，对于 CPU 和 GPU）？

我看了一点LSTMBlock操作（lstm_ops.h，lstm_ops.cc，lstm_ops_gpu.cu.cc，blas_gemm.h，blas_gemm.cc），这似乎将 Eigen 用于 CPU 案例和ctx->op_device_context()->stream()->ThenBlasGemm()GPU 案例。

如果ThenBlasGemmAPI 也可用于 CPU 但它似乎仅适用于 GPU，那就太好了。（也许我应该为此做一个功能请求？）

但是，将 Eigen 用于 CPU 案例（或者一般来说也是 gemm）似乎并不是在所有情况下都最好做的事情。我正在查看MatMul操作（matmul_op.h、matmul_op.cc、mkl_matmul_op.cc、gemm_functors.h），似乎有很多特殊情况（例如，使用 gemm 或 gemv，如果可用，使用 MKL 等）。我认为尝试复制所有这些（主要是复制和粘贴代码）并不是一个好主意。

是否有任何易于使用的 APImatmul可以在我的自定义操作中使用？（也许我应该为此做一个功能请求？）

或者，有没有办法MatMulOp::Compute在我自己的操作中调用现有的？

标签： c++tensorflowblascublas