Skip to content

[Backend] Remove inner tensor type inside tensordesc#9851

Merged
peterbell10 merged 16 commits into
triton-lang:mainfrom
antiagainst:detach-tensordesc-shared-layout
Apr 9, 2026
Merged

[Backend] Remove inner tensor type inside tensordesc#9851
peterbell10 merged 16 commits into
triton-lang:mainfrom
antiagainst:detach-tensordesc-shared-layout

Conversation

@antiagainst
Copy link
Copy Markdown
Member

@antiagainst antiagainst commented Mar 25, 2026

This commit changes tt.tesnordesc type to be directly
composed of shape, element type, and shared layout.
It drops the previous tensor type wrapper around them.
The main reason is that the shared memory layout encoding
attaching inside tensor type was confusing since the layout
describes shared memory, not the tensor.

!tt.tensordesc<tensor<128x64xf16, #shared>> now becomes
!tt.tensordesc<128x64xf16, #shared>. This is also closer to
ttg.memdesc.

The shared memory layout encoding attribute was previously embedded
inside tensordesc's blockType as a RankedTensorType encoding.
This was confusing since the layout describes shared memory, not the
tensor.

Add an explicit optional $sharedLayout attribute to TensorDescType and
TensorDescIm2ColType, and a getSharedLayout() method on TensorDescInterface.
Syntax change: `!tt.tensordesc<tensor<128x64xf16, #enc>>` becomes
`!tt.tensordesc<tensor<128x64xf16>, #enc>`
@antiagainst antiagainst force-pushed the detach-tensordesc-shared-layout branch from a7c57b5 to ed050a7 Compare March 25, 2026 22:51
@lezcano
Copy link
Copy Markdown
Contributor

lezcano commented Mar 26, 2026

can you review this one @peterbell10 ?

Copy link
Copy Markdown
Contributor

@peterbell10 peterbell10 left a comment

Choose a reason for hiding this comment

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

If we want to go in this direction, I think the correct format should be:

!tt.tensordesc<1x2x128xf32, #shared1>

i.e. no RankedTensorType at all, in the same way that memdesc works.

@antiagainst antiagainst marked this pull request as ready for review March 26, 2026 17:29
@antiagainst antiagainst changed the title [Backend] Change to explicit shared layout attribute on tensordesc [Backend] Remove inner tensor type inside tensordesc Mar 26, 2026
@antiagainst
Copy link
Copy Markdown
Member Author

Yup makes sense. I modified to drop the intermediate tensor type wrapper entirely. Please take another look @peterbell10.

@antiagainst antiagainst requested a review from peterbell10 March 26, 2026 21:54
Copy link
Copy Markdown
Contributor

@peterbell10 peterbell10 left a comment

Choose a reason for hiding this comment

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

Mostly LGTM, just a few questions.

Comment thread include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUTypes.td Outdated
Comment thread python/src/gluon_ir.cc Outdated
Comment thread python/src/ir.cc Outdated
Comment thread python/src/ir.cc Outdated
Comment on lines +198 to +199
auto blockSize =
ttng::getTMABlockShape(blockType, /*packedSize=*/false, tmaMode);
auto shapePerCTA = ttg::getShapePerCTA(encoding, descTy.getShape());
auto blockSize = ttng::getTMABlockShape(encoding, shapePerCTA,
/*packedSize=*/false, tmaMode);
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Why this change?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I see, the non-RankedTensor overload takes shapePerCTA. Can't you just keep using descTy.getBlockType() since you added it as an interface method?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

The interface getBlockType() just returns a RankedTensorType without encoding attached now (given we want to avoid associate shared layout with tensors). But the getTMABlockShape variant that takes RankedTensorType still queries layout from the tensor so it won't work. Actually I think we should just remove it given now (done w/ 90946c5) no users.

Copy link
Copy Markdown
Contributor

@peterbell10 peterbell10 Apr 8, 2026

Choose a reason for hiding this comment

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

Okay, then give getTMABlockShape an overload takingTensorDescInterface?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

SG. Done with 6757366

@antiagainst antiagainst force-pushed the detach-tensordesc-shared-layout branch from a7f5f75 to e66e7d9 Compare April 1, 2026 04:55
@lezcano lezcano requested a review from peterbell10 April 9, 2026 09:35
@peterbell10 peterbell10 merged commit 354df74 into triton-lang:main Apr 9, 2026
9 checks passed
plognjen pushed a commit to plognjen/triton that referenced this pull request Apr 14, 2026
This commit changes `tt.tensordesc` type to be directly
composed of shape, element type, and shared layout.
It drops the previous tensor type wrapper around them.
The main reason is that the shared memory layout encoding
attaching inside tensor type was confusing since the layout
describes shared memory, not the tensor.

`!tt.tensordesc<tensor<128x64xf16, #shared>>` now becomes
`!tt.tensordesc<128x64xf16, #shared>`. This is also closer to
`ttg.memdesc`.
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.

4 participants