[TLE] Add alias support for tle.gpu.alloc to enable shared memory aliasing

[henghengxiedaima]

PR Description

Overview

Add shared-memory alias support to tle.gpu.alloc, allowing multiple descriptors to share the same underlying physical shared memory allocation without allocating new memory.

API

def alloc(
    shape: tuple,
    dtype: tl.dtype,
    layout: Optional[tle.shared_layout] = None,
    scope: tle.scope = tle.smem,
    init_value: Optional[tl.tensor] = None,
    alias: Optional[tle.buffered_tensor] = None,
    alias_offset_bytes: int = 0,
    nv_mma_shared_layout=True,
    _semantic=None,  # internal
) -> tle.buffered_tensor

New parameters:

• alias: An existing shared-memory tle.buffered_tensor whose backing storage will be reused. When set, no new shared memory is allocated.
• alias_offset_bytes: Static byte offset from the source descriptor's view base. Must be a compile-time, non-negative integer.

Lowering Path

tle.gpu.alloc(alias=...)
    │
    ▼  Python builder: create_memdesc_alias()
    │
    ▼  MLIR: tle.memdesc_alias
    │
    ▼  Shared-memory allocation/view analysis (no new allocation — reuses source)
    │
    ▼  NVIDIA MemDescAliasOpConversion
    │
    ▼  llvm.getelementptr (source shared-memory base + byte offset),
       followed by reconstruction of the result descriptor

Validation Rules & Limitations

• Shared memory only: alias mode currently supports tle.smem only. Both source and alias must be shared-memory buffers.
• No init_value: alias and init_value are mutually exclusive.
• Type of alias: alias must be a tle.buffered_tensor.
• Compile-time offset: alias_offset_bytes must be a compile-time, non-negative integer, fit in i32, and be aligned to the result element byte width.
• Static sizes: Source and result memdesc byte sizes must be statically known. The result byte range must fit within the source view.
• Memory space match: Source and result memory spaces must match. A mutable alias cannot be created from an immutable source.
• Lowering: The LLVM lowering is currently registered in the NVIDIA path (third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/).

Test

• Unit tests: python/test/tle/unit/test_tle_alloc_alias.py
• E2E test: python/test/tle/integration/test_tle_alias_e2e.py — validates that alias descriptors observe writes made through the source descriptor, confirming shared physical memory.

[CLAassistant]

All committers have signed the CLA.

[sunnycase]

Thank you for adding shared-memory alias support to tle.gpu.alloc. Could you please address the following items before this is merged?

1. Guard TLE-specific changes outside TLE-owned directories

   Please wrap all TLE-specific native changes outside the TLE directories with the established compile-time guard:

   #ifdef __TLE__
   // TLE-specific code
   #endif

   In particular, please check the TLE includes and MemDescAliasOp handling added under lib/Dialect/TritonGPU/**, as well as the lowering registration under third_party/nvidia/**. Files under third_party/tle/** do not need this guard.

2. Use English comments only

   Please replace the Chinese comments in test_alias_e2e.py with English and ensure that no Chinese code comments remain in this PR.

3. Expand the PR Summary

   Please document the tle.gpu.alloc API and the new alias mode explicitly. The current API definition is:

   def alloc(
       shape: tuple,
       dtype: tl.dtype,
       layout: Optional[tle.shared_layout] = None,
       scope: tle.scope = tle.smem,
       init_value: Optional[tl.tensor] = None,
       alias: Optional[tle.buffered_tensor] = None,
       alias_offset_bytes: int = 0,
       nv_mma_shared_layout=True,
       _semantic=None,  # internal
   ) -> tle.buffered_tensor

   Please also summarize the lowering path, for example:

   tle.gpu.alloc(alias=...) → Python builder create_memdesc_alias → tle.memdesc_alias → shared-memory allocation/view analysis (no new allocation) → NVIDIA MemDescAliasOpConversion → source shared-memory base plus a byte-offset llvm.getelementptr, followed by reconstruction of the result descriptor.

   Finally, please list the current limitations and validation rules, including:

   • alias mode currently supports shared memory only, and alias must be a shared-memory tle.buffered_tensor;
   • alias cannot be combined with init_value;
   • alias_offset_bytes must be a compile-time, non-negative integer, fit in i32, and be aligned to the result element byte width;
   • source and result memdesc byte sizes must be static, and the result byte range must fit within the source view;
   • source and result memory spaces must match, and a mutable alias cannot be created from an immutable source;
   • the LLVM lowering is currently registered in the NVIDIA path.

Thank you!

[zhzhcookie]

The test file cannot be placed here.

[zhzhcookie]

#ifdef __TLE__

[henghengxiedaima]

Thanks for the review! All three items have been addressed — please take another look.

[henghengxiedaima]

Done. Wrapped the tle::MemDescAliasOpreference in#ifdef TLE(and the same forMemDescWGMMAViewOp in OptimizeDotOperands); the TritonGPUToLLVM.cpp one was already guarded.