[llvm] [AMDGPU][GlobalISel] Add register bank legalize rules for PST rsq. (PR #196116)
via llvm-commits
llvm-commits at lists.llvm.org
Wed May 6 09:50:03 PDT 2026
https://github.com/anjenner created https://github.com/llvm/llvm-project/pull/196116
None
>From ce93701a564b3af6f858dc92adc72798663c843c Mon Sep 17 00:00:00 2001
From: Andrew Jenner <Andrew.Jenner at amd.com>
Date: Wed, 6 May 2026 12:56:36 -0400
Subject: [PATCH] [AMDGPU][GlobalISel] Add register bank legalize rules for PST
rsq.
---
.../AMDGPU/AMDGPURegBankLegalizeRules.cpp | 6 +-
.../CodeGen/AMDGPU/llvm.amdgcn.rsq.f16.ll | 10 ++-
llvm/test/CodeGen/AMDGPU/llvm.amdgcn.rsq.ll | 71 +++++++++++++++++++
3 files changed, 83 insertions(+), 4 deletions(-)
diff --git a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
index 118d75dd6bbe3..690acfb8b9e38 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
@@ -1726,9 +1726,11 @@ RegBankLegalizeRules::RegBankLegalizeRules(const GCNSubtarget &_ST,
.Div(S32, {{Vgpr32}, {IntrId, Vgpr32, Vgpr32, Vgpr32}});
addRulesForIOpcs({amdgcn_rsq, amdgcn_rsq_clamp}, Standard)
- .Uni(S16, {{UniInVgprS16}, {IntrId, Vgpr16}})
+ .Uni(S16, {{Sgpr16}, {IntrId, Sgpr16}}, hasPST)
+ .Uni(S16, {{UniInVgprS16}, {IntrId, Vgpr16}}, !hasPST)
.Div(S16, {{Vgpr16}, {IntrId, Vgpr16}})
- .Uni(S32, {{UniInVgprS32}, {IntrId, Vgpr32}})
+ .Uni(S32, {{Sgpr32}, {IntrId, Sgpr32}}, hasPST)
+ .Uni(S32, {{UniInVgprS32}, {IntrId, Vgpr32}}, !hasPST)
.Div(S32, {{Vgpr32}, {IntrId, Vgpr32}})
.Uni(S64, {{UniInVgprS64}, {IntrId, Vgpr64}})
.Div(S64, {{Vgpr64}, {IntrId, Vgpr64}});
diff --git a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.rsq.f16.ll b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.rsq.f16.ll
index 7ac8caf71af22..dc36c027c746e 100644
--- a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.rsq.f16.ll
+++ b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.rsq.f16.ll
@@ -156,7 +156,10 @@ define amdgpu_kernel void @rsq_f16(
; GISEL-GFX12-TRUE16-NEXT: s_load_u16 s2, s[2:3], 0x0
; GISEL-GFX12-TRUE16-NEXT: s_mov_b32 s3, 0x31016000
; GISEL-GFX12-TRUE16-NEXT: s_wait_kmcnt 0x0
-; GISEL-GFX12-TRUE16-NEXT: v_rsq_f16_e32 v0.l, s2
+; GISEL-GFX12-TRUE16-NEXT: v_s_rsq_f16 s2, s2
+; GISEL-GFX12-TRUE16-NEXT: s_wait_alu depctr_va_sdst(0)
+; GISEL-GFX12-TRUE16-NEXT: s_delay_alu instid0(TRANS32_DEP_1)
+; GISEL-GFX12-TRUE16-NEXT: v_mov_b16_e32 v0.l, s2
; GISEL-GFX12-TRUE16-NEXT: s_mov_b32 s2, -1
; GISEL-GFX12-TRUE16-NEXT: buffer_store_b16 v0, off, s[0:3], null
; GISEL-GFX12-TRUE16-NEXT: s_endpgm
@@ -168,7 +171,10 @@ define amdgpu_kernel void @rsq_f16(
; GISEL-GFX12-FAKE16-NEXT: s_load_u16 s2, s[2:3], 0x0
; GISEL-GFX12-FAKE16-NEXT: s_mov_b32 s3, 0x31016000
; GISEL-GFX12-FAKE16-NEXT: s_wait_kmcnt 0x0
-; GISEL-GFX12-FAKE16-NEXT: v_rsq_f16_e32 v0, s2
+; GISEL-GFX12-FAKE16-NEXT: v_s_rsq_f16 s2, s2
+; GISEL-GFX12-FAKE16-NEXT: s_wait_alu depctr_va_sdst(0)
+; GISEL-GFX12-FAKE16-NEXT: s_delay_alu instid0(TRANS32_DEP_1)
+; GISEL-GFX12-FAKE16-NEXT: v_mov_b32_e32 v0, s2
; GISEL-GFX12-FAKE16-NEXT: s_mov_b32 s2, -1
; GISEL-GFX12-FAKE16-NEXT: buffer_store_b16 v0, off, s[0:3], null
; GISEL-GFX12-FAKE16-NEXT: s_endpgm
diff --git a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.rsq.ll b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.rsq.ll
index 1cd24e876eec4..74b89953aa730 100644
--- a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.rsq.ll
+++ b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.rsq.ll
@@ -3,6 +3,7 @@
; RUN: llc -mtriple=amdgcn -mcpu=tonga < %s | FileCheck -check-prefixes=VI %s
; RUN: llc -mtriple=amdgcn -global-isel=1 -new-reg-bank-select < %s | FileCheck -check-prefixes=SI-GISEL %s
; RUN: llc -mtriple=amdgcn -global-isel=1 -new-reg-bank-select -mcpu=tonga -mattr=-flat-for-global < %s | FileCheck -check-prefixes=VI-GISEL %s
+; RUN: llc -mtriple=amdgcn -global-isel=1 -new-reg-bank-select -mcpu=gfx1200 < %s | FileCheck -check-prefixes=GFX12-GISEL %s
declare float @llvm.amdgcn.rsq.f32(float) #0
declare double @llvm.amdgcn.rsq.f64(double) #0
@@ -51,6 +52,18 @@ define amdgpu_kernel void @rsq_f32(ptr addrspace(1) %out, float %src) #1 {
; VI-GISEL-NEXT: s_mov_b32 s2, -1
; VI-GISEL-NEXT: buffer_store_dword v0, off, s[0:3], 0
; VI-GISEL-NEXT: s_endpgm
+;
+; GFX12-GISEL-LABEL: rsq_f32:
+; GFX12-GISEL: ; %bb.0:
+; GFX12-GISEL-NEXT: s_load_b96 s[0:2], s[4:5], 0x24
+; GFX12-GISEL-NEXT: v_mov_b32_e32 v1, 0
+; GFX12-GISEL-NEXT: s_wait_kmcnt 0x0
+; GFX12-GISEL-NEXT: v_s_rsq_f32 s2, s2
+; GFX12-GISEL-NEXT: s_wait_alu depctr_va_sdst(0)
+; GFX12-GISEL-NEXT: s_delay_alu instid0(TRANS32_DEP_1)
+; GFX12-GISEL-NEXT: v_mov_b32_e32 v0, s2
+; GFX12-GISEL-NEXT: global_store_b32 v1, v0, s[0:1]
+; GFX12-GISEL-NEXT: s_endpgm
%rsq = call float @llvm.amdgcn.rsq.f32(float %src) #0
store float %rsq, ptr addrspace(1) %out, align 4
ret void
@@ -97,6 +110,16 @@ define amdgpu_kernel void @rsq_f32_constant_4.0(ptr addrspace(1) %out) #1 {
; VI-GISEL-NEXT: s_waitcnt lgkmcnt(0)
; VI-GISEL-NEXT: buffer_store_dword v0, off, s[0:3], 0
; VI-GISEL-NEXT: s_endpgm
+;
+; GFX12-GISEL-LABEL: rsq_f32_constant_4.0:
+; GFX12-GISEL: ; %bb.0:
+; GFX12-GISEL-NEXT: s_load_b64 s[0:1], s[4:5], 0x24
+; GFX12-GISEL-NEXT: v_s_rsq_f32 s2, 4.0
+; GFX12-GISEL-NEXT: s_delay_alu instid0(TRANS32_DEP_1)
+; GFX12-GISEL-NEXT: v_dual_mov_b32 v1, 0 :: v_dual_mov_b32 v0, s2
+; GFX12-GISEL-NEXT: s_wait_kmcnt 0x0
+; GFX12-GISEL-NEXT: global_store_b32 v1, v0, s[0:1]
+; GFX12-GISEL-NEXT: s_endpgm
%rsq = call float @llvm.amdgcn.rsq.f32(float 4.0) #0
store float %rsq, ptr addrspace(1) %out, align 4
ret void
@@ -142,6 +165,16 @@ define amdgpu_kernel void @rsq_f32_constant_100.0(ptr addrspace(1) %out) #1 {
; VI-GISEL-NEXT: s_waitcnt lgkmcnt(0)
; VI-GISEL-NEXT: buffer_store_dword v0, off, s[0:3], 0
; VI-GISEL-NEXT: s_endpgm
+;
+; GFX12-GISEL-LABEL: rsq_f32_constant_100.0:
+; GFX12-GISEL: ; %bb.0:
+; GFX12-GISEL-NEXT: s_load_b64 s[0:1], s[4:5], 0x24
+; GFX12-GISEL-NEXT: v_s_rsq_f32 s2, 0x42c80000
+; GFX12-GISEL-NEXT: s_delay_alu instid0(TRANS32_DEP_1)
+; GFX12-GISEL-NEXT: v_dual_mov_b32 v1, 0 :: v_dual_mov_b32 v0, s2
+; GFX12-GISEL-NEXT: s_wait_kmcnt 0x0
+; GFX12-GISEL-NEXT: global_store_b32 v1, v0, s[0:1]
+; GFX12-GISEL-NEXT: s_endpgm
%rsq = call float @llvm.amdgcn.rsq.f32(float 100.0) #0
store float %rsq, ptr addrspace(1) %out, align 4
ret void
@@ -189,6 +222,15 @@ define amdgpu_kernel void @rsq_f64(ptr addrspace(1) %out, double %src) #1 {
; VI-GISEL-NEXT: s_mov_b32 s3, 0xf000
; VI-GISEL-NEXT: buffer_store_dwordx2 v[0:1], off, s[0:3], 0
; VI-GISEL-NEXT: s_endpgm
+;
+; GFX12-GISEL-LABEL: rsq_f64:
+; GFX12-GISEL: ; %bb.0:
+; GFX12-GISEL-NEXT: s_load_b128 s[0:3], s[4:5], 0x24
+; GFX12-GISEL-NEXT: v_mov_b32_e32 v2, 0
+; GFX12-GISEL-NEXT: s_wait_kmcnt 0x0
+; GFX12-GISEL-NEXT: v_rsq_f64_e32 v[0:1], s[2:3]
+; GFX12-GISEL-NEXT: global_store_b64 v2, v[0:1], s[0:1]
+; GFX12-GISEL-NEXT: s_endpgm
%rsq = call double @llvm.amdgcn.rsq.f64(double %src) #0
store double %rsq, ptr addrspace(1) %out, align 4
ret void
@@ -235,6 +277,15 @@ define amdgpu_kernel void @rsq_f64_constant_4.0(ptr addrspace(1) %out) #1 {
; VI-GISEL-NEXT: s_waitcnt lgkmcnt(0)
; VI-GISEL-NEXT: buffer_store_dwordx2 v[0:1], off, s[0:3], 0
; VI-GISEL-NEXT: s_endpgm
+;
+; GFX12-GISEL-LABEL: rsq_f64_constant_4.0:
+; GFX12-GISEL: ; %bb.0:
+; GFX12-GISEL-NEXT: v_rsq_f64_e32 v[0:1], 4.0
+; GFX12-GISEL-NEXT: s_load_b64 s[0:1], s[4:5], 0x24
+; GFX12-GISEL-NEXT: v_mov_b32_e32 v2, 0
+; GFX12-GISEL-NEXT: s_wait_kmcnt 0x0
+; GFX12-GISEL-NEXT: global_store_b64 v2, v[0:1], s[0:1]
+; GFX12-GISEL-NEXT: s_endpgm
%rsq = call double @llvm.amdgcn.rsq.f64(double 4.0) #0
store double %rsq, ptr addrspace(1) %out, align 4
ret void
@@ -280,6 +331,15 @@ define amdgpu_kernel void @rsq_f64_constant_100.0(ptr addrspace(1) %out) #1 {
; VI-GISEL-NEXT: s_waitcnt lgkmcnt(0)
; VI-GISEL-NEXT: buffer_store_dwordx2 v[0:1], off, s[0:3], 0
; VI-GISEL-NEXT: s_endpgm
+;
+; GFX12-GISEL-LABEL: rsq_f64_constant_100.0:
+; GFX12-GISEL: ; %bb.0:
+; GFX12-GISEL-NEXT: v_rsq_f64_e32 v[0:1], 0x40590000
+; GFX12-GISEL-NEXT: s_load_b64 s[0:1], s[4:5], 0x24
+; GFX12-GISEL-NEXT: v_mov_b32_e32 v2, 0
+; GFX12-GISEL-NEXT: s_wait_kmcnt 0x0
+; GFX12-GISEL-NEXT: global_store_b64 v2, v[0:1], s[0:1]
+; GFX12-GISEL-NEXT: s_endpgm
%rsq = call double @llvm.amdgcn.rsq.f64(double 100.0) #0
store double %rsq, ptr addrspace(1) %out, align 4
ret void
@@ -316,6 +376,17 @@ define amdgpu_kernel void @rsq_undef_f32(ptr addrspace(1) %out) #1 {
; VI-GISEL-NEXT: v_rsq_f32_e32 v0, s0
; VI-GISEL-NEXT: buffer_store_dword v0, off, s[0:3], 0
; VI-GISEL-NEXT: s_endpgm
+;
+; GFX12-GISEL-LABEL: rsq_undef_f32:
+; GFX12-GISEL: ; %bb.0:
+; GFX12-GISEL-NEXT: s_load_b64 s[0:1], s[4:5], 0x24
+; GFX12-GISEL-NEXT: v_mov_b32_e32 v1, 0
+; GFX12-GISEL-NEXT: s_wait_kmcnt 0x0
+; GFX12-GISEL-NEXT: v_s_rsq_f32 s2, s0
+; GFX12-GISEL-NEXT: s_delay_alu instid0(TRANS32_DEP_1)
+; GFX12-GISEL-NEXT: v_mov_b32_e32 v0, s2
+; GFX12-GISEL-NEXT: global_store_b32 v1, v0, s[0:1]
+; GFX12-GISEL-NEXT: s_endpgm
%rsq = call float @llvm.amdgcn.rsq.f32(float poison)
store float %rsq, ptr addrspace(1) %out, align 4
ret void
More information about the llvm-commits
mailing list