|
| 1 | +#ifndef INFINI_OPS_ASCEND_EMBEDDING_KERNEL_H_ |
| 2 | +#define INFINI_OPS_ASCEND_EMBEDDING_KERNEL_H_ |
| 3 | + |
| 4 | +#include <cassert> |
| 5 | + |
| 6 | +#include "acl/acl.h" |
| 7 | +#include "aclnn/aclnn_base.h" |
| 8 | +#include "aclnnop/aclnn_embedding.h" |
| 9 | +#include "base/embedding.h" |
| 10 | +#include "native/ascend/common.h" |
| 11 | +#include "native/ascend/workspace_pool_.h" |
| 12 | +#include "operator.h" |
| 13 | + |
| 14 | +namespace infini::ops { |
| 15 | + |
| 16 | +template <> |
| 17 | +class Operator<Embedding, Device::Type::kAscend> : public Embedding { |
| 18 | + public: |
| 19 | + Operator(const Tensor input_ids, const Tensor weight, Tensor out) |
| 20 | + : Embedding(input_ids, weight, out), |
| 21 | + input_ids_cache_(input_ids), |
| 22 | + weight_cache_(weight), |
| 23 | + out_cache_(out) { |
| 24 | + assert((weight_dtype_ == DataType::kFloat16 || |
| 25 | + weight_dtype_ == DataType::kBFloat16 || |
| 26 | + weight_dtype_ == DataType::kFloat32) && |
| 27 | + "`Embedding`: Ascend path supports `float16`, `bfloat16`, and " |
| 28 | + "`float32` weights"); |
| 29 | + } |
| 30 | + |
| 31 | + ~Operator() { |
| 32 | + if (!ascend::IsAclRuntimeAlive()) return; |
| 33 | + |
| 34 | + input_ids_cache_.release(); |
| 35 | + weight_cache_.release(); |
| 36 | + out_cache_.release(); |
| 37 | + } |
| 38 | + |
| 39 | + void operator()(const Tensor input_ids, const Tensor weight, |
| 40 | + Tensor out) const override { |
| 41 | + auto stream = static_cast<aclrtStream>(stream_); |
| 42 | + |
| 43 | + auto t_weight = weight_cache_.get(const_cast<void*>(weight.data())); |
| 44 | + auto t_input_ids = |
| 45 | + input_ids_cache_.get(const_cast<void*>(input_ids.data())); |
| 46 | + auto t_out = out_cache_.get(out.data()); |
| 47 | + |
| 48 | + if (!executor_) { |
| 49 | + auto ret = aclnnEmbeddingGetWorkspaceSize(t_weight, t_input_ids, t_out, |
| 50 | + &ws_size_, &executor_); |
| 51 | + assert(ret == ACL_SUCCESS && "`aclnnEmbeddingGetWorkspaceSize` failed"); |
| 52 | + aclSetAclOpExecutorRepeatable(executor_); |
| 53 | + } else { |
| 54 | + aclSetInputTensorAddr(executor_, 0, t_weight, |
| 55 | + const_cast<void*>(weight.data())); |
| 56 | + aclSetInputTensorAddr(executor_, 1, t_input_ids, |
| 57 | + const_cast<void*>(input_ids.data())); |
| 58 | + aclSetOutputTensorAddr(executor_, 0, t_out, out.data()); |
| 59 | + } |
| 60 | + |
| 61 | + auto& arena = ascend::GetWorkspacePool().Ensure(stream, ws_size_); |
| 62 | + auto ret = aclnnEmbedding(arena.buf, ws_size_, executor_, stream); |
| 63 | + assert(ret == ACL_SUCCESS && "`aclnnEmbedding` failed"); |
| 64 | + } |
| 65 | + |
| 66 | + private: |
| 67 | + mutable ascend::AclTensorCache input_ids_cache_; |
| 68 | + |
| 69 | + mutable ascend::AclTensorCache weight_cache_; |
| 70 | + |
| 71 | + mutable ascend::AclTensorCache out_cache_; |
| 72 | + |
| 73 | + mutable aclOpExecutor* executor_ = nullptr; |
| 74 | + |
| 75 | + mutable uint64_t ws_size_ = 0; |
| 76 | +}; |
| 77 | + |
| 78 | +} // namespace infini::ops |
| 79 | + |
| 80 | +#endif // INFINI_OPS_ASCEND_EMBEDDING_KERNEL_H_ |
0 commit comments