[TLE] Add tle.gpu.set_layout primitive

[ph0375]

PR Category
Compiler / TLE primitive

Type of Change
New Feature

Description
This PR adds tle.gpu.set_layout, a TLE primitive that lets kernel authors attach an explicit, user-chosen TritonGPU distributed encoding to a tensor value, instead of relying entirely on the compiler's automatic layout inference. The requested layout is treated as a hard constraint and propagated through the surrounding computation, taking precedence over layouts the compiler would otherwise infer.

API
def set_layout(value: tl.tensor, layout: distributed_encoding) -> tl.tensor

layout must be a compile-time distributed_encoding instance. Four concrete encoding types are provided:

①BlockEncoding(size_per_thread, threads_per_warp, warps_per_cta, order, cga_layout=None)
②MmaEncoding(version, warps_per_cta, instr_shape, cga_layout=None)
③DotOperandEncoding(operand_index, parent, k_width)
④SlicedEncoding(dim, parent)

set_layout only accepts block tensors (value.type.is_block()); scalars are rejected. layout must be resolvable at compile time (it is unwrapped via _unwrap_if_constexpr and must implement to_ir()).

Lowering path
1、Python frontend call: tle.gpu.set_layout(value, layout)
2、Pybind builder call: builder.create_tle_gpu_set_layout(...) (python/src/ir.cc)
3、TLE dialect op: tle.gpu.set_layout (Tle_SetLayoutOp — pure, shape/element-type-preserving)
4、Triton→TritonGPU pass: propagateTleEncodingHints marks the requested encoding as "hard" and threads it through scf.for/scf.while, dot operands/accumulators, and other layout-sensitive ops
5、RemoveLayoutConversions::resolveConflicts prefers hard encodings over inferred ones when a value has multiple candidate encodings
6、Coalesce pass skips auto-coalescing for memory ops carrying an explicit TLE memory encoding
7、Final lowering via TleSetLayoutOpPattern: elided if the source is already in the requested layout, otherwise rewritten to a tagged ttg.ConvertLayoutOp

Limitations / validation rules
Only block tensors are supported; scalar values are rejected at the frontend.
layout must be a compile-time constant implementing to_ir(); it cannot be a runtime value.
BlockEncoding/MmaEncoding validate rank consistency and that order is a valid permutation at construction time — malformed layouts fail immediately in Python rather than surfacing as an opaque MLIR verifier error later.
DotOperandEncoding.parent and SlicedEncoding.parent must themselves be distributed_encoding instances (not raw MLIR attributes).
If a value ends up with two different hard (explicitly requested) encodings that don't match, this is a genuine conflict and fails compilation with a diagnostic (mergeTleEncodingInfo in TritonToTritonGPUPass.cpp) rather than silently picking one.
Memory ops (load/store) whose pointer chain carries an explicit encoding are skipped by the coalescing pass, so set_layout on a pointer-producing chain is not silently overridden by auto-coalescing.
All TLE-specific additions outside third_party/tle/** (i.e. in the shared Triton/TritonGPU codebase: Utility.h/.cpp, TritonGPUConversion.cpp, TritonToTritonGPUPass.cpp, Coalesce.cpp, RemoveLayoutConversions.cpp, python/src/ir.cc) are guarded with #ifdef TLE so non-TLE builds are unaffected. Files under third_party/tle/** do not carry this guard, per existing convention.

Tests
Python frontend tests (test_tle.py): BlockEncoding/SlicedEncoding lowering (test_explicit_distributed_encoding_frontend), and MmaEncoding/DotOperandEncoding lowering across a chained-dot kernel (test_explicit_dot_encoding_frontend).
MLIR filecheck tests: a chain of tt.dots sharing a hard-encoded accumulator produces no spurious ttg.convert_layout (test_tle_explicit_dot_encoding_propagation.mlir); an explicitly-encoded store is left un-coalesced with the requested encoding preserved as an attribute (test_tle_explicit_memory_encoding_coalesce.mlir).

Issue
None

Progress
Change is properly reviewed (1 reviewer required, 2 recommended).
Change is responded to an issue.
Change is fully covered by a UT.

Performance
This PR's own tests target correctness of the primitive (no spurious conversions, no dropped explicit encodings) rather than a specific kernel's speedup.