[Mlir-commits] [clang] [llvm] [mlir] [AMDGPU] Add sched_vgpr/sgpr/agpr intrinsics for GPR register hints (PR #233773)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 20 03:19:31 PDT 2026
https://github.com/jli-melchior created https://github.com/llvm/llvm-project/pull/233773
## Motivation
In performance-critical GPU kernels (e.g. GEMM, attention, MoE), the exact physical register assignment can have a measurable impact on performance.
This PR adds intrinsics that carry an exact register-number hint through the compilation pipeline, giving explicit control over physical register assignment for all three GPR classes (VGPR, SGPR, AGPR). The hint is advisory by default, and can be made strict when correctness depends on exact placement.
In practice, using these intrinsics to control register placement in a production [MoE kernel](https://github.com/ROCm/aiter/blob/f9e2d3743658d549a4edd0d04959a0bd3fee9cf2/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py) yields over 10% performance improvement.
This is the GPR-only counterpart to the bank-level scheduling in https://github.com/llvm/llvm-project/pull/227953.
## Summary
Add `llvm.amdgcn.sched.{vgpr,sgpr,agpr}` intrinsics and supporting infrastructure to allow explicit control over which physical GPR a value is allocated to. This is the GPR-only counterpart to the bank scheduling in https://github.com/llvm/llvm-project/pull/227953.
- **LLVM intrinsics**: `llvm.amdgcn.sched.{vgpr,sgpr,agpr}(value, regnum, strict)` — request that `value` be allocated to the specified register number. When `strict=1`, the hint is treated as a hard constraint.
- **Clang builtins**: `__builtin_amdgcn_sched_{vgpr,sgpr,agpr}` with Sema validation of register number bounds.
- **MLIR ROCDL ops**: `rocdl.sched.{vgpr,sgpr,agpr}` with strict/soft mode attributes.
- **ISel**: Intrinsics lower to `V_SCHEDULE_VGPR_B*`, `S_SCHEDULE_SGPR_B*`, `A_SCHEDULE_AGPR_B*` pseudos (B32/B64/B128/B256/B512).
- **AMDGPUSchedGPR pass**: Pre-RA pass that reads the regnum/strict operands, sets `AMDGPURI::{VGPR,SGPR,AGPR}RegHint` allocation hints on vregs, propagates hints through COPY/REG_SEQUENCE chains, and replaces pseudos with plain COPYs.
- **SIRegisterInfo**: `getRegAllocationHints` handles the new hint types, reordering candidates so the allocator prefers the exact physical register. Strict hints return true to make the preference a hard constraint.
## Test plan
- [x] `clang/test/CodeGenOpenCL/builtins-amdgcn-sched-gpr.cl` — Clang codegen
- [x] `llvm/test/CodeGen/AMDGPU/sched-gpr.ll` — end-to-end: v42/v100/v512 appear in output
- [x] `llvm/test/CodeGen/AMDGPU/sched-gpr-hints.mir` — MIR: pseudos lowered to COPY
- [x] `llvm/test/CodeGen/AMDGPU/llc-pipeline.ll` / `llc-pipeline-npm.ll` — pipeline ordering
- [x] `mlir/test/Dialect/LLVMIR/rocdl.mlir` — MLIR parse/print
>From 41303030ff50c6962839df03277ae784cb3adf8b Mon Sep 17 00:00:00 2001
From: Jun Li <jli10004 at amd.com>
Date: Tue, 20 Oct 2026 09:47:44 +0000
Subject: [PATCH 1/3] [AMDGPU] Add sched_vgpr/sgpr/agpr intrinsics, builtins,
and MLIR ops
Add exact GPR register scheduling hints for VGPR, SGPR, and AGPR:
- Add llvm.amdgcn.sched.{vgpr,sgpr,agpr} intrinsics with (value, regnum, strict) signature
- Add __builtin_amdgcn_sched_{vgpr,sgpr,agpr} clang builtins with Sema checks
- Add ISel lowering to V_SCHEDULE_VGPR_B*, S_SCHEDULE_SGPR_B*, A_SCHEDULE_AGPR_B* pseudos (32-512 bit)
- Add VGPRRegHint/SGPRRegHint/AGPRRegHint register allocation hint types in SIRegisterInfo
- Add MLIR ROCDL rocdl.sched.{vgpr,sgpr,agpr} ops with strict/soft modes
- Add CodeGen and MIR tests
Depends on #227953 (sched_bank) which provides the AMDGPUSchedBank pass
infrastructure that lowers these pseudos into register allocation hints.
Co-Authored-By: Claude Opus 4 <noreply at anthropic.com>
---
clang/include/clang/Basic/BuiltinsAMDGPU.td | 5 +
clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp | 26 +++++
clang/lib/Sema/SemaAMDGPU.cpp | 28 +++++
.../builtins-amdgcn-sched-gpr.cl | 68 ++++++++++++
llvm/include/llvm/IR/IntrinsicsAMDGPU.td | 22 ++++
llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp | 57 ++++++++++
llvm/lib/Target/AMDGPU/SIInstructions.td | 48 +++++++++
llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp | 21 ++++
llvm/lib/Target/AMDGPU/SIRegisterInfo.h | 11 +-
llvm/test/CodeGen/AMDGPU/sched-gpr-hints.mir | 66 ++++++++++++
llvm/test/CodeGen/AMDGPU/sched-gpr.ll | 101 ++++++++++++++++++
mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 73 +++++++++++++
mlir/test/Dialect/LLVMIR/rocdl.mlir | 24 +++++
mlir/test/Target/LLVMIR/rocdl.mlir | 29 +++++
14 files changed, 578 insertions(+), 1 deletion(-)
create mode 100644 clang/test/CodeGenOpenCL/builtins-amdgcn-sched-gpr.cl
create mode 100644 llvm/test/CodeGen/AMDGPU/sched-gpr-hints.mir
create mode 100644 llvm/test/CodeGen/AMDGPU/sched-gpr.ll
diff --git a/clang/include/clang/Basic/BuiltinsAMDGPU.td b/clang/include/clang/Basic/BuiltinsAMDGPU.td
index 34a1bd2d301a37..466dfc8e99e758 100644
--- a/clang/include/clang/Basic/BuiltinsAMDGPU.td
+++ b/clang/include/clang/Basic/BuiltinsAMDGPU.td
@@ -1012,6 +1012,11 @@ def __builtin_amdgcn_cvt_sr_f16_f32 : AMDGPUBuiltin<"_ExtVector<2, _Float16>(_Ex
//===----------------------------------------------------------------------===//
def __builtin_amdgcn_s_cluster_barrier : AMDGPUBuiltin<"void()", [], "gfx1250-insts">;
+// Exact GPR register scheduling hints.
+def __builtin_amdgcn_sched_vgpr : AMDGPUBuiltin<"int(int, _Constant int, _Constant bool)", [Const, CustomTypeChecking]>;
+def __builtin_amdgcn_sched_sgpr : AMDGPUBuiltin<"int(int, _Constant int, _Constant bool)", [Const, CustomTypeChecking]>;
+def __builtin_amdgcn_sched_agpr : AMDGPUBuiltin<"int(int, _Constant int, _Constant bool)", [Const, CustomTypeChecking], "mai-insts">;
+
def __builtin_amdgcn_flat_prefetch : AMDGPUBuiltin<"void(void const address_space<0> *, _Constant int)", [Const], "vmem-pref-insts">;
def __builtin_amdgcn_global_prefetch : AMDGPUBuiltin<"void(void const address_space<1> *, _Constant int)", [Const], "vmem-pref-insts">;
diff --git a/clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp b/clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp
index 2f29e9557ba56f..c27a8e8f2a9fcf 100644
--- a/clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp
+++ b/clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp
@@ -2283,6 +2283,32 @@ Value *CodeGenFunction::EmitAMDGPUBuiltinExpr(unsigned BuiltinID,
case AMDGPU::BI__builtin_amdgcn_permlane_xor:
return emitBuiltinWithOneOverloadedType<3>(*this, E,
Intrinsic::amdgcn_permlane_xor);
+ case AMDGPU::BI__builtin_amdgcn_sched_vgpr:
+ case AMDGPU::BI__builtin_amdgcn_sched_sgpr:
+ case AMDGPU::BI__builtin_amdgcn_sched_agpr: {
+ llvm::Value *Val = EmitScalarExpr(E->getArg(0));
+ llvm::Value *RegNum = EmitScalarExpr(E->getArg(1));
+ llvm::Value *Strict;
+ if (E->getNumArgs() > 2) {
+ Strict = EmitScalarExpr(E->getArg(2));
+ Strict = Builder.CreateZExt(Strict, RegNum->getType());
+ } else {
+ // Default: strict
+ Strict = Builder.getInt32(1);
+ }
+ Intrinsic::ID IID;
+ switch (BuiltinID) {
+ case AMDGPU::BI__builtin_amdgcn_sched_vgpr:
+ IID = Intrinsic::amdgcn_sched_vgpr; break;
+ case AMDGPU::BI__builtin_amdgcn_sched_sgpr:
+ IID = Intrinsic::amdgcn_sched_sgpr; break;
+ case AMDGPU::BI__builtin_amdgcn_sched_agpr:
+ IID = Intrinsic::amdgcn_sched_agpr; break;
+ default: llvm_unreachable("unexpected builtin");
+ }
+ llvm::Function *F = CGM.getIntrinsic(IID, Val->getType());
+ return Builder.CreateCall(F, {Val, RegNum, Strict});
+ }
default:
return nullptr;
}
diff --git a/clang/lib/Sema/SemaAMDGPU.cpp b/clang/lib/Sema/SemaAMDGPU.cpp
index b1f6fa33649bd1..556526394d2872 100644
--- a/clang/lib/Sema/SemaAMDGPU.cpp
+++ b/clang/lib/Sema/SemaAMDGPU.cpp
@@ -465,6 +465,34 @@ bool SemaAMDGPU::CheckAMDGCNBuiltinFunctionCall(const TargetInfo &TI,
/*High=*/0) ||
SemaRef.BuiltinConstantArgRange(TheCall, /*ArgNum=*/2, /*Low=*/0,
/*High=*/0);
+ case AMDGPU::BI__builtin_amdgcn_sched_vgpr:
+ case AMDGPU::BI__builtin_amdgcn_sched_sgpr:
+ case AMDGPU::BI__builtin_amdgcn_sched_agpr: {
+ if (SemaRef.checkArgCountRange(TheCall, 2, 3))
+ return true;
+ Expr *DataArg = TheCall->getArg(0);
+ QualType DataTy = DataArg->getType();
+ if (DataTy->isAnyComplexType() ||
+ !(DataTy->isArithmeticType() ||
+ (DataTy->isVectorType() &&
+ DataTy->castAs<VectorType>()
+ ->getElementType()
+ ->isArithmeticType()))) {
+ SemaRef.Diag(DataArg->getBeginLoc(),
+ diag::err_typecheck_cond_expect_int_float)
+ << DataTy << DataArg->getSourceRange();
+ return true;
+ }
+ TheCall->setType(DataTy);
+ int MaxReg;
+ switch (BuiltinID) {
+ case AMDGPU::BI__builtin_amdgcn_sched_vgpr: MaxReg = 1023; break;
+ case AMDGPU::BI__builtin_amdgcn_sched_sgpr: MaxReg = 105; break;
+ case AMDGPU::BI__builtin_amdgcn_sched_agpr: MaxReg = 255; break;
+ default: llvm_unreachable("unexpected builtin");
+ }
+ return SemaRef.BuiltinConstantArgRange(TheCall, 1, 0, MaxReg);
+ }
default:
return false;
}
diff --git a/clang/test/CodeGenOpenCL/builtins-amdgcn-sched-gpr.cl b/clang/test/CodeGenOpenCL/builtins-amdgcn-sched-gpr.cl
new file mode 100644
index 00000000000000..2d889f27f38a4e
--- /dev/null
+++ b/clang/test/CodeGenOpenCL/builtins-amdgcn-sched-gpr.cl
@@ -0,0 +1,68 @@
+// NOTE: Assertions have been autogenerated by utils/update_cc_test_checks.py
+// RUN: %clang_cc1 -cl-std=CL2.0 -O0 -triple amdgcn-unknown-unknown -target-cpu gfx1250 -emit-llvm -o - %s | FileCheck %s
+// REQUIRES: amdgpu-registered-target
+
+typedef int __attribute__((ext_vector_type(2))) int2;
+typedef int __attribute__((ext_vector_type(4))) int4;
+typedef int __attribute__((ext_vector_type(8))) int8;
+typedef float __attribute__((ext_vector_type(2))) float2;
+typedef float __attribute__((ext_vector_type(4))) float4;
+
+// === sched_vgpr tests ===
+
+// CHECK-LABEL: @test_sched_vgpr_i32
+// CHECK: call i32 @llvm.amdgcn.sched.vgpr.i32(i32 %{{.*}}, i32 42, i32 0)
+void test_sched_vgpr_i32(global int *out, int x) {
+ *out = __builtin_amdgcn_sched_vgpr(x, 42, false);
+}
+
+// CHECK-LABEL: @test_sched_vgpr_f32_strict
+// CHECK: call float @llvm.amdgcn.sched.vgpr.f32(float %{{.*}}, i32 100, i32 1)
+void test_sched_vgpr_f32_strict(global float *out, float x) {
+ *out = __builtin_amdgcn_sched_vgpr(x, 100, true);
+}
+
+// CHECK-LABEL: @test_sched_vgpr_i64
+// CHECK: call i64 @llvm.amdgcn.sched.vgpr.i64(i64 %{{.*}}, i32 200, i32 0)
+void test_sched_vgpr_i64(global long *out, long x) {
+ *out = __builtin_amdgcn_sched_vgpr(x, 200, false);
+}
+
+// CHECK-LABEL: @test_sched_vgpr_v2i32
+// CHECK: call <2 x i32> @llvm.amdgcn.sched.vgpr.v2i32(<2 x i32> %{{.*}}, i32 50, i32 1)
+void test_sched_vgpr_v2i32(global int2 *out, int2 x) {
+ *out = __builtin_amdgcn_sched_vgpr(x, 50, true);
+}
+
+// CHECK-LABEL: @test_sched_vgpr_v4i32
+// CHECK: call <4 x i32> @llvm.amdgcn.sched.vgpr.v4i32(<4 x i32> %{{.*}}, i32 512, i32 1)
+void test_sched_vgpr_v4i32(global int4 *out, int4 x) {
+ *out = __builtin_amdgcn_sched_vgpr(x, 512, true);
+}
+
+// CHECK-LABEL: @test_sched_vgpr_default_strict
+// CHECK: call i32 @llvm.amdgcn.sched.vgpr.i32(i32 %{{.*}}, i32 42, i32 1)
+void test_sched_vgpr_default_strict(global int *out, int x) {
+ // 2 args: strict=true by default
+ *out = __builtin_amdgcn_sched_vgpr(x, 42);
+}
+
+// === sched_sgpr tests ===
+
+// CHECK-LABEL: @test_sched_sgpr_i32
+// CHECK: call i32 @llvm.amdgcn.sched.sgpr.i32(i32 %{{.*}}, i32 10, i32 0)
+void test_sched_sgpr_i32(global int *out, int x) {
+ *out = __builtin_amdgcn_sched_sgpr(x, 10, false);
+}
+
+// CHECK-LABEL: @test_sched_sgpr_i32_strict
+// CHECK: call i32 @llvm.amdgcn.sched.sgpr.i32(i32 %{{.*}}, i32 50, i32 1)
+void test_sched_sgpr_i32_strict(global int *out, int x) {
+ *out = __builtin_amdgcn_sched_sgpr(x, 50, true);
+}
+
+// CHECK-LABEL: @test_sched_sgpr_default_strict
+// CHECK: call i32 @llvm.amdgcn.sched.sgpr.i32(i32 %{{.*}}, i32 20, i32 1)
+void test_sched_sgpr_default_strict(global int *out, int x) {
+ *out = __builtin_amdgcn_sched_sgpr(x, 20);
+}
diff --git a/llvm/include/llvm/IR/IntrinsicsAMDGPU.td b/llvm/include/llvm/IR/IntrinsicsAMDGPU.td
index ed80bc4628b6fa..2400bc1d6b5e4b 100644
--- a/llvm/include/llvm/IR/IntrinsicsAMDGPU.td
+++ b/llvm/include/llvm/IR/IntrinsicsAMDGPU.td
@@ -4478,4 +4478,26 @@ def int_amdgcn_global_load_monitor_b128 : AMDGPULoadMonitor<global_ptr_ty>;
/// incoming block.
def int_amdgcn_dead: DefaultAttrsIntrinsic<[llvm_any_ty], [],
[IntrNoMem]>;
+
+// Hint the register allocator to place the result in a specific VGPR/SGPR/AGPR
+// register number. These are identity functions at the IR level; the register
+// number and strict flag arguments are consumed by the pre-RA AMDGPUSchedBank
+// pass which attaches a register allocation hint.
+// regnum: the target HW register index (e.g., 42 means v42/s42/a42).
+// strict: 1 = allocator must obey if possible, 0 = advisory hint.
+// For wide values (64/128/256 bit), regnum is the starting register index.
+def int_amdgcn_sched_vgpr : DefaultAttrsIntrinsic<
+ [llvm_any_ty], [LLVMMatchType<0>, llvm_i32_ty, llvm_i32_ty],
+ [IntrNoMem, IntrSpeculatable, IntrWillReturn,
+ ImmArg<ArgIndex<1>>, ImmArg<ArgIndex<2>>]>;
+
+def int_amdgcn_sched_sgpr : DefaultAttrsIntrinsic<
+ [llvm_any_ty], [LLVMMatchType<0>, llvm_i32_ty, llvm_i32_ty],
+ [IntrNoMem, IntrSpeculatable, IntrWillReturn,
+ ImmArg<ArgIndex<1>>, ImmArg<ArgIndex<2>>]>;
+
+def int_amdgcn_sched_agpr : DefaultAttrsIntrinsic<
+ [llvm_any_ty], [LLVMMatchType<0>, llvm_i32_ty, llvm_i32_ty],
+ [IntrNoMem, IntrSpeculatable, IntrWillReturn,
+ ImmArg<ArgIndex<1>>, ImmArg<ArgIndex<2>>]>;
}
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp b/llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp
index dc2abb60567ffa..a6c20d61be838f 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp
@@ -3353,6 +3353,63 @@ void AMDGPUDAGToDAGISel::SelectINTRINSIC_WO_CHAIN(SDNode *N) {
case Intrinsic::amdgcn_interp_p1_f16:
SelectInterpP1F16(N);
return;
+ case Intrinsic::amdgcn_sched_vgpr:
+ case Intrinsic::amdgcn_sched_sgpr:
+ case Intrinsic::amdgcn_sched_agpr: {
+ SDValue Src = N->getOperand(1);
+ auto *RegNumC = dyn_cast<ConstantSDNode>(N->getOperand(2));
+ auto *StrictC = dyn_cast<ConstantSDNode>(N->getOperand(3));
+ if (!RegNumC || !StrictC) {
+ SelectCode(N);
+ return;
+ }
+ unsigned Bits = N->getValueType(0).getSizeInBits();
+ unsigned PseudoOpc;
+ if (IntrID == Intrinsic::amdgcn_sched_vgpr) {
+ switch (Bits) {
+ case 32: PseudoOpc = AMDGPU::V_SCHEDULE_VGPR_B32; break;
+ case 64: PseudoOpc = AMDGPU::V_SCHEDULE_VGPR_B64; break;
+ case 128: PseudoOpc = AMDGPU::V_SCHEDULE_VGPR_B128; break;
+ case 256: PseudoOpc = AMDGPU::V_SCHEDULE_VGPR_B256; break;
+ case 512: PseudoOpc = AMDGPU::V_SCHEDULE_VGPR_B512; break;
+ default:
+ ReplaceUses(SDValue(N, 0), Src);
+ CurDAG->RemoveDeadNode(N);
+ return;
+ }
+ } else if (IntrID == Intrinsic::amdgcn_sched_sgpr) {
+ switch (Bits) {
+ case 32: PseudoOpc = AMDGPU::S_SCHEDULE_SGPR_B32; break;
+ case 64: PseudoOpc = AMDGPU::S_SCHEDULE_SGPR_B64; break;
+ case 128: PseudoOpc = AMDGPU::S_SCHEDULE_SGPR_B128; break;
+ case 256: PseudoOpc = AMDGPU::S_SCHEDULE_SGPR_B256; break;
+ case 512: PseudoOpc = AMDGPU::S_SCHEDULE_SGPR_B512; break;
+ default:
+ ReplaceUses(SDValue(N, 0), Src);
+ CurDAG->RemoveDeadNode(N);
+ return;
+ }
+ } else {
+ switch (Bits) {
+ case 32: PseudoOpc = AMDGPU::A_SCHEDULE_AGPR_B32; break;
+ case 64: PseudoOpc = AMDGPU::A_SCHEDULE_AGPR_B64; break;
+ case 128: PseudoOpc = AMDGPU::A_SCHEDULE_AGPR_B128; break;
+ case 256: PseudoOpc = AMDGPU::A_SCHEDULE_AGPR_B256; break;
+ case 512: PseudoOpc = AMDGPU::A_SCHEDULE_AGPR_B512; break;
+ default:
+ ReplaceUses(SDValue(N, 0), Src);
+ CurDAG->RemoveDeadNode(N);
+ return;
+ }
+ }
+ SDValue RegNumImm = CurDAG->getTargetConstant(RegNumC->getZExtValue(),
+ SDLoc(N), MVT::i32);
+ SDValue StrictImm = CurDAG->getTargetConstant(StrictC->getZExtValue(),
+ SDLoc(N), MVT::i32);
+ CurDAG->SelectNodeTo(N, PseudoOpc, N->getVTList(),
+ {Src, RegNumImm, StrictImm});
+ return;
+ }
case Intrinsic::amdgcn_permlane16_swap:
case Intrinsic::amdgcn_permlane32_swap: {
if ((IntrID == Intrinsic::amdgcn_permlane16_swap &&
diff --git a/llvm/lib/Target/AMDGPU/SIInstructions.td b/llvm/lib/Target/AMDGPU/SIInstructions.td
index c9e0e4c79f8ccd..af21dd0a6483e6 100644
--- a/llvm/lib/Target/AMDGPU/SIInstructions.td
+++ b/llvm/lib/Target/AMDGPU/SIInstructions.td
@@ -129,6 +129,23 @@ def V_CNDMASK_B64_PSEUDO : VOP3Common <(outs VReg_64:$vdst),
let usesCustomInserter = 1;
}
+// Exact VGPR register hint pseudos. Identity copies carrying a target VGPR
+// register number and strict flag. The pre-RA AMDGPUSchedBank pass reads
+// these operands, sets a register allocation hint, and replaces with COPY.
+let isAsCheapAsAMove = 1, isReMaterializable = 1, hasSideEffects = 0,
+ mayLoad = 0, mayStore = 0, Size = 0 in {
+ def V_SCHEDULE_VGPR_B32 : VPseudoInstSI<(outs VGPR_32:$vdst),
+ (ins VGPR_32:$src, i32imm:$regnum, i32imm:$strict)>;
+ def V_SCHEDULE_VGPR_B64 : VPseudoInstSI<(outs VReg_64:$vdst),
+ (ins VReg_64:$src, i32imm:$regnum, i32imm:$strict)>;
+ def V_SCHEDULE_VGPR_B128 : VPseudoInstSI<(outs VReg_128:$vdst),
+ (ins VReg_128:$src, i32imm:$regnum, i32imm:$strict)>;
+ def V_SCHEDULE_VGPR_B256 : VPseudoInstSI<(outs VReg_256:$vdst),
+ (ins VReg_256:$src, i32imm:$regnum, i32imm:$strict)>;
+ def V_SCHEDULE_VGPR_B512 : VPseudoInstSI<(outs VReg_512:$vdst),
+ (ins VReg_512:$src, i32imm:$regnum, i32imm:$strict)>;
+}
+
// 64-bit vector move instruction. This is mainly used by the
// SIFoldOperands pass to enable folding of inline immediates.
def V_MOV_B64_PSEUDO : VPseudoInstSI <(outs VReg_64:$vdst),
@@ -217,6 +234,37 @@ def STRICT_WQM : PseudoInstSI <(outs unknown:$vdst), (ins unknown:$src0)>;
} // End let hasSideEffects = 0, mayLoad = 0, mayStore = 0, Uses = [EXEC]
+// Exact SGPR register hint pseudos. Placed outside the Uses = [EXEC] block
+// because SGPR operations do not use EXEC.
+let isAsCheapAsAMove = 1, isReMaterializable = 1, hasSideEffects = 0,
+ mayLoad = 0, mayStore = 0, Size = 0 in {
+ def S_SCHEDULE_SGPR_B32 : SPseudoInstSI<(outs SGPR_32:$sdst),
+ (ins SGPR_32:$src, i32imm:$regnum, i32imm:$strict)>;
+ def S_SCHEDULE_SGPR_B64 : SPseudoInstSI<(outs SGPR_64:$sdst),
+ (ins SGPR_64:$src, i32imm:$regnum, i32imm:$strict)>;
+ def S_SCHEDULE_SGPR_B128 : SPseudoInstSI<(outs SGPR_128:$sdst),
+ (ins SGPR_128:$src, i32imm:$regnum, i32imm:$strict)>;
+ def S_SCHEDULE_SGPR_B256 : SPseudoInstSI<(outs SGPR_256:$sdst),
+ (ins SGPR_256:$src, i32imm:$regnum, i32imm:$strict)>;
+ def S_SCHEDULE_SGPR_B512 : SPseudoInstSI<(outs SGPR_512:$sdst),
+ (ins SGPR_512:$src, i32imm:$regnum, i32imm:$strict)>;
+}
+
+// Exact AGPR register hint pseudos. Only valid on targets with MAI instructions.
+let isAsCheapAsAMove = 1, isReMaterializable = 1, hasSideEffects = 0,
+ mayLoad = 0, mayStore = 0, Size = 0 in {
+ def A_SCHEDULE_AGPR_B32 : VPseudoInstSI<(outs AGPR_32:$vdst),
+ (ins AGPR_32:$src, i32imm:$regnum, i32imm:$strict)>;
+ def A_SCHEDULE_AGPR_B64 : VPseudoInstSI<(outs AReg_64:$vdst),
+ (ins AReg_64:$src, i32imm:$regnum, i32imm:$strict)>;
+ def A_SCHEDULE_AGPR_B128 : VPseudoInstSI<(outs AReg_128:$vdst),
+ (ins AReg_128:$src, i32imm:$regnum, i32imm:$strict)>;
+ def A_SCHEDULE_AGPR_B256 : VPseudoInstSI<(outs AReg_256:$vdst),
+ (ins AReg_256:$src, i32imm:$regnum, i32imm:$strict)>;
+ def A_SCHEDULE_AGPR_B512 : VPseudoInstSI<(outs AReg_512:$vdst),
+ (ins AReg_512:$src, i32imm:$regnum, i32imm:$strict)>;
+}
+
def WWM_COPY : SPseudoInstSI <
(outs unknown:$dst), (ins unknown:$src)> {
let hasSideEffects = 0;
diff --git a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
index 0989263433f41c..0bc47414f52232 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
@@ -55,6 +55,7 @@ static cl::opt<unsigned> StressSGPRLimit(
"amdgpu-stress-sgpr", cl::Hidden, cl::init(0),
cl::desc("Limit SGPRs to N registers by reserving the rest"));
+
std::array<std::vector<int16_t>, 32> SIRegisterInfo::RegSplitParts;
std::array<std::array<uint16_t, 32>, 9> SIRegisterInfo::SubRegFromChannelTable;
@@ -4329,6 +4330,26 @@ bool SIRegisterInfo::getRegAllocationHints(Register VirtReg,
}
return false;
}
+ case AMDGPURI::VGPRRegHint:
+ case AMDGPURI::StrictVGPRRegHint:
+ case AMDGPURI::SGPRRegHint:
+ case AMDGPURI::StrictSGPRRegHint:
+ case AMDGPURI::AGPRRegHint:
+ case AMDGPURI::StrictAGPRRegHint: {
+ // Exact register preference from llvm.amdgcn.sched.{vgpr,sgpr,agpr}.
+ // The target HW register index is carried in Hint.second. Prefer the
+ // physreg whose HW index matches.
+ unsigned TargetIdx = Hint.second;
+ for (MCPhysReg PhysReg : Order) {
+ if (MRI.isReserved(PhysReg))
+ continue;
+ if (getHWRegIndex(PhysReg) == TargetIdx)
+ Hints.insert(PhysReg);
+ }
+ return Hint.first == AMDGPURI::StrictVGPRRegHint ||
+ Hint.first == AMDGPURI::StrictSGPRRegHint ||
+ Hint.first == AMDGPURI::StrictAGPRRegHint;
+ }
default:
return TargetRegisterInfo::getRegAllocationHints(VirtReg, Order, Hints, MF,
VRM);
diff --git a/llvm/lib/Target/AMDGPU/SIRegisterInfo.h b/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
index b70adfdd921012..a1a49f23ad52d4 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
@@ -38,7 +38,16 @@ struct SGPRSpillBuilder;
/// Register allocation hint types. Helps eliminate unneeded COPY with True16
namespace AMDGPURI {
-enum { Size16 = 1, Size32 = 2 };
+enum {
+ Size16 = 1,
+ Size32 = 2,
+ VGPRRegHint = 3,
+ StrictVGPRRegHint = 4,
+ SGPRRegHint = 5,
+ StrictSGPRRegHint = 6,
+ AGPRRegHint = 7,
+ StrictAGPRRegHint = 8,
+};
} // end namespace AMDGPURI
diff --git a/llvm/test/CodeGen/AMDGPU/sched-gpr-hints.mir b/llvm/test/CodeGen/AMDGPU/sched-gpr-hints.mir
new file mode 100644
index 00000000000000..d480eea0ce596b
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/sched-gpr-hints.mir
@@ -0,0 +1,66 @@
+# RUN: llc -mtriple=amdgcn -mcpu=gfx1250 -run-pass=amdgpu-sched-bank -o - %s | FileCheck %s
+
+# The AMDGPUSchedBank pass attaches exact register allocation hints (not
+# printed in MIR) and lowers each schedule pseudo to a plain COPY.
+
+---
+# CHECK-LABEL: name: lower_sched_vgpr_b32
+# CHECK: %1:vgpr_32 = COPY %0
+# CHECK-NOT: V_SCHEDULE_VGPR_B32
+name: lower_sched_vgpr_b32
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $vgpr0
+ %0:vgpr_32 = COPY $vgpr0
+ %1:vgpr_32 = V_SCHEDULE_VGPR_B32 %0, 42, 1, implicit $exec
+ %2:vgpr_32 = V_ADD_U32_e32 1, %1, implicit $exec
+ $vgpr0 = COPY %2
+ SI_RETURN_TO_EPILOG $vgpr0
+...
+
+---
+# CHECK-LABEL: name: lower_sched_vgpr_b64
+# CHECK: %1:vreg_64_align2 = COPY %0
+# CHECK-NOT: V_SCHEDULE_VGPR_B64
+name: lower_sched_vgpr_b64
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $vgpr0_vgpr1
+ %0:vreg_64_align2 = COPY $vgpr0_vgpr1
+ %1:vreg_64_align2 = V_SCHEDULE_VGPR_B64 %0, 100, 0, implicit $exec
+ $vgpr0_vgpr1 = COPY %1
+ SI_RETURN_TO_EPILOG $vgpr0_vgpr1
+...
+
+---
+# CHECK-LABEL: name: lower_sched_sgpr_b32
+# CHECK: %1:sgpr_32 = COPY %0
+# CHECK-NOT: S_SCHEDULE_SGPR_B32
+name: lower_sched_sgpr_b32
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $sgpr0
+ %0:sgpr_32 = COPY $sgpr0
+ %1:sgpr_32 = S_SCHEDULE_SGPR_B32 %0, 10, 1
+ %2:sgpr_32 = S_ADD_I32 %1, 1, implicit-def $scc
+ $sgpr0 = COPY %2
+ SI_RETURN_TO_EPILOG $sgpr0
+...
+
+---
+# CHECK-LABEL: name: lower_sched_sgpr_b64
+# CHECK: %1:sgpr_64 = COPY %0
+# CHECK-NOT: S_SCHEDULE_SGPR_B64
+name: lower_sched_sgpr_b64
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $sgpr0_sgpr1
+ %0:sgpr_64 = COPY $sgpr0_sgpr1
+ %1:sgpr_64 = S_SCHEDULE_SGPR_B64 %0, 20, 0
+ $sgpr0_sgpr1 = COPY %1
+ SI_RETURN_TO_EPILOG $sgpr0_sgpr1
+...
diff --git a/llvm/test/CodeGen/AMDGPU/sched-gpr.ll b/llvm/test/CodeGen/AMDGPU/sched-gpr.ll
new file mode 100644
index 00000000000000..846b642d0b0cde
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/sched-gpr.ll
@@ -0,0 +1,101 @@
+; RUN: llc -mtriple=amdgcn -mcpu=gfx1250 < %s | FileCheck -check-prefix=GCN %s
+
+; Verify that llvm.amdgcn.schedule.{vgpr,sgpr} steer register allocation to
+; the exact requested register number. Kernels are pinned to occupancy 1
+; (waves-per-eu=1,1 with a single-wave workgroup) so the whole 1024-VGPR file
+; is addressable and high-numbered registers are available to the allocator.
+
+declare i32 @llvm.amdgcn.sched.vgpr.i32(i32, i32, i32)
+declare i64 @llvm.amdgcn.sched.vgpr.i64(i64, i32, i32)
+declare <16 x float> @llvm.amdgcn.sched.vgpr.v16f32(<16 x float>, i32, i32)
+
+; --------------------------------------------------------------------------
+; sched_vgpr: strict hint to v42
+
+; GCN-LABEL: {{^}}vgpr42:
+; GCN: v42
+define amdgpu_kernel void @vgpr42(ptr addrspace(1) %out, i32 %x) #0 {
+ %h = call i32 @llvm.amdgcn.sched.vgpr.i32(i32 %x, i32 42, i32 1)
+ %y = add i32 %h, 1
+ store i32 %y, ptr addrspace(1) %out
+ ret void
+}
+
+; --------------------------------------------------------------------------
+; sched_vgpr: strict hint to v100
+
+; GCN-LABEL: {{^}}vgpr100:
+; GCN: v100
+define amdgpu_kernel void @vgpr100(ptr addrspace(1) %out, i32 %x) #0 {
+ %h = call i32 @llvm.amdgcn.sched.vgpr.i32(i32 %x, i32 100, i32 1)
+ %y = add i32 %h, 1
+ store i32 %y, ptr addrspace(1) %out
+ ret void
+}
+
+; --------------------------------------------------------------------------
+; sched_vgpr: strict hint to v0
+
+; GCN-LABEL: {{^}}vgpr0:
+; GCN: v0,
+define amdgpu_kernel void @vgpr0(ptr addrspace(1) %out, i32 %x) #0 {
+ %h = call i32 @llvm.amdgcn.sched.vgpr.i32(i32 %x, i32 0, i32 1)
+ %y = add i32 %h, 1
+ store i32 %y, ptr addrspace(1) %out
+ ret void
+}
+
+; --------------------------------------------------------------------------
+; sched_vgpr: strict hint to v512 (bank 2 territory)
+
+; GCN-LABEL: {{^}}vgpr512:
+; GCN: v{{(512|0 /\*v512\*/)}}
+
+define amdgpu_kernel void @vgpr512(ptr addrspace(1) %out, i32 %x) #0 {
+ %h = call i32 @llvm.amdgcn.sched.vgpr.i32(i32 %x, i32 512, i32 1)
+ %y = add i32 %h, 1
+ store i32 %y, ptr addrspace(1) %out
+ ret void
+}
+
+; --------------------------------------------------------------------------
+; sched_vgpr: 64-bit value, strict hint to v[200:201]
+
+; GCN-LABEL: {{^}}vgpr200_64:
+; GCN: v[200:201]
+define amdgpu_kernel void @vgpr200_64(ptr addrspace(1) %out, i64 %x) #0 {
+ %h = call i64 @llvm.amdgcn.sched.vgpr.i64(i64 %x, i32 200, i32 1)
+ %y = add i64 %h, 1
+ store i64 %y, ptr addrspace(1) %out
+ ret void
+}
+
+; --------------------------------------------------------------------------
+; sched_vgpr: 512-bit value (v16f32), strict hint to v[512:527]
+
+; GCN-LABEL: {{^}}vgpr512_v16f32:
+; GCN: v{{(\[512:513\]|\[0:1\] /\*v\[512:513\]\*/)}}
+
+define amdgpu_kernel void @vgpr512_v16f32(ptr addrspace(1) %out, ptr addrspace(1) %in) #0 {
+ %x = load <16 x float>, ptr addrspace(1) %in
+ %h = call <16 x float> @llvm.amdgcn.sched.vgpr.v16f32(<16 x float> %x, i32 512, i32 1)
+ store <16 x float> %h, ptr addrspace(1) %out
+ ret void
+}
+
+; --------------------------------------------------------------------------
+; Without hint: register allocator uses default low-numbered registers.
+; Verify no high-numbered VGPRs appear.
+
+; GCN-LABEL: {{^}}nohint:
+; GCN-NOT: v42
+; GCN-NOT: v100
+; GCN-NOT: v512
+; GCN: s_endpgm
+define amdgpu_kernel void @nohint(ptr addrspace(1) %out, i32 %x) #0 {
+ %y = add i32 %x, 1
+ store i32 %y, ptr addrspace(1) %out
+ ret void
+}
+
+attributes #0 = { "amdgpu-flat-work-group-size"="32,32" "amdgpu-waves-per-eu"="1,1" }
diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 716fd92e810ace..fe3dfc164264bd 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -813,6 +813,79 @@ def ROCDL_IglpOpt : ROCDL_ConcreteNonMemIntrOp<"iglp.opt", [], 0, [0], ["variant
}];
}
+class ROCDL_ScheduleGPROp<string mnemonic, string intrName, string gprKind>
+ : ROCDL_IntrOp<mnemonic, [], [0],
+ [AllTypesMatch<["res", "value"]>], 1, 0, 0, 0, [], []>,
+ Arguments<(ins LLVM_Type:$value, I32Attr:$regnum,
+ OptionalAttr<UnitAttr>:$soft)> {
+ let results = (outs LLVM_Type:$res);
+ let assemblyFormat = [{
+ $value `,` $regnum (`soft` $soft^)? attr-dict `:` type($value)
+ }];
+ string llvmBuilder = [{
+ auto schedOp = cast<ROCDL::Schedule}] # gprKind # [{Op>(opInst);
+ auto *value = moduleTranslation.lookupValue(schedOp.getValue());
+ int32_t regnum = schedOp.getRegnum();
+ int32_t strict = schedOp.getSoftAttr() ? 0 : 1;
+ auto *fn = llvm::Intrinsic::getOrInsertDeclaration(
+ moduleTranslation.getLLVMModule(),
+ llvm::Intrinsic::amdgcn_sched_}] # intrName # [{,
+ {value->getType()});
+ auto *inst = builder.CreateCall(fn,
+ {value, builder.getInt32(regnum), builder.getInt32(strict)});
+ moduleTranslation.mapValue(opInst.getResult(0), inst);
+ }];
+}
+
+def ROCDL_ScheduleVGPROp : ROCDL_ScheduleGPROp<"sched.vgpr", "vgpr", "VGPR"> {
+ let description = [{
+ VGPR register scheduling hint. Identity at runtime; guides the register
+ allocator to place `value` in the specified VGPR register number.
+
+ For wide values (64/128/256 bit), `regnum` is the starting register index.
+ Default is strict (allocator must obey). Pass `soft` for an advisory hint.
+
+ Example:
+ ```mlir
+ %0 = rocdl.sched.vgpr %val, 42 : f32 // strict: allocate to v42
+ %1 = rocdl.sched.vgpr %val, 10 soft : f32 // soft: prefer v10
+ ```
+ }];
+}
+
+def ROCDL_ScheduleSGPROp : ROCDL_ScheduleGPROp<"sched.sgpr", "sgpr", "SGPR"> {
+ let description = [{
+ SGPR register scheduling hint. Identity at runtime; guides the register
+ allocator to place `value` in the specified SGPR register number.
+
+ For wide values (64/128/256 bit), `regnum` is the starting register index.
+ Default is strict (allocator must obey). Pass `soft` for an advisory hint.
+
+ Example:
+ ```mlir
+ %0 = rocdl.sched.sgpr %val, 10 : i32 // strict: allocate to s10
+ %1 = rocdl.sched.sgpr %val, 4 soft : i32 // soft: prefer s4
+ ```
+ }];
+}
+
+def ROCDL_ScheduleAGPROp : ROCDL_ScheduleGPROp<"sched.agpr", "agpr", "AGPR"> {
+ let description = [{
+ AGPR register scheduling hint. Identity at runtime; guides the register
+ allocator to place `value` in the specified AGPR register number.
+ Only valid on targets with MAI instructions (e.g., gfx90a, gfx940).
+
+ For wide values (64/128/256 bit), `regnum` is the starting register index.
+ Default is strict (allocator must obey). Pass `soft` for an advisory hint.
+
+ Example:
+ ```mlir
+ %0 = rocdl.sched.agpr %val, 8 : f32 // strict: allocate to a8
+ %1 = rocdl.sched.agpr %val, 0 soft : f32 // soft: prefer a0
+ ```
+ }];
+}
+
//===---------------------------------------------------------------------===//
// Xdlops intrinsics
diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir
index 1b5b6419368819..ff62db8d69a253 100644
--- a/mlir/test/Dialect/LLVMIR/rocdl.mlir
+++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir
@@ -168,6 +168,30 @@ func.func @rocdl_iglp_opt() {
llvm.return
}
+func.func @rocdl_sched_vgpr(%val : f32, %ival : i32) -> f32 {
+ // CHECK: rocdl.sched.vgpr %{{.*}}, 42 : f32
+ %0 = rocdl.sched.vgpr %val, 42 : f32
+ // CHECK: rocdl.sched.vgpr %{{.*}}, 10 soft : i32
+ %1 = rocdl.sched.vgpr %ival, 10 soft : i32
+ llvm.return %0 : f32
+}
+
+func.func @rocdl_sched_sgpr(%val : f32, %ival : i32) -> f32 {
+ // CHECK: rocdl.sched.sgpr %{{.*}}, 20 : f32
+ %0 = rocdl.sched.sgpr %val, 20 : f32
+ // CHECK: rocdl.sched.sgpr %{{.*}}, 5 soft : i32
+ %1 = rocdl.sched.sgpr %ival, 5 soft : i32
+ llvm.return %0 : f32
+}
+
+func.func @rocdl_sched_agpr(%val : f32, %ival : i32) -> f32 {
+ // CHECK: rocdl.sched.agpr %{{.*}}, 8 : f32
+ %0 = rocdl.sched.agpr %val, 8 : f32
+ // CHECK: rocdl.sched.agpr %{{.*}}, 0 soft : i32
+ %1 = rocdl.sched.agpr %ival, 0 soft : i32
+ llvm.return %0 : f32
+}
+
func.func @rocdl.setprio() {
// CHECK: rocdl.s.setprio
rocdl.s.setprio 0
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index 2f067b8ba98fda..193f02bf674941 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -451,6 +451,35 @@ llvm.func @rocdl.iglp.opt() {
llvm.return
}
+llvm.func @rocdl.sched.vgpr(%val : f32, %ival : i32) -> f32 {
+ // CHECK-LABEL: rocdl.sched.vgpr
+ // strict (default): regnum 42, strict 1
+ // CHECK: call float @llvm.amdgcn.sched.vgpr.f32(float %{{.*}}, i32 42, i32 1)
+ %0 = rocdl.sched.vgpr %val, 42 : f32
+ // soft: regnum 10, strict 0
+ // CHECK: call i32 @llvm.amdgcn.sched.vgpr.i32(i32 %{{.*}}, i32 10, i32 0)
+ %1 = rocdl.sched.vgpr %ival, 10 soft : i32
+ llvm.return %0 : f32
+}
+
+llvm.func @rocdl.sched.sgpr(%val : f32, %ival : i32) -> f32 {
+ // CHECK-LABEL: rocdl.sched.sgpr
+ // CHECK: call float @llvm.amdgcn.sched.sgpr.f32(float %{{.*}}, i32 20, i32 1)
+ %0 = rocdl.sched.sgpr %val, 20 : f32
+ // CHECK: call i32 @llvm.amdgcn.sched.sgpr.i32(i32 %{{.*}}, i32 5, i32 0)
+ %1 = rocdl.sched.sgpr %ival, 5 soft : i32
+ llvm.return %0 : f32
+}
+
+llvm.func @rocdl.sched.agpr(%val : f32, %ival : i32) -> f32 {
+ // CHECK-LABEL: rocdl.sched.agpr
+ // CHECK: call float @llvm.amdgcn.sched.agpr.f32(float %{{.*}}, i32 8, i32 1)
+ %0 = rocdl.sched.agpr %val, 8 : f32
+ // CHECK: call i32 @llvm.amdgcn.sched.agpr.i32(i32 %{{.*}}, i32 0, i32 0)
+ %1 = rocdl.sched.agpr %ival, 0 soft : i32
+ llvm.return %0 : f32
+}
+
llvm.func @rocdl.xdlops(%arg0 : f32, %arg1 : f32,
%arg2 : vector<32 x f32>, %arg3: i32,
%arg4 : vector<16 x f32>, %arg5 : vector<4xf32>,
>From 0aed547bb5ab54ca745881f116a405b65e05eb26 Mon Sep 17 00:00:00 2001
From: Jun Li <jli10004 at amd.com>
Date: Tue, 20 Oct 2026 10:01:17 +0000
Subject: [PATCH 2/3] [AMDGPU] Add AMDGPUSchedGPR pass to lower
sched_vgpr/sgpr/agpr pseudos
Add a pre-RA machine function pass that lowers V_SCHEDULE_VGPR_B*,
S_SCHEDULE_SGPR_B*, and A_SCHEDULE_AGPR_B* pseudos into plain COPYs
while attaching AMDGPURI register allocation hints to guide the
register allocator to the requested physical register numbers.
The pass propagates hints through COPY/REG_SEQUENCE/INSERT_SUBREG/
SUBREG_TO_REG so the preference survives coalescing. Strict hints
cause getRegAllocationHints to return true, making the allocator
treat the hint as a hard constraint.
Co-Authored-By: Claude Opus 4 <noreply at anthropic.com>
---
llvm/lib/Target/AMDGPU/AMDGPU.h | 3 +
llvm/lib/Target/AMDGPU/AMDGPUPassRegistry.def | 1 +
llvm/lib/Target/AMDGPU/AMDGPUSchedGPR.cpp | 212 ++++++++++++++++++
llvm/lib/Target/AMDGPU/AMDGPUSchedGPR.h | 22 ++
.../lib/Target/AMDGPU/AMDGPUTargetMachine.cpp | 4 +
llvm/lib/Target/AMDGPU/CMakeLists.txt | 1 +
llvm/test/CodeGen/AMDGPU/sched-gpr-hints.mir | 4 +-
7 files changed, 245 insertions(+), 2 deletions(-)
create mode 100644 llvm/lib/Target/AMDGPU/AMDGPUSchedGPR.cpp
create mode 100644 llvm/lib/Target/AMDGPU/AMDGPUSchedGPR.h
diff --git a/llvm/lib/Target/AMDGPU/AMDGPU.h b/llvm/lib/Target/AMDGPU/AMDGPU.h
index 75a83c55db2bc9..12fb9095ceafc6 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPU.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPU.h
@@ -651,6 +651,9 @@ class AMDGPURewriteAGPRCopyMFMAPass
void initializeAMDGPURewriteAGPRCopyMFMALegacyPass(PassRegistry &);
extern char &AMDGPURewriteAGPRCopyMFMALegacyID;
+void initializeAMDGPUSchedGPRPass(PassRegistry &);
+extern char &AMDGPUSchedGPRID;
+
void initializeAMDGPUUniformIntrinsicCombineLegacyPass(PassRegistry &);
extern char &AMDGPUUniformIntrinsicCombineLegacyPassID;
FunctionPass *createAMDGPUUniformIntrinsicCombineLegacyPass();
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUPassRegistry.def b/llvm/lib/Target/AMDGPU/AMDGPUPassRegistry.def
index 372d5f5acab21a..97111a43198666 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUPassRegistry.def
+++ b/llvm/lib/Target/AMDGPU/AMDGPUPassRegistry.def
@@ -139,6 +139,7 @@ MACHINE_FUNCTION_PASS("amdgpu-preload-kern-arg-prolog", AMDGPUPreloadKernArgProl
MACHINE_FUNCTION_PASS("amdgpu-prepare-agpr-alloc", AMDGPUPrepareAGPRAllocPass())
MACHINE_FUNCTION_PASS("amdgpu-nsa-reassign", GCNNSAReassignPass())
MACHINE_FUNCTION_PASS("amdgpu-wait-sgpr-hazards", AMDGPUWaitSGPRHazardsPass())
+MACHINE_FUNCTION_PASS("amdgpu-sched-gpr", AMDGPUSchedGPRPass())
MACHINE_FUNCTION_PASS("gcn-create-vopd", GCNCreateVOPDPass())
MACHINE_FUNCTION_PASS("gcn-dpp-combine", GCNDPPCombinePass())
MACHINE_FUNCTION_PASS("print<amdgpu-next-use-analysis>", AMDGPUNextUseAnalysisPrinterPass(errs()))
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSchedGPR.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSchedGPR.cpp
new file mode 100644
index 00000000000000..85b4776db6a58a
--- /dev/null
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSchedGPR.cpp
@@ -0,0 +1,212 @@
+//===-- AMDGPUSchedGPR.cpp ------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file
+/// Lower the V_SCHEDULE_VGPR_B*, S_SCHEDULE_SGPR_B*, and A_SCHEDULE_AGPR_B*
+/// pseudos produced from llvm.amdgcn.sched.{vgpr,sgpr,agpr}. For each pseudo
+/// this pass:
+/// 1. reads the requested register number and strict flag from the immediate
+/// operands,
+/// 2. attaches an AMDGPURI::{VGPR,SGPR,AGPR}RegHint register allocation hint
+/// to the destination (and source) vreg so the allocator prefers that exact
+/// physical register,
+/// 3. propagates the hint through COPY / REG_SEQUENCE / INSERT_SUBREG /
+/// SUBREG_TO_REG so it survives coalescing, and
+/// 4. replaces the pseudo with a plain COPY.
+///
+/// The hint is advisory unless strict: SIRegisterInfo::getRegAllocationHints
+/// only reorders the allocation candidates, so a non-strict hint falls back to
+/// the default order and the pass never forces a spill.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPUSchedGPR.h"
+#include "AMDGPU.h"
+#include "GCNSubtarget.h"
+#include "SIInstrInfo.h"
+#include "SIRegisterInfo.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/InitializePasses.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "amdgpu-sched-gpr"
+
+enum ScheduleGPRKind { SGK_None, SGK_VGPR, SGK_SGPR, SGK_AGPR };
+
+static ScheduleGPRKind getScheduleGPRKind(unsigned Opcode) {
+ switch (Opcode) {
+ case AMDGPU::V_SCHEDULE_VGPR_B32:
+ case AMDGPU::V_SCHEDULE_VGPR_B64:
+ case AMDGPU::V_SCHEDULE_VGPR_B128:
+ case AMDGPU::V_SCHEDULE_VGPR_B256:
+ case AMDGPU::V_SCHEDULE_VGPR_B512:
+ return SGK_VGPR;
+ case AMDGPU::S_SCHEDULE_SGPR_B32:
+ case AMDGPU::S_SCHEDULE_SGPR_B64:
+ case AMDGPU::S_SCHEDULE_SGPR_B128:
+ case AMDGPU::S_SCHEDULE_SGPR_B256:
+ case AMDGPU::S_SCHEDULE_SGPR_B512:
+ return SGK_SGPR;
+ case AMDGPU::A_SCHEDULE_AGPR_B32:
+ case AMDGPU::A_SCHEDULE_AGPR_B64:
+ case AMDGPU::A_SCHEDULE_AGPR_B128:
+ case AMDGPU::A_SCHEDULE_AGPR_B256:
+ case AMDGPU::A_SCHEDULE_AGPR_B512:
+ return SGK_AGPR;
+ default:
+ return SGK_None;
+ }
+}
+
+/// Walk the SSA use chain of \p Reg and propagate the hint through
+/// value-preserving pseudos so the preference reaches the vregs that survive
+/// coalescing.
+static void propagateHint(MachineRegisterInfo &MRI, Register Reg,
+ unsigned HintValue, unsigned HintKind,
+ SmallDenseSet<unsigned, 32> &Visited) {
+ if (!Reg.isVirtual() || !Visited.insert(Reg.id()).second)
+ return;
+
+ for (MachineInstr &UseMI : MRI.use_nodbg_instructions(Reg)) {
+ Register DefReg;
+ switch (UseMI.getOpcode()) {
+ case TargetOpcode::COPY:
+ case TargetOpcode::REG_SEQUENCE:
+ case TargetOpcode::INSERT_SUBREG:
+ case TargetOpcode::SUBREG_TO_REG:
+ DefReg = UseMI.getOperand(0).getReg();
+ break;
+ default:
+ continue;
+ }
+ if (!DefReg.isVirtual())
+ continue;
+
+ // Do not clobber an existing, conflicting hint.
+ std::pair<unsigned, Register> Existing = MRI.getRegAllocationHint(DefReg);
+ if (Existing.first != 0 &&
+ (Existing.first != HintKind || Existing.second != HintValue))
+ continue;
+
+ MRI.setRegAllocationHint(DefReg, HintKind, HintValue);
+ propagateHint(MRI, DefReg, HintValue, HintKind, Visited);
+ }
+}
+
+static bool runSchedGPR(MachineFunction &MF) {
+ MachineRegisterInfo &MRI = MF.getRegInfo();
+ const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
+ const SIInstrInfo *TII = ST.getInstrInfo();
+ bool Changed = false;
+
+ SmallVector<MachineInstr *, 16> ToErase;
+
+ for (MachineBasicBlock &MBB : MF) {
+ for (MachineInstr &MI : MBB) {
+ ScheduleGPRKind Kind = getScheduleGPRKind(MI.getOpcode());
+ if (Kind == SGK_None)
+ continue;
+
+ Register DstReg = MI.getOperand(0).getReg();
+ const MachineOperand &SrcMO = MI.getOperand(1);
+ Register SrcReg = SrcMO.getReg();
+ unsigned RegNum = MI.getOperand(2).getImm();
+ bool Strict = MI.getOperand(3).getImm() != 0;
+
+ unsigned HintKind;
+ switch (Kind) {
+ case SGK_VGPR:
+ HintKind = Strict ? AMDGPURI::StrictVGPRRegHint
+ : AMDGPURI::VGPRRegHint;
+ break;
+ case SGK_SGPR:
+ HintKind = Strict ? AMDGPURI::StrictSGPRRegHint
+ : AMDGPURI::SGPRRegHint;
+ break;
+ case SGK_AGPR:
+ HintKind = Strict ? AMDGPURI::StrictAGPRRegHint
+ : AMDGPURI::AGPRRegHint;
+ break;
+ default:
+ llvm_unreachable("unexpected schedule GPR kind");
+ }
+
+ LLVM_DEBUG(dbgs() << " sched."
+ << (Kind == SGK_VGPR ? "vgpr"
+ : Kind == SGK_SGPR ? "sgpr"
+ : "agpr")
+ << ": " << printReg(DstReg) << " <- "
+ << printReg(SrcReg) << " reg " << RegNum << " strict "
+ << Strict << '\n');
+
+ // Hint both the destination and the source so the preference survives
+ // whichever side coalescing keeps.
+ if (DstReg.isVirtual())
+ MRI.setRegAllocationHint(DstReg, HintKind, RegNum);
+ if (SrcReg.isVirtual())
+ MRI.setRegAllocationHint(SrcReg, HintKind, RegNum);
+
+ SmallDenseSet<unsigned, 32> Visited;
+ if (DstReg.isVirtual())
+ propagateHint(MRI, DstReg, RegNum, HintKind, Visited);
+
+ // Replace the pseudo with a plain COPY.
+ BuildMI(MBB, MI, MI.getDebugLoc(), TII->get(TargetOpcode::COPY), DstReg)
+ .addReg(SrcReg, getRegState(SrcMO), SrcMO.getSubReg());
+
+ ToErase.push_back(&MI);
+ Changed = true;
+ }
+ }
+
+ for (MachineInstr *MI : ToErase)
+ MI->eraseFromParent();
+
+ return Changed;
+}
+
+namespace {
+
+class AMDGPUSchedGPR : public MachineFunctionPass {
+public:
+ static char ID;
+
+ AMDGPUSchedGPR() : MachineFunctionPass(ID) {}
+
+ bool runOnMachineFunction(MachineFunction &MF) override {
+ return runSchedGPR(MF);
+ }
+
+ StringRef getPassName() const override { return "AMDGPU Sched GPR"; }
+
+ void getAnalysisUsage(AnalysisUsage &AU) const override {
+ AU.setPreservesCFG();
+ MachineFunctionPass::getAnalysisUsage(AU);
+ }
+};
+
+} // end anonymous namespace
+
+INITIALIZE_PASS(AMDGPUSchedGPR, DEBUG_TYPE, "AMDGPU Sched GPR", false, false)
+
+char AMDGPUSchedGPR::ID = 0;
+
+char &llvm::AMDGPUSchedGPRID = AMDGPUSchedGPR::ID;
+
+PreservedAnalyses
+AMDGPUSchedGPRPass::run(MachineFunction &MF,
+ MachineFunctionAnalysisManager &MFAM) {
+ if (!runSchedGPR(MF))
+ return PreservedAnalyses::all();
+
+ return getMachineFunctionPassPreservedAnalyses().preserveSet<CFGAnalyses>();
+}
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSchedGPR.h b/llvm/lib/Target/AMDGPU/AMDGPUSchedGPR.h
new file mode 100644
index 00000000000000..fd7c0a824c98ee
--- /dev/null
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSchedGPR.h
@@ -0,0 +1,22 @@
+//===- AMDGPUSchedGPR.h -----------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_AMDGPU_AMDGPUSCHEDGPR_H
+#define LLVM_LIB_TARGET_AMDGPU_AMDGPUSCHEDGPR_H
+
+#include "llvm/CodeGen/MachinePassManager.h"
+
+namespace llvm {
+class AMDGPUSchedGPRPass : public OptionalPassInfoMixin<AMDGPUSchedGPRPass> {
+public:
+ PreservedAnalyses run(MachineFunction &MF,
+ MachineFunctionAnalysisManager &MFAM);
+};
+} // namespace llvm
+
+#endif // LLVM_LIB_TARGET_AMDGPU_AMDGPUSCHEDGPR_H
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index 6a5b27e5775c62..137ff98b173e64 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -35,6 +35,7 @@
#include "AMDGPURemoveIncompatibleFunctions.h"
#include "AMDGPUReserveWWMRegs.h"
#include "AMDGPUResourceUsageAnalysis.h"
+#include "AMDGPUSchedGPR.h"
#include "AMDGPUSplitModule.h"
#include "AMDGPUTargetObjectFile.h"
#include "AMDGPUTargetTransformInfo.h"
@@ -733,6 +734,7 @@ extern "C" LLVM_ABI LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAMDGPUTarget() {
initializeAMDGPULowerIntrinsicsLegacyPass(*PR);
initializeAMDGPUReserveWWMRegsLegacyPass(*PR);
initializeAMDGPURewriteAGPRCopyMFMALegacyPass(*PR);
+ initializeAMDGPUSchedGPRPass(*PR);
initializeAMDGPURewriteOutArgumentsPass(*PR);
initializeAMDGPURewriteUndefForPHILegacyPass(*PR);
initializeSIAnnotateControlFlowLegacyPass(*PR);
@@ -1846,6 +1848,7 @@ void GCNPassConfig::addFastRegAlloc() {
}
void GCNPassConfig::addPreRegAlloc() {
+ addPass(&AMDGPUSchedGPRID);
if (getOptLevel() != CodeGenOptLevel::None)
addPass(&AMDGPUPrepareAGPRAllocLegacyID);
if (getOptLevel() >= CodeGenOptLevel::Default && EnableMachinePipeliner)
@@ -2708,6 +2711,7 @@ Error AMDGPUCodeGenPassBuilder::addOptimizedRegAlloc(PassManagerWrapper &PMW) {
}
void AMDGPUCodeGenPassBuilder::addPreRegAlloc(PassManagerWrapper &PMW) {
+ addMachineFunctionPass(AMDGPUSchedGPRPass(), PMW);
if (getOptLevel() != CodeGenOptLevel::None)
addMachineFunctionPass(AMDGPUPrepareAGPRAllocPass(), PMW);
if (getOptLevel() >= CodeGenOptLevel::Default && EnableMachinePipeliner)
diff --git a/llvm/lib/Target/AMDGPU/CMakeLists.txt b/llvm/lib/Target/AMDGPU/CMakeLists.txt
index b6c5f63379344c..7baf68c021ae1c 100644
--- a/llvm/lib/Target/AMDGPU/CMakeLists.txt
+++ b/llvm/lib/Target/AMDGPU/CMakeLists.txt
@@ -113,6 +113,7 @@ add_llvm_target(AMDGPUCodeGen
AMDGPUResourceUsageAnalysis.cpp
AMDGPURewriteAGPRCopyMFMA.cpp
AMDGPURewriteOutArguments.cpp
+ AMDGPUSchedGPR.cpp
AMDGPURewriteUndefForPHI.cpp
AMDGPUSelectionDAGInfo.cpp
AMDGPUSetWavePriority.cpp
diff --git a/llvm/test/CodeGen/AMDGPU/sched-gpr-hints.mir b/llvm/test/CodeGen/AMDGPU/sched-gpr-hints.mir
index d480eea0ce596b..a1e178ece02ed9 100644
--- a/llvm/test/CodeGen/AMDGPU/sched-gpr-hints.mir
+++ b/llvm/test/CodeGen/AMDGPU/sched-gpr-hints.mir
@@ -1,6 +1,6 @@
-# RUN: llc -mtriple=amdgcn -mcpu=gfx1250 -run-pass=amdgpu-sched-bank -o - %s | FileCheck %s
+# RUN: llc -mtriple=amdgcn -mcpu=gfx1250 -run-pass=amdgpu-sched-gpr -o - %s | FileCheck %s
-# The AMDGPUSchedBank pass attaches exact register allocation hints (not
+# The AMDGPUSchedGPR pass attaches exact register allocation hints (not
# printed in MIR) and lowers each schedule pseudo to a plain COPY.
---
>From 72113c3a42c257ac5eda5ce13a533fd537f0969f Mon Sep 17 00:00:00 2001
From: Jun Li <jli10004 at amd.com>
Date: Tue, 20 Oct 2026 10:04:02 +0000
Subject: [PATCH 3/3] [AMDGPU] Add AMDGPUSchedGPR to llc pipeline tests
Update llc-pipeline.ll and llc-pipeline-npm.ll to include the
amdgpu-sched-gpr pass in the expected pass ordering for all
optimization levels (O0-O3).
Co-Authored-By: Claude Opus 4 <noreply at anthropic.com>
---
llvm/test/CodeGen/AMDGPU/llc-pipeline-npm.ll | 3 +++
llvm/test/CodeGen/AMDGPU/llc-pipeline.ll | 5 +++++
2 files changed, 8 insertions(+)
diff --git a/llvm/test/CodeGen/AMDGPU/llc-pipeline-npm.ll b/llvm/test/CodeGen/AMDGPU/llc-pipeline-npm.ll
index 3e44687d1856ed..73d5e1ea7e9583 100644
--- a/llvm/test/CodeGen/AMDGPU/llc-pipeline-npm.ll
+++ b/llvm/test/CodeGen/AMDGPU/llc-pipeline-npm.ll
@@ -68,6 +68,7 @@
; GCN-O0-NEXT: function
; GCN-O0-NEXT: machine-function
; GCN-O0-NEXT: reg-usage-propagation
+; GCN-O0-NEXT: amdgpu-sched-gpr
; GCN-O0-NEXT: phi-node-elimination
; GCN-O0-NEXT: si-lower-control-flow
; GCN-O0-NEXT: two-address-instruction
@@ -215,6 +216,7 @@
; GCN-O2-NEXT: function
; GCN-O2-NEXT: machine-function
; GCN-O2-NEXT: reg-usage-propagation
+; GCN-O2-NEXT: amdgpu-sched-gpr
; GCN-O2-NEXT: amdgpu-prepare-agpr-alloc
; GCN-O2-NEXT: detect-dead-lanes
; GCN-O2-NEXT: dead-mi-elimination
@@ -404,6 +406,7 @@
; GCN-O3-NEXT: function
; GCN-O3-NEXT: machine-function
; GCN-O3-NEXT: reg-usage-propagation
+; GCN-O3-NEXT: amdgpu-sched-gpr
; GCN-O3-NEXT: amdgpu-prepare-agpr-alloc
; GCN-O3-NEXT: detect-dead-lanes
; GCN-O3-NEXT: dead-mi-elimination
diff --git a/llvm/test/CodeGen/AMDGPU/llc-pipeline.ll b/llvm/test/CodeGen/AMDGPU/llc-pipeline.ll
index a97c1955c94f5c..e9e4b8d364ec95 100644
--- a/llvm/test/CodeGen/AMDGPU/llc-pipeline.ll
+++ b/llvm/test/CodeGen/AMDGPU/llc-pipeline.ll
@@ -104,6 +104,7 @@
; GCN-O0-NEXT: Finalize ISel and expand pseudo-instructions
; GCN-O0-NEXT: Local Stack Slot Allocation
; GCN-O0-NEXT: Register Usage Information Propagation
+; GCN-O0-NEXT: AMDGPU Sched GPR
; GCN-O0-NEXT: Eliminate PHI nodes for register allocation
; GCN-O0-NEXT: SI Lower control flow pseudo instructions
; GCN-O0-NEXT: Two-Address instruction pass
@@ -344,6 +345,7 @@
; GCN-O1-NEXT: Remove dead machine instructions
; GCN-O1-NEXT: SI Shrink Instructions
; GCN-O1-NEXT: Register Usage Information Propagation
+; GCN-O1-NEXT: AMDGPU Sched GPR
; GCN-O1-NEXT: AMDGPU Prepare AGPR Alloc
; GCN-O1-NEXT: Detect Dead Lanes
; GCN-O1-NEXT: Remove dead machine instructions
@@ -673,6 +675,7 @@
; GCN-O1-OPTS-NEXT: Remove dead machine instructions
; GCN-O1-OPTS-NEXT: SI Shrink Instructions
; GCN-O1-OPTS-NEXT: Register Usage Information Propagation
+; GCN-O1-OPTS-NEXT: AMDGPU Sched GPR
; GCN-O1-OPTS-NEXT: AMDGPU Prepare AGPR Alloc
; GCN-O1-OPTS-NEXT: Detect Dead Lanes
; GCN-O1-OPTS-NEXT: Remove dead machine instructions
@@ -1006,6 +1009,7 @@
; GCN-O2-NEXT: Remove dead machine instructions
; GCN-O2-NEXT: SI Shrink Instructions
; GCN-O2-NEXT: Register Usage Information Propagation
+; GCN-O2-NEXT: AMDGPU Sched GPR
; GCN-O2-NEXT: AMDGPU Prepare AGPR Alloc
; GCN-O2-NEXT: Detect Dead Lanes
; GCN-O2-NEXT: Remove dead machine instructions
@@ -1355,6 +1359,7 @@
; GCN-O3-NEXT: Remove dead machine instructions
; GCN-O3-NEXT: SI Shrink Instructions
; GCN-O3-NEXT: Register Usage Information Propagation
+; GCN-O3-NEXT: AMDGPU Sched GPR
; GCN-O3-NEXT: AMDGPU Prepare AGPR Alloc
; GCN-O3-NEXT: Detect Dead Lanes
; GCN-O3-NEXT: Remove dead machine instructions
More information about the Mlir-commits
mailing list