我想编写一个自定义的OP(在C++中),既适用于CPU,也适用于GPU(CUDA)。
这很直接,除了:如何调用blas(s)gemm(在这两种情况下,对于cpu和gpu)?
我看了一下
LSTMBlock
运算器
lstm_ops.h
,
lstm_ops.cc
,
lstm_ops_gpu.cu.cc
,
blas_gemm.h
,
blas_gemm.cc
,这似乎使用了eign作为cpu的外壳,并且
ctx->op_device_context()->stream()->ThenBlasGemm()
在GPU的情况下。
如果
ThenBlasGemm
api本来也可以用于cpu,但它似乎只是gpu。(也许我应该为此做一个功能请求?)
但是,在所有情况下,使用eigen来处理cpu情况(或者通常使用gemm)并不是最好的做法。我在看
MatMul
运算器
matmul_op.h
,
matmul_op.cc
,
mkl_matmul_op.cc
,
gemm_functors.h
)而且似乎有许多特殊情况(例如,使用gemm或gemv、使用mkl(如果可用)等)。我不认为尝试复制所有这些(主要是复制和粘贴代码)是一个好主意。
是否有易于使用的api
matmul
我可以在我的自定义操作中使用哪个?(也许我应该为此做一个功能请求?)
或者,有没有办法把现有的
MatMulOp::Compute
在我自己的手术室里?
