Skip to content

feat!: complete embedding PyTorch Python API #821

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
voltjia wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

feat!: complete embedding PyTorch Python API #821

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 47 additions & 13 deletions src/base/embedding.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
#define INFINI_OPS_BASE_EMBEDDING_H_

#include <cstddef>
#include <optional>

#include "data_type.h"
#include "operator.h"
Expand All @@ -11,7 +12,9 @@ namespace infini::ops {

class Embedding : public Operator<Embedding> {
public:
Embedding(const Tensor input, const Tensor weight, const int64_t padding_idx,
Embedding(const Tensor input, const Tensor weight,
const std::optional<int64_t> padding_idx,
const std::optional<double> max_norm, const double norm_type,
const bool scale_grad_by_freq, const bool sparse, Tensor out)
: input_shape_{input.shape()},
weight_shape_{weight.shape()},
Expand All @@ -26,6 +29,8 @@ class Embedding : public Operator<Embedding> {
vocab_size_{weight.size(0)},
embedding_dim_{weight.size(1)},
padding_idx_{padding_idx},
max_norm_{max_norm},
norm_type_{norm_type},
scale_grad_by_freq_{scale_grad_by_freq},
sparse_{sparse} {
assert(weight.ndim() == 2 && "`Embedding` requires 2D `weight`");
Expand All @@ -49,28 +54,53 @@ class Embedding : public Operator<Embedding> {
"`Embedding` supports float32, float16, and bfloat16 weights only");
assert(out_dtype_ == weight_dtype_ &&
"`Embedding` output dtype must match `weight` dtype");
assert(padding_idx_ >= -static_cast<int64_t>(vocab_size_) &&
padding_idx_ < static_cast<int64_t>(vocab_size_) &&
assert((!padding_idx_.has_value() ||
(*padding_idx_ >= -static_cast<int64_t>(vocab_size_) &&
*padding_idx_ < static_cast<int64_t>(vocab_size_))) &&
"`Embedding` padding_idx must be within the weight rows");
}

Embedding(const Tensor input, const Tensor weight, Tensor out)
: Embedding{input, weight, -1, false, false, out} {}
: Embedding{input, weight, std::nullopt, std::nullopt,
2.0, false, false, out} {}

/// \deprecated Use the overload that also accepts `max_norm` and
/// `norm_type` instead.
[[deprecated("Use the PyTorch-compatible overload instead.")]]
Embedding(const Tensor input, const Tensor weight, const int64_t padding_idx,
const bool scale_grad_by_freq, const bool sparse, Tensor out)
: Embedding{input, weight, padding_idx,
std::nullopt, 2.0, scale_grad_by_freq,
sparse, out} {}

virtual void operator()(const Tensor input, const Tensor weight,
const int64_t padding_idx,
const bool scale_grad_by_freq, const bool sparse,
Tensor out) const = 0;
const std::optional<int64_t> padding_idx,
const std::optional<double> max_norm,
const double norm_type, const bool scale_grad_by_freq,
const bool sparse, Tensor out) const = 0;

void operator()(const Tensor input, const Tensor weight, Tensor out) const {
(*this)(input, weight, -1, false, false, out);
(*this)(input, weight, std::nullopt, std::nullopt, 2.0, false, false, out);
}

/// \deprecated Use the overload that also accepts `max_norm` and
/// `norm_type` instead.
[[deprecated("Use the PyTorch-compatible overload instead.")]]
void operator()(const Tensor input, const Tensor weight,
const int64_t padding_idx, const bool scale_grad_by_freq,
const bool sparse, Tensor out) const {
(*this)(input, weight, std::optional<int64_t>{padding_idx}, std::nullopt,
2.0, scale_grad_by_freq, sparse, out);
}

template <typename TensorLike>
static auto MakeReturnValue(const TensorLike& input, const TensorLike& weight,
const int64_t /*padding_idx*/ = -1,
const bool /*scale_grad_by_freq*/ = false,
const bool /*sparse*/ = false) {
static auto MakeReturnValue(
const TensorLike& input, const TensorLike& weight,
const std::optional<int64_t> /*padding_idx*/ = std::nullopt,
const std::optional<double> /*max_norm*/ = std::nullopt,
const double /*norm_type*/ = 2.0,
const bool /*scale_grad_by_freq*/ = false,
const bool /*sparse*/ = false) {
auto out_shape = input.shape();
out_shape.push_back(weight.size(1));

Expand Down Expand Up @@ -112,7 +142,11 @@ class Embedding : public Operator<Embedding> {

Tensor::Size embedding_dim_{0};

int64_t padding_idx_{0};
std::optional<int64_t> padding_idx_{};

std::optional<double> max_norm_{};

double norm_type_{2.0};

bool scale_grad_by_freq_{false};

Expand Down
72 changes: 57 additions & 15 deletions src/native/ascend/ops/embedding/kernel.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
#ifndef INFINI_OPS_ASCEND_EMBEDDING_KERNEL_H_
#define INFINI_OPS_ASCEND_EMBEDDING_KERNEL_H_

#include <algorithm>
#include <cassert>

#include "acl/acl.h"
#include "aclnn/aclnn_base.h"
#include "aclnnop/aclnn_embedding.h"
#include "aclnnop/aclnn_embedding_renorm.h"
#include "base/embedding.h"
#include "native/ascend/common.h"
#include "native/ascend/workspace_pool_.h"
Expand All @@ -16,9 +18,12 @@ namespace infini::ops {
template <>
class Operator<Embedding, Device::Type::kAscend> : public Embedding {
public:
Operator(const Tensor input, const Tensor weight, const int64_t padding_idx,
Operator(const Tensor input, const Tensor weight,
const std::optional<int64_t> padding_idx,
const std::optional<double> max_norm, const double norm_type,
const bool scale_grad_by_freq, const bool sparse, Tensor out)
: Embedding(input, weight, padding_idx, scale_grad_by_freq, sparse, out),
: Embedding(input, weight, padding_idx, max_norm, norm_type,
scale_grad_by_freq, sparse, out),
input_cache_(input),
weight_cache_(weight),
out_cache_(out) {
Expand All @@ -30,7 +35,16 @@ class Operator<Embedding, Device::Type::kAscend> : public Embedding {
}

Operator(const Tensor input, const Tensor weight, Tensor out)
: Operator(input, weight, -1, false, false, out) {}
: Operator(input, weight, std::nullopt, std::nullopt, 2.0, false, false,
out) {}

/// \deprecated Use the overload that also accepts `max_norm` and
/// `norm_type` instead.
[[deprecated("Use the PyTorch-compatible overload instead.")]]
Operator(const Tensor input, const Tensor weight, const int64_t padding_idx,
const bool scale_grad_by_freq, const bool sparse, Tensor out)
: Operator(input, weight, padding_idx, std::nullopt, 2.0,
scale_grad_by_freq, sparse, out) {}

~Operator() {
if (!ascend::IsAclRuntimeAlive()) return;
Expand All @@ -41,7 +55,8 @@ class Operator<Embedding, Device::Type::kAscend> : public Embedding {
}

void operator()(const Tensor input, const Tensor weight,
const int64_t /*padding_idx*/,
const std::optional<int64_t> /*padding_idx*/,
const std::optional<double> max_norm, const double norm_type,
const bool /*scale_grad_by_freq*/, const bool /*sparse*/,
Tensor out) const override {
auto stream = static_cast<aclrtStream>(stream_);
Expand All @@ -50,21 +65,44 @@ class Operator<Embedding, Device::Type::kAscend> : public Embedding {
auto t_input = input_cache_.get(const_cast<void*>(input.data()));
auto t_out = out_cache_.get(out.data());

if (!executor_) {
auto ret = aclnnEmbeddingGetWorkspaceSize(t_weight, t_input, t_out,
&ws_size_, &executor_);
if (max_norm.has_value() && !renorm_executor_) {
auto ret = aclnnEmbeddingRenormGetWorkspaceSize(
t_weight, t_input, *max_norm, norm_type, &renorm_ws_size_,
&renorm_executor_);
assert(ret == ACL_SUCCESS &&
"`aclnnEmbeddingRenormGetWorkspaceSize` failed");
aclSetAclOpExecutorRepeatable(renorm_executor_);
} else if (max_norm.has_value()) {
aclSetInputTensorAddr(renorm_executor_, 0, t_weight,
const_cast<void*>(weight.data()));
aclSetInputTensorAddr(renorm_executor_, 1, t_input,
const_cast<void*>(input.data()));
}

if (!embedding_executor_) {
auto ret = aclnnEmbeddingGetWorkspaceSize(
t_weight, t_input, t_out, &embedding_ws_size_, &embedding_executor_);
assert(ret == ACL_SUCCESS && "`aclnnEmbeddingGetWorkspaceSize` failed");
aclSetAclOpExecutorRepeatable(executor_);
aclSetAclOpExecutorRepeatable(embedding_executor_);
} else {
aclSetInputTensorAddr(executor_, 0, t_weight,
aclSetInputTensorAddr(embedding_executor_, 0, t_weight,
const_cast<void*>(weight.data()));
aclSetInputTensorAddr(executor_, 1, t_input,
aclSetInputTensorAddr(embedding_executor_, 1, t_input,
const_cast<void*>(input.data()));
aclSetOutputTensorAddr(executor_, 0, t_out, out.data());
aclSetOutputTensorAddr(embedding_executor_, 0, t_out, out.data());
}

const auto workspace_size = std::max(renorm_ws_size_, embedding_ws_size_);
auto& arena = ascend::GetWorkspacePool().Ensure(stream, workspace_size);

if (max_norm.has_value()) {
auto ret = aclnnEmbeddingRenorm(arena.buf, renorm_ws_size_,
renorm_executor_, stream);
assert(ret == ACL_SUCCESS && "`aclnnEmbeddingRenorm` failed");
}

auto& arena = ascend::GetWorkspacePool().Ensure(stream, ws_size_);
auto ret = aclnnEmbedding(arena.buf, ws_size_, executor_, stream);
auto ret = aclnnEmbedding(arena.buf, embedding_ws_size_,
embedding_executor_, stream);
assert(ret == ACL_SUCCESS && "`aclnnEmbedding` failed");
}

Expand All @@ -75,9 +113,13 @@ class Operator<Embedding, Device::Type::kAscend> : public Embedding {

mutable ascend::AclTensorCache out_cache_;

mutable aclOpExecutor* executor_ = nullptr;
mutable aclOpExecutor* renorm_executor_ = nullptr;

mutable aclOpExecutor* embedding_executor_ = nullptr;

mutable uint64_t renorm_ws_size_ = 0;

mutable uint64_t ws_size_ = 0;
mutable uint64_t embedding_ws_size_ = 0;
};

} // namespace infini::ops
Expand Down
Loading
Loading