-
Notifications
You must be signed in to change notification settings - Fork 66
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
#8118: ported ttnn::embedding to C++
- Loading branch information
Showing
5 changed files
with
141 additions
and
100 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,62 @@ | ||
// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc. | ||
// | ||
// SPDX-License-Identifier: Apache-2.0 | ||
|
||
#pragma once | ||
|
||
#include <pybind11/pybind11.h> | ||
#include <pybind11/stl.h> | ||
|
||
#include "ttnn/cpp/pybind11/decorators.hpp" | ||
#include "ttnn/operations/embedding.hpp" | ||
|
||
namespace py = pybind11; | ||
|
||
namespace ttnn { | ||
namespace operations { | ||
namespace embedding { | ||
|
||
void py_module(py::module& module) { | ||
bind_registered_operation( | ||
module, | ||
ttnn::embedding, | ||
R"doc( | ||
embedding(inxput_tensor: ttnn.Tensor, weight: ttnn.Tensor, *, pad_token: Optional[int] = None, layout: ttnn.Layout = ttnn.ROW_MAJOR_LAYOUT, memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG) -> ttnn.Tensor | ||
Retrieves word embeddings using input_tensor. The input_tensor is a list of indices, and the embedding matrix, and the output is the corresponding word embeddings. | ||
Args: | ||
* :attr:`input_tensor`: the indices ttnn.Tensor | ||
* :attr:`weight`: the embeddings ttnn.Tensor that correspond to the indices ttnn.Tensor | ||
Keyword Args: | ||
* :attr:`pad_token`: the padding token. Default is None. | ||
* :attr:`layout`: the layout of the input and output tensors. Default is ttnn.ROW_MAJOR_LAYOUT. | ||
* :attr:`memory_config`: the memory configuration of the output tensor. Default is ttnn.DRAM_MEMORY_CONFIG. | ||
Example:: | ||
>>> device_id = 0 | ||
>>> device = ttnn.open_device(device_id=device_id) | ||
>>> input_tensor = ttnn.to_device(ttnn.from_torch(torch.tensor([[1, 2, 4, 5], [4, 3, 2, 9]]), dtype=ttnn.uint32), device) | ||
>>> # an embedding matrix containing 10 tensors of size 4 | ||
>>> weight = ttnn.to_device(ttnn.from_torch(torch.rand(10, 4), dtype=ttnn.bfloat16), device) | ||
>>> ttnn.embedding(input_tensor, weight) | ||
ttnn.Tensor([ [[1, 0.106445, 0.988281, 0.59375], | ||
[0.212891, 0.964844, 0.199219, 0.996094], | ||
[3.78362e-38, 0, 7.89785e-39, 0], | ||
[8.04479e-38, 0, 1.25815e-38, 0]], | ||
[[2.71833e-38, 0, 3.59995e-38, 0], | ||
[7.60398e-38, 0, 1.83671e-38, 0], | ||
[2.22242e-38, 0, 1.88263e-38, 0], | ||
[1.35917e-38, 0, 4.49994e-39, 0]]], dtype=bfloat16 ))doc", | ||
ttnn::pybind_arguments_t{ | ||
py::arg("input_tensor"), | ||
py::arg("weight"), | ||
py::arg("pad_token") = std::nullopt, | ||
py::arg("layout") = ttnn::ROW_MAJOR_LAYOUT, | ||
py::arg("memory_config") = std::nullopt}); | ||
} | ||
|
||
} // namespace embedding | ||
} // namespace operations | ||
} // namespace ttnn |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,72 @@ | ||
// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc. | ||
// | ||
// SPDX-License-Identifier: Apache-2.0 | ||
|
||
#pragma once | ||
|
||
#include "tt_eager/tt_dnn/op_library/embeddings/embeddings_op.hpp" | ||
#include "tt_eager/tt_dnn/op_library/run_operation.hpp" | ||
#include "ttnn/decorators.hpp" | ||
#include "ttnn/operations/core.hpp" | ||
#include "ttnn/validation.hpp" | ||
|
||
namespace ttnn { | ||
|
||
namespace operations { | ||
|
||
namespace embedding { | ||
|
||
using EmbeddingsType = tt::tt_metal::EmbeddingsType; | ||
|
||
struct Embedding { | ||
static const std::array<ttnn::TensorSchema, 2> input_tensor_schemas() { | ||
return { | ||
ttnn::TensorSchema{2, 2, {ttnn::uint32}, {ttnn::ROW_MAJOR_LAYOUT}, true, false, false, false}, | ||
ttnn::TensorSchema{2, 4, {ttnn::bfloat16}, {ttnn::ROW_MAJOR_LAYOUT}, true, false, false, false}}; | ||
} | ||
|
||
template <typename... Args> | ||
static auto input_tensors_to_validate(const Tensor& input_tensor, const Tensor& weight, Args&&... args) { | ||
return std::make_tuple(input_tensor, weight); | ||
} | ||
|
||
static Tensor execute( | ||
const Tensor& input_tensor_arg, | ||
const Tensor& weight_arg, | ||
const std::optional<int>& pad_token = std::nullopt, | ||
const Layout& layout = ttnn::ROW_MAJOR_LAYOUT, | ||
const std::optional<MemoryConfig>& memory_config = std::nullopt) { | ||
auto embeddings_type = EmbeddingsType::GENERIC; | ||
if (pad_token.has_value()) { | ||
embeddings_type = EmbeddingsType::PADDED; | ||
} | ||
|
||
auto hidden_embedding_dim = weight_arg.get_shape()[-1]; | ||
auto padded_hidden_embedding_dim = weight_arg.get_shape().with_tile_padding()[-1]; | ||
auto weight = ttnn::unsqueeze_to_4D(weight_arg); | ||
|
||
auto batch_size = input_tensor_arg.get_shape()[0]; | ||
auto sentence_size = input_tensor_arg.get_shape()[-1]; | ||
auto input_tensor = ttnn::reshape(input_tensor_arg, ttnn::Shape{{batch_size, 1, 1, sentence_size}}); | ||
|
||
bool tilized = layout == ttnn::TILE_LAYOUT; | ||
auto embeddings = operation::run( | ||
tt::tt_metal::Embeddings{ | ||
.output_mem_config = memory_config.value_or(input_tensor.memory_config()), | ||
.tilized = tilized, | ||
.embeddings_type = embeddings_type, | ||
.pad_token = pad_token, | ||
.output_dtype = weight.get_dtype()}, | ||
{input_tensor, weight}) | ||
.at(0); | ||
embeddings = ttnn::reshape(embeddings, ttnn::Shape{{batch_size, sentence_size, hidden_embedding_dim}}); | ||
return embeddings; | ||
} | ||
}; | ||
|
||
} // namespace embedding | ||
} // namespace operations | ||
|
||
constexpr auto embedding = ttnn::register_operation<ttnn::operations::embedding::Embedding>("ttnn::embedding"); | ||
|
||
} // namespace ttnn |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters