[llvm] AMDGPU/GlobalISel: RegBankLegalize rules for tensor load/store to lds (PR #202363)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Petar Avramovic (petar-avramovic)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/202363.diff


4 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeHelper.cpp (+4-1) 
- (modified) llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp (+7) 
- (modified) llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.h (+1) 
- (modified) llvm/test/CodeGen/AMDGPU/llvm.amdgcn.tensor.load.store.ll (+1-1) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeHelper.cpp b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeHelper.cpp
index 3b29b52af02fb..26d2d2cb01f51 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeHelper.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeHelper.cpp
@@ -1534,6 +1534,7 @@ LLT RegBankLegalizeHelper::getTyFromID(RegBankLLTMappingApplyID ID) {
     return LLT::fixed_vector(4, 32);
   case VgprV8S32:
   case UniInVgprV8S32:
+  case SgprV8S32_ReadFirstLane:
     return LLT::fixed_vector(8, 32);
   case VgprV2S64:
   case UniInVgprV2S64:
@@ -1664,6 +1665,7 @@ RegBankLegalizeHelper::getRegBankFromID(RegBankLLTMappingApplyID ID) {
   case SgprV4S32:
   case SgprV4S32_WF:
   case SgprV4S32_ReadFirstLane:
+  case SgprV8S32_ReadFirstLane:
   case SgprB32:
   case SgprB64:
   case SgprB96:
@@ -2104,7 +2106,8 @@ bool RegBankLegalizeHelper::applyMappingSrc(
       Op.setReg(NewSGPR);
       break;
     }
-    case SgprV4S32_ReadFirstLane: {
+    case SgprV4S32_ReadFirstLane:
+    case SgprV8S32_ReadFirstLane: {
       assert(Ty == getTyFromID(MethodIDs[i]));
       if (RB == SgprRB)
         break;
diff --git a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
index 86eb3217fa996..5d84eac8c6d33 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp
@@ -2149,6 +2149,13 @@ RegBankLegalizeRules::RegBankLegalizeRules(const GCNSubtarget &_ST,
                     amdgcn_global_store_async_from_lds_b128})
       .Any({{}, {{}, {IntrId, VgprP1, VgprP3}}});
 
+  addRulesForIOpcs({amdgcn_tensor_load_to_lds, amdgcn_tensor_store_from_lds})
+      .Any({{},
+            {{},
+             {IntrId, SgprV4S32_ReadFirstLane, SgprV8S32_ReadFirstLane,
+              SgprV4S32_ReadFirstLane, SgprV4S32_ReadFirstLane,
+              SgprV8S32_ReadFirstLane}}});
+
   addRulesForIOpcs({amdgcn_cluster_load_b32})
       .Any({{UniB32}, {{UniInVgprB32}, {IntrId, SgprP1, Imm, SgprB32_M0}}})
       .Any({{DivB32, _, UniP1}, {{VgprB32}, {IntrId, SgprP1, Imm, SgprB32_M0}}})
diff --git a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.h b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.h
index d9b1dd6289376..bce8356528f92 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.h
@@ -279,6 +279,7 @@ enum RegBankLLTMappingApplyID {
   SgprB32_ReadFirstLane,
   SgprB64_ReadFirstLane,
   SgprV4S32_ReadFirstLane,
+  SgprV8S32_ReadFirstLane,
 
   // Src only modifiers: extends
   Sgpr32AExt,
diff --git a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.tensor.load.store.ll b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.tensor.load.store.ll
index 9f30255a07095..5a878ab14e46e 100644
--- a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.tensor.load.store.ll
+++ b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.tensor.load.store.ll
@@ -1,6 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -global-isel=0 -mtriple=amdgcn -mcpu=gfx1250 < %s | FileCheck -check-prefixes=GFX1250,GFX1250-SDAG %s
-; RUN: llc -global-isel=1 -mtriple=amdgcn -mcpu=gfx1250 < %s | FileCheck -check-prefixes=GFX1250,GFX1250-GISEL %s
+; RUN: llc -global-isel=1 -new-reg-bank-select -mtriple=amdgcn -mcpu=gfx1250 < %s | FileCheck -check-prefixes=GFX1250,GFX1250-GISEL %s
 
 ; %D4 should be zero-initialized for gfx1250, which only supports 4 groups of tensor descriptor
 declare void @llvm.amdgcn.tensor.load.to.lds(<4 x i32> %D0, <8 x i32> %D1, <4 x i32> %D2, <4 x i32> %D3, <8 x i32> %D4, i32 %cpol)

``````````

</details>


https://github.com/llvm/llvm-project/pull/202363