[llvm] [AMDGPU][GISel] Add register bank legalization rules for amdgcn_raw_buffer_load_async_lds (PR #201406)

[via llvm-commits]

https://github.com/anjenner updated https://github.com/llvm/llvm-project/pull/201406

>From 8034bd7e7856d986f1b5794417c711026f92818c Mon Sep 17 00:00:00 2001
From: Andrew Jenner <Andrew.Jenner at amd.com>
Date: Wed, 3 Jun 2026 12:39:47 -0400
Subject: [PATCH 1/3] [AMDGPU][GISel] Add register bank legalization rules for
 amdgcn_raw_buffer_load_async_lds, amdgcn_struct_buffer_load_async_lds,
 amdgcn_raw_ptr_buffer_load_async_lds, and
 amdgcn_struct_ptr_buffer_load_async_lds.

---
 .../Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp   | 14 ++++++++++++++
 llvm/test/CodeGen/AMDGPU/async-buffer-loads.ll     |  2 +-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
index e5bab2bbdfa25..d70008d903399 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
@@ -1906,6 +1906,14 @@ RegBankLegalizeRules::RegBankLegalizeRules(const GCNSubtarget &_ST,
   addRulesForIOpcs({amdgcn_raw_buffer_load_lds})
       .Any({{_}, {{}, {IntrId, SgprV4S32, SgprP3, Imm, Vgpr32, Sgpr32}}});
 
+  addRulesForIOpcs({amdgcn_raw_buffer_load_async_lds})
+      .Any({{_}, {{}, {IntrId, SgprV4S32, SgprB32_M0, Imm, Vgpr32, Sgpr32}}});
+
+  addRulesForIOpcs({amdgcn_struct_buffer_load_async_lds})
+      .Any(
+	 {{_},
+	  {{}, {IntrId, SgprV4S32, SgprB32_M0, Imm, Vgpr32, Vgpr32, Sgpr32}}});
+
   addRulesForIOpcs({amdgcn_struct_buffer_load_lds})
       .Any({{_},
             {{}, {IntrId, SgprV4S32, SgprP3, Imm, Vgpr32, Vgpr32, Sgpr32}}});
@@ -1913,6 +1921,12 @@ RegBankLegalizeRules::RegBankLegalizeRules(const GCNSubtarget &_ST,
   addRulesForIOpcs({amdgcn_raw_ptr_buffer_load_lds})
       .Any({{_}, {{}, {IntrId, SgprP8, SgprP3, Imm, Vgpr32, Sgpr32}}});
 
+  addRulesForIOpcs({amdgcn_raw_ptr_buffer_load_async_lds})
+      .Any({{}, {{}, {IntrId, SgprP8, SgprB32_M0, Imm, VgprB32, SgprB32}}});
+
+  addRulesForIOpcs({amdgcn_struct_ptr_buffer_load_async_lds})
+      .Any({{_}, {{}, {IntrId, SgprP8, SgprP3, Imm, Vgpr32, Vgpr32, Sgpr32}}});
+
   addRulesForIOpcs({amdgcn_struct_ptr_buffer_load_lds})
       .Any({{_}, {{}, {IntrId, SgprP8, SgprP3, Imm, Vgpr32, Vgpr32, Sgpr32}}});
 
diff --git a/llvm/test/CodeGen/AMDGPU/async-buffer-loads.ll b/llvm/test/CodeGen/AMDGPU/async-buffer-loads.ll
index 184fd61c355f0..129b4186a3fb0 100644
--- a/llvm/test/CodeGen/AMDGPU/async-buffer-loads.ll
+++ b/llvm/test/CodeGen/AMDGPU/async-buffer-loads.ll
@@ -1,6 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
 ; RUN: llc -global-isel=0 -mtriple=amdgcn -mcpu=gfx900 < %s | FileCheck %s
-; RUN: llc -global-isel=1 -mtriple=amdgcn -mcpu=gfx900 < %s | FileCheck %s
+; RUN: llc -global-isel=1 -mtriple=amdgcn -mcpu=gfx900 -new-reg-bank-select < %s | FileCheck %s
 
 define float @raw.buffer.load(<4 x i32> inreg %rsrc, ptr addrspace(3) inreg %lds) {
 ; CHECK-LABEL: raw.buffer.load:

>From 7d5aa41fe8e9d7a8a87ddbb6e3a532e75c5beeef Mon Sep 17 00:00:00 2001
From: Andrew Jenner <Andrew.Jenner at amd.com>
Date: Thu, 4 Jun 2026 05:13:00 -0400
Subject: [PATCH 2/3] Fix formatting.

---
 llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
index d70008d903399..fde656e7c15c9 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
@@ -1911,8 +1911,8 @@ RegBankLegalizeRules::RegBankLegalizeRules(const GCNSubtarget &_ST,
 
   addRulesForIOpcs({amdgcn_struct_buffer_load_async_lds})
       .Any(
-	 {{_},
-	  {{}, {IntrId, SgprV4S32, SgprB32_M0, Imm, Vgpr32, Vgpr32, Sgpr32}}});
+          {{_},
+           {{}, {IntrId, SgprV4S32, SgprB32_M0, Imm, Vgpr32, Vgpr32, Sgpr32}}});
 
   addRulesForIOpcs({amdgcn_struct_buffer_load_lds})
       .Any({{_},

>From 8f4b607848d21aa153559617c37866058340ff69 Mon Sep 17 00:00:00 2001
From: Andrew Jenner <Andrew.Jenner at amd.com>
Date: Thu, 4 Jun 2026 06:09:44 -0400
Subject: [PATCH 3/3] Use SgprB32_M0 instead of SgprP3 for
 amdgcn_struct_ptr_buffer_load_async_lds.

---
 llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
index fde656e7c15c9..8baac5f465d6b 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
@@ -1925,7 +1925,7 @@ RegBankLegalizeRules::RegBankLegalizeRules(const GCNSubtarget &_ST,
       .Any({{}, {{}, {IntrId, SgprP8, SgprB32_M0, Imm, VgprB32, SgprB32}}});
 
   addRulesForIOpcs({amdgcn_struct_ptr_buffer_load_async_lds})
-      .Any({{_}, {{}, {IntrId, SgprP8, SgprP3, Imm, Vgpr32, Vgpr32, Sgpr32}}});
+      .Any({{_}, {{}, {IntrId, SgprP8, SgprB32_M0, Imm, Vgpr32, Vgpr32, Sgpr32}}});
 
   addRulesForIOpcs({amdgcn_struct_ptr_buffer_load_lds})
       .Any({{_}, {{}, {IntrId, SgprP8, SgprP3, Imm, Vgpr32, Vgpr32, Sgpr32}}});