DPS version of matmul in TTNN

[azecevicTT]

Is there a reason why matmul (and linear) aren't modeled as DPS ops in TTNN? We model them as DPS ops in TTNN MLIR dialect, is there a plan to make them DPS ops in the future or should we change them to non DPS to reflect TTNN state.

[sdjordjevicTT]

Hi @bbradelTT, we discovered this during our onboarding in the TT-MLIR compiler. Can you give us some feedback on whether this inconsistency in the matmul op is a limitation or just needs to be implemented from your side?

[bbradelTT]

It just needs to be implemented from our side. The underlying device op has output tensors passed in to it that are created via operation::get_workers_for_op_output. It should be straightforward to expose that via an optional parameter. Having said that, the output tensor would need to be the right format, otherwise we would need to fail in validation.

How high of a priority would this work be for you and when would you need it by?

[sdjordjevicTT]

Thanks for the quick reply @bbradelTT! It isn't critical for now, but it would be great to have support for DPS matmul sometime in the near future, hence I will mark this as P1 if that sounds ok with you?

[bbradelTT]

@sdjordjevicTT sounds good. @asandhupatlaTT will work on it.

[sdjordjevicTT]

Great, thanks! Let's keep this issue open to track the progress.

[bbradelTT]

@asandhupatlaTT
Please try to see if it makes it easier if we convert to compute_output_specs for struct Matmul.
See ttnn/cpp/ttnn/operations/reduction/generic/device/reduce_op.hpp for an example.

[bbradelTT]

@sdjordjevicTT @azecevicTT which TTNN ops are DPS ops already? Using those as guides will help speed up development.

[sdjordjevicTT]

All the ops that define the optional output tensor in their APIs should be DPS.

For example, you can take a look at element-wise ops API definition:
https://github.com/tenstorrent/tt-metal/blob/8901511f737c15d181629ce4de93fae4a11da3b8/ttnn/cpp/ttnn/operations/eltwise/binary/binary.hpp#L29C1-L29C69

Or take a look at the embedding op and its optional output tensor:
https://github.com/tenstorrent/tt-metal/blob/8901511f737c15d181629ce4de93fae4a11da3b8/ttnn/cpp/ttnn/operations/embedding/embedding.hpp#L26C1-L26C70

[asandhupatlaTT]

fixed at #15285