[µTLX] fix h100 persistent matmul

[plotfi]

This makes TLX H100 persistent matmul for µTLX, but there was quite a lot of vibe coding involved so this must be audited somewhat.

This makes the following example notebook work:

https://colab.research.google.com/drive/1zANW7SP8dG9I7SXvWkH_9pYFRZYPAu_y?usp=sharing

Claude Summary of what it did:

Legality Check Fixes for ConvertTritonToTritonGPU

Problem

The uTLX plugin creates TTNG ops (WarpGroupDotOp, WarpGroupDotWaitOp) and TLX ops (RequireLayoutOp, ReleaseLayoutOp) during the builder phase with unencoded tensor types. The ConvertTritonToTritonGPU pass must convert these to use encoded types (e.g., #blocked, #nvidia_mma), but several issues prevented this.

Fix 1: Custom patterns for TTNG ops with InferTypeOpInterface

uTLXConversionPatterns.cpp lines 128-158 — WarpGroupDotPattern and WarpGroupDotWaitPattern
GenericOpPattern calls replaceOpWithNewOp<Op>(op, retTypes, adaptor.getOperands(), op->getAttrs()), passing explicit result types. But WarpGroupDotOp has TypesMatchWith and WarpGroupDotWaitOp has AllTypesMatch — both implement InferTypeOpInterface, meaning result types are inferred from operands, not accepted as parameters. The builder ignores the explicit types, so the result stays unencoded.
Fix: Custom patterns that call the op's named-parameter constructor (no explicit result types), letting type inference derive encoded results from converted operands:

• WarpGroupDotPattern: passes adaptor.getA(), adaptor.getB(), adaptor.getC(), adaptor.getUseC() + attribute args
• WarpGroupDotWaitPattern: passes adaptor.getInputs(), op.getPendings()

Fix 2: Dynamic legality for TTNG ops

uTLXConversionPatterns.cpp lines 966-972
The TTNG dialect is blanket-legal (addLegalDialect<TritonNvidiaGPUDialect>()), so the conversion framework skips TTNG ops even when they have unencoded tensors.
Fix: Override with addDynamicallyLegalOp for WarpGroupDotOp and WarpGroupDotWaitOp using allTensorsEncoded — these ops are only legal when all their tensor operands/results have encodings.

Fix 3: RequireLayoutOp/ReleaseLayoutOp to ttg.convert_layout

uTLXConversionPatterns.cpp lines 86-122 — RequireLayoutToConvertPattern and ReleaseLayoutToConvertPattern
These TLX ops change tensor encodings (e.g., #blocked to #nvidia_mma). They sit between WarpGroupDotOp and WarpGroupDotWaitOp in the SSA chain, creating unencoded/encoded boundaries that cause source materializations.
Using GenericOpPattern would create new RequireLayoutOp/ReleaseLayoutOp instances (which are also illegal, causing an infinite conversion loop). Instead, these patterns replace them with ttg.convert_layout, which is legal in the TritonGPU dialect.

Fix 4: Dynamic legality for TLX ops

uTLXConversionPatterns.cpp lines 976-986
TLX RequireLayout/ReleaseLayout ops that operate on MemDescType (not RankedTensorType) should pass through unchanged. Only tensor-typed instances need conversion.
Fix: addDynamicallyLegalOp returns false (illegal) when any operand or result is RankedTensorType, true otherwise.

Fix 5: Post-conversion cleanup walk

uTLXConversionPatterns.cpp lines 1040-1060
If any tensor-typed RequireLayoutOp/ReleaseLayoutOp survive applyPartialConversion (e.g., because the TLX PropagateLayout pass is not in the pipeline), they would reach LLVM translation and cause unrealized_conversion_cast failures.
Fix: A post-conversion mod.walk lowers surviving tensor-typed instances to ttg.convert_layout as a safety net.

Fix 6: ensureModuleAttrs in NewOps.cpp

NewOps.cpp lines 17-38 — ensureModuleAttrs() helper
NvidiaMmaEncoding layout verifiers require ttg.num-warps, ttg.threads-per-warp, and ttg.num-ctas module attributes. These weren't being set early enough during the builder phase.
Fix: ensureModuleAttrs() eagerly sets these attributes on the parent ModuleOp when createRequireNvMmaLayout is called.

Pattern registration

uTLXConversionPatterns.cpp lines 1020-1023 — All four custom patterns are registered:

patterns.add<RequireLayoutToConvertPattern,
             ReleaseLayoutToConvertPattern,
             WarpGroupDotPattern,
             WarpGroupDotWaitPattern>(typeConverter, context);

Remaining issue

The failed to legalize unresolved materialization error may still persist for tlx.require_layout inside SCF for-loop bodies. The MLIR conversion framework inserts source materializations when it converts SCF block args from unencoded to #blocked, but the RequireLayoutOp that uses those block args may not get its pattern applied at the right time. This is a framework-level timing issue with applyPartialConversion.