Skip to content
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.

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

Update eltwise flatbuffer to include extra params, move eltwise composite to own files #1121

Merged
merged 1 commit into from
Oct 31, 2024

Conversation

jnie-TT
Copy link
Contributor

@jnie-TT jnie-TT commented Oct 31, 2024

Related #1093

As we're bringing up more ops, having a single flatbuffer table that describes parameters for eltwise ops isn't sufficient to cover all cases. For example, unary composite ops like clamp could have 2 extra float parameters min and max. These ops are however still modelled as eltwise operations in ttnn, so we should model them as eltwise in runtime as well.

This PR introduces a new union type EltwiseOpParams, which is a collection of extra parameters required for specific eltwise ops. These extra params can later be extracted in runtime to input into the ttnn op. This will allow us to add eltwise ops that have extra parameters in the future.

For example, if we were to add a ClampOp as an EltwiseOp, we could create special params for it:

table ClampOpParams {
  min: float;
  max: float;
}

union EltwiseOpParams {
  ClampOpParams,
}

Then handle it in TTNNToFlatbuffer.cpp with something like:

::flatbuffers::Offset<::tt::target::ttnn::ClampOpParams>
createEltwiseOpParams(FlatbufferObjectCache &cache, ClampOp op) {
  auto min = op.getMin().convertToFloat();
  auto max = op.getMax().convertToFloat();
  return ::tt::target::ttnn::CreateClampOpParams(*cache.fbb, min, max);
}

template <typename EltwiseOp>
::flatbuffers::Offset<::tt::target::ttnn::EltwiseOp>
createEltwiseOp(FlatbufferObjectCache &cache, EltwiseOp op) {
  ...
  else if constexpr (std::is_same_v<EltwiseOp, ClampOp>) {
    type = ::tt::target::ttnn::EltwiseOpType::Clamp;
    paramsType = ::tt::target::ttnn::EltwiseOpParams::ClampOpParams;
    params = createEltwiseOpParams(cache, op).Union();
  }

Then in runtime extract the required params.

This PR also moves unary/binary composite ops to their own files to match the implementation of ttnn and to reduce the size of unary.cpp and binary.cpp (especially as we add more composite ops).

FYI @mmanzoorTT


namespace tt::runtime::ttnn::operations::unary::composite {

static void runEltwiseUnaryCompositeOP(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do we mean when we say "Composite" op in our runtime?

Copy link
Contributor Author

@jnie-TT jnie-TT Oct 31, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It reflects the scheme of ttnn (in src/tt-metal/ttnn/cpp/ttnn/operations/eltwise/unary/unary_composite.hpp). They have a concept of unaryCompositeOp that has a different execution function signature than ordinary unary ops (same for binary_composite). Alot of these ops have extra parameters aside from the input tensor like min/max for clamp.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see ok makes sense.

DEBUG_ASSERT((*rhs)->is_allocated());

// Switch the order of operands if the second operand requires broadcast
if ((*rhs)->volume() < (*lhs)->volume()) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should guard this underneath a workaround.

Copy link
Contributor Author

@jnie-TT jnie-TT Oct 31, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good, FYI @uazizTT I think you added this, is this still required in runtime? Are there plans to do this in the compiler?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Created issue #1124

Copy link
Contributor

@uazizTT uazizTT Oct 31, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good, FYI @uazizTT I think you added this, is this still required in runtime? Are there plans to do this in the compiler?

We don't intend to fix it in the compiler as it should be fixed in tt-metal, this is the reported issue in tt-metal to track it. tenstorrent/tt-metal#13566

Once tt-metal fixes it on their end, this workaround can be removed, but until then it can be guarded as a workaround in runtime.

@jnie-TT jnie-TT force-pushed the jnie/eltwise_params branch from 278a774 to 0af336b Compare October 31, 2024 14:52
… composite to its own file to match ttnn implementation
@jnie-TT jnie-TT force-pushed the jnie/eltwise_params branch from 0af336b to 0b5043d Compare October 31, 2024 15:12
Copy link
Contributor

@kmabeeTT kmabeeTT left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

@jnie-TT jnie-TT merged commit 4dc7502 into main Oct 31, 2024
13 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants