Skip to content

[TLE][BACKEND] Add AMD tile and cumsum lowering - #939

Open
zihaomu wants to merge 3 commits into
flagos-ai:mainfrom
zihaomu:feat/tle-lite-amd-extract-insert-cumsum
Open

[TLE][BACKEND] Add AMD tile and cumsum lowering#939
zihaomu wants to merge 3 commits into
flagos-ai:mainfrom
zihaomu:feat/tle-lite-amd-extract-insert-cumsum

Conversation

@zihaomu

@zihaomu zihaomu commented Aug 8, 2026

Copy link
Copy Markdown

What

  • Add AMD LLVM lowering for tle.extract_tile, tle.insert_tile, and tle.exclusive_cumsum.
  • Register AMD shared-memory scratch allocation for the three operations.
  • Select thread-id, barrier, and shuffle lowering through target capabilities instead of backend identity.
  • Keep backend-owned TargetInfoBase copies source-compatible with the shared capability API.
  • Report residual unsupported TLE operations with their names and AMD target architecture.

Why

The shared TLE lowering previously emitted NVIDIA-specific thread-id, barrier, and shuffle operations unless the target identified as HCU. AMD therefore could not lower these portable tile operations even though its target interface already provides the required primitives.

This change lets the AMD backend reuse the shared lowering while preserving existing behavior for NVIDIA, HCU, Iluvatar, Metax, and MThreads.

Testing

  • gfx1201 extract/insert/cumsum tests: 13 passed, 4 skipped
    • Skips are NVIDIA PTX, HCU ISA, and the follow-up AMD local-pointer lowering cases.
  • AMD residual-op diagnostic tests: 2 passed
  • AMD triton-opt + FileCheck regression for extract_tile, insert_tile, and exclusive_cumsum: passed
  • Ninja build and libtriton.so link: passed
  • Ruff 0.9.1, YAPF 0.43.0, clang-format 19.1.6, and git diff --check: passed
  • merge-tree against current upstream main: conflict-free

Cross-backend validation

An HCU-equivalent wheel build compiled the changed TargetInfo.cpp and TritonGPUToLLVM.cpp. It later stopped in the unchanged TargetUtils.cpp because the local LLVM package does not contain HCU-specific GPU enums (GK_GFX928, GK_GFX936, and related targets). The upstream HCU/Iluvatar/Metax/MThreads CI jobs remain the final cross-backend gate.

zihaomu added 3 commits August 8, 2026 17:57
Route thread-id, barrier, and shuffle lowering through target behavior capabilities instead of backend identity. Keep backend-owned TargetInfoBase copies source-compatible and preserve the existing HCU behavior.
Register shared-memory scratch allocation and AMD LLVM patterns for extract_tile, insert_tile, and exclusive_cumsum. Report any unsupported TLE operations that remain after AMD lowering with their operation names and target architecture.
Exercise dynamic and static tile operations plus the cumsum correctness matrix on AMD. Add an AMD FileCheck lowering regression and unit coverage for residual unsupported TLE diagnostics.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant