Я хочу написать собственный оператор (в C ++), как для CPU, так и для GPU (в CUDA).Это довольно просто, за исключением: как мне назвать BLAS (s) gemm (в обоих случаях для CPU и GPU)?
Я немного посмотрел через LSTMBlock op ( lstm_ops.h, lstm_ops.cc , lstm_ops_gpu.cu.cc , blas_gemm.h , blas_gemm.cc ), и это кажетсяиспользовать Eigen для случая с процессором и ctx->op_device_context()->stream()->ThenBlasGemm() для случая с графическим процессором.
Было бы неплохо, если бы API ThenBlasGemm был бы также доступен для процессора, но, похоже, это был только графический процессор.(Может быть, я должен сделать запрос функции для этого?)
Но использование Eigen для процессора (или также gemm в целом), кажется, не лучшая вещь для всех случаев.Я просматривал MatMul op ( matmul_op.h , matmul_op.cc , mkl_matmul_op.cc , gemm_functors.h )и, кажется, существует много особых случаев (например, использование gemm или gemv, использование MKL, если доступно, и т. д.).Я не думаю, что было бы хорошей идеей попытаться воспроизвести все это (что в основном будет копировать и вставлять код).
Есть ли какой-либо простой в использовании API для matmul, который яможно использовать в моей пользовательской операции?(Может быть, я должен сделать запрос функции для этого?)
Или, может быть, есть ли способ вызвать существующий MatMulOp::Compute внутри моей собственной операции?