[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:20:21 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-amdgpu
Author: jli (jli-melchior)
<details>
<summary>Changes</summary>
## 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
---
Patch is 45.11 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/233773.diff
22 Files Affected:
- (modified) clang/include/clang/Basic/BuiltinsAMDGPU.td (+5)
- (modified) clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp (+26)
- (modified) clang/lib/Sema/SemaAMDGPU.cpp (+28)
- (added) clang/test/CodeGenOpenCL/builtins-amdgcn-sched-gpr.cl (+68)
- (modified) llvm/include/llvm/IR/IntrinsicsAMDGPU.td (+22)
- (modified) llvm/lib/Target/AMDGPU/AMDGPU.h (+3)
- (modified) llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp (+57)
- (modified) llvm/lib/Target/AMDGPU/AMDGPUPassRegistry.def (+1)
- (added) llvm/lib/Target/AMDGPU/AMDGPUSchedGPR.cpp (+212)
- (added) llvm/lib/Target/AMDGPU/AMDGPUSchedGPR.h (+22)
- (modified) llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp (+4)
- (modified) llvm/lib/Target/AMDGPU/CMakeLists.txt (+1)
- (modified) llvm/lib/Target/AMDGPU/SIInstructions.td (+48)
- (modified) llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp (+21)
- (modified) llvm/lib/Target/AMDGPU/SIRegisterInfo.h (+10-1)
- (modified) llvm/test/CodeGen/AMDGPU/llc-pipeline-npm.ll (+3)
- (modified) llvm/test/CodeGen/AMDGPU/llc-pipeline.ll (+5)
- (added) llvm/test/CodeGen/AMDGPU/sched-gpr-hints.mir (+66)
- (added) llvm/test/CodeGen/AMDGPU/sched-gpr.ll (+101)
- (modified) mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td (+73)
- (modified) mlir/test/Dialect/LLVMIR/rocdl.mlir (+24)
- (modified) mlir/test/Target/LLVMIR/rocdl.mlir (+29)
``````````diff
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/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/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/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 "AMD...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/233773
More information about the Mlir-commits
mailing list