[IR] Loosen tensor memory encoding checks added in #7748 #7784

peterbell10 · 2025-08-05T20:12:24Z

Commits in this PR

Revert "[LAYOUTS] Implement toLinearLayout for TensorMemoryEncodingAttr ([LAYOUTS] Implement toLinearLayout for TensorMemoryEncodingAttr #7748)"

This reverts commit 40335eb.
Reapply "[LAYOUTS] Implement toLinearLayout for TensorMemoryEncodingAttr ([LAYOUTS] Implement toLinearLayout for TensorMemoryEncodingAttr #7748)"

This reverts commit e6eb871.
Re-enable float8 tensor memory
Make shape errors more informative
Respond to comments

PR chain

👉 [IR] Loosen tensor memory encoding checks added in #7748 #7784 👈 YOU ARE HERE

…tr (#7748)" This reverts commit 40335eb. git-pr-chain: revert_layouts_implement_tolinearlayout__9b92

…ttr (#7748)" This reverts commit e6eb871.

peterbell10

Narrowed this down a bit as I think some of the verification failures were genuine.

peterbell10 · 2025-08-05T20:47:13Z

lib/Dialect/TritonGPU/IR/Types.cpp

-      return emitError() << "bitwidth must be 16 or 32";
+    if (bitwidth > 32) {
+      return emitError() << "bitwidth must be <= 32";
    }


Transmitting smaller dtypes through tmem seems to work fine, so removing this restriction.

it should error out ATM in ld/st?

triton/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TensorMemoryToLLVM.cpp

Lines 254 to 259 in dd58234

if (auto attr = dyn_cast<triton::nvidia_gpu::TensorMemoryEncodingAttr>(

memType.getEncoding())) {

info.blockM = attr.getBlockM();

info.blockN = attr.getBlockN();

assert((!attr.getUnpacked() || info.numElementsPer32B <= 2) &&

"unsupported unpacked layout");

In this case the tmem layout is packed, so that assert doesn't apply.

ah, I see, we can have packed for bitwidth=8

lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp

lezcano · 2025-08-05T22:19:45Z

lib/Dialect/TritonGPU/IR/Types.cpp

+    if (!enc.getUnpacked() && bitwidth > 16) {
+      return emitError() << "bitwidth must be <= 16 for packed tensor memory";


This was something I wanted to make more strict regardless. Any chance we could change the use case to leave it as it is?

I'm looking at the use case and it's passing fp8 as the lhs in tmem for a tcgen05_mma so there is no other way to do it. We should just fix load and store if they really are broken.

lezcano

Not sure about the two changes in the verifier... I did mean to write those. In particular the bitwidth == 8 case was failing in a test

lezcano

Can we then just support bitwidth=8 for the packed case? Otherwise lgtm

lezcano · 2025-08-06T00:15:19Z

lib/Dialect/TritonGPU/IR/Types.cpp

-      return emitError() << "bitwidth must be 16 or 32";
+    if (bitwidth > 32) {
+      return emitError() << "bitwidth must be <= 32";
    }


ah, I see, we can have packed for bitwidth=8

lezcano · 2025-08-06T00:22:44Z

lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp

-    return emitError() << "blockM must be 64 or 128";
+    return emitError() << "blockM must be 64 or 128 but got " << blockM;
  }
  if (!llvm::isPowerOf2_32(blockN) || blockN > 512) {


Unrelated, but I just realised that this should be blockN < 512 * (isUnpacked ? (32 / bitwidth) : 1). I can add it to another PR if you don't want to fix it in this one tho.

we could probably remove this restriction as it is more a restriction on allocation size

ThomasRaoux · 2025-08-06T00:28:18Z

lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp

-    return emitError() << "blockM must be 64 or 128";
+    return emitError() << "blockM must be 64 or 128 but got " << blockM;
  }
  if (!llvm::isPowerOf2_32(blockN) || blockN > 512) {


we could probably remove this restriction as it is more a restriction on allocation size

Revert "[LAYOUTS] Implement toLinearLayout for TensorMemoryEncodingAt…

e6eb871

…tr (#7748)" This reverts commit 40335eb. git-pr-chain: revert_layouts_implement_tolinearlayout__9b92

peterbell10 requested review from lezcano and ptillet as code owners August 5, 2025 20:12

Reapply "[LAYOUTS] Implement toLinearLayout for TensorMemoryEncodingA…

701d0a2

…ttr (#7748)" This reverts commit e6eb871.

ThomasRaoux approved these changes Aug 5, 2025

View reviewed changes

Re-enable float8 tensor memory

b3f7ce4

peterbell10 commented Aug 5, 2025

View reviewed changes

peterbell10 changed the title ~~Revert "[LAYOUTS] Implement toLinearLayout for TensorMemoryEncodingAttr (#7748)"~~ [IR] Loosen tensor memory encoding checks added in #7748 Aug 5, 2025

peterbell10 enabled auto-merge (squash) August 5, 2025 20:58

peterbell10 force-pushed the pb/pr-chain/revert_layouts_implement_tolinearlayout__9b92 branch from 4cd6aba to b3f7ce4 Compare August 5, 2025 21:24

Make shape errors more informative

c996054

lezcano reviewed Aug 5, 2025

View reviewed changes

lezcano approved these changes Aug 6, 2025

View reviewed changes

lezcano reviewed Aug 6, 2025

View reviewed changes

ThomasRaoux approved these changes Aug 6, 2025

View reviewed changes

Respond to comments

8a45fbc

peterbell10 merged commit bbdbbd1 into main Aug 6, 2025
9 checks passed

peterbell10 deleted the pb/pr-chain/revert_layouts_implement_tolinearlayout__9b92 branch August 6, 2025 02:36

ptillet pushed a commit that referenced this pull request Aug 7, 2025

[IR] Loosen tensor memory encoding checks added in #7748 (#7784)

53b0dd9

	if (auto attr = dyn_cast<triton::nvidia_gpu::TensorMemoryEncodingAttr>(
	memType.getEncoding())) {
	info.blockM = attr.getBlockM();
	info.blockN = attr.getBlockN();
	assert((!attr.getUnpacked() \|\| info.numElementsPer32B <= 2) &&
	"unsupported unpacked layout");

		if (!enc.getUnpacked() && bitwidth > 16) {
		return emitError() << "bitwidth must be <= 16 for packed tensor memory";

[IR] Loosen tensor memory encoding checks added in #7748 #7784

[IR] Loosen tensor memory encoding checks added in #7748 #7784

Uh oh!

Conversation

peterbell10 commented Aug 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Commits in this PR

Uh oh!

peterbell10 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

peterbell10 Aug 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lezcano left a comment

Choose a reason for hiding this comment

Uh oh!

lezcano left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

peterbell10 commented Aug 5, 2025 •

edited

Loading

peterbell10 Aug 6, 2025 •

edited

Loading