Handle G_TRUNC on S1 SGPR

Author: gandhi56

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalize.cpp

@@ -477,14 +477,38 @@ bool AMDGPURegBankLegalize::runOnMachineFunction(MachineFunction &MF) {
       }
       if (MI.getOpcode() == AMDGPU::G_ANYEXT) {
         Combiner.tryCombineS1AnyExt(MI);
+        // Also handle anyext where src and dst have same type (nop anyext)
+        if (!MI.getParent())
+          continue; // Already erased by tryCombineS1AnyExt
+        Register Dst = MI.getOperand(0).getReg();
+        Register Src = MI.getOperand(1).getReg();
+        if (MRI.getType(Dst) == MRI.getType(Src)) {
+          while (!MRI.use_empty(Dst)) {
+            auto &Use = *MRI.use_begin(Dst);
+            Use.setReg(Src);
+          }
+          MI.eraseFromParent();
+        }
+        continue;
+      }
+      // Handle G_TRUNC that produces sgpr S1 - convert to sgpr S32
+      // Only eliminate if dead (no uses), otherwise legalization will handle it
+      if (MI.getOpcode() == AMDGPU::G_TRUNC) {
+        Register Dst = MI.getOperand(0).getReg();
+        if (MRI.getType(Dst) == LLT::scalar(1) && MRI.getRegBankOrNull(Dst) &&
+            MRI.getRegBankOrNull(Dst)->getID() == AMDGPU::SGPRRegBankID &&
+            MRI.use_empty(Dst)) {
+          // Dead sgpr S1 trunc, just erase it
+          MI.eraseFromParent();
+        }
         continue;
       }
     }
   }
 
-  // assert(!getAnySgprS1(MRI).isValid() &&
-  //        "Registers with sgpr reg bank and S1 LLT are not legal after "
-  //        "AMDGPURegBankLegalize. Should lower to sgpr S32");
+  assert(!getAnySgprS1(MRI).isValid() &&
+         "Registers with sgpr reg bank and S1 LLT are not legal after "
+         "AMDGPURegBankLegalize. Should lower to sgpr S32");
 
   return true;
 }

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeHelper.cpp

@@ -506,9 +506,13 @@ void RegBankLegalizeHelper::lowerScalarizeV2S16(MachineInstr &MI) {
   auto Op2 = B.buildUnmerge({SgprRB, S16}, MI.getOperand(2).getReg());
 
   // Perform scalar additions on S16 values
-  Register Lo = B.buildInstr(MI.getOpcode(), {SgprRB_S16}, {Op1.getReg(0), Op2.getReg(0)}).getReg(0);
-  Register Hi = B.buildInstr(MI.getOpcode(), {SgprRB_S16}, {Op1.getReg(1), Op2.getReg(1)}).getReg(0);
-  
+  Register Lo =
+      B.buildInstr(MI.getOpcode(), {SgprRB_S16}, {Op1.getReg(0), Op2.getReg(0)})
+          .getReg(0);
+  Register Hi =
+      B.buildInstr(MI.getOpcode(), {SgprRB_S16}, {Op1.getReg(1), Op2.getReg(1)})
+          .getReg(0);
+
   // Pack the results back into V2S16
   B.buildBuildVectorTrunc(MI.getOperand(0).getReg(), {Lo, Hi});
   MI.eraseFromParent();

llvm/test/CodeGen/AMDGPU/GlobalISel/regbankselect-add-overflow.s32.mir

@@ -14,7 +14,6 @@ body: |
     ; CHECK-NEXT: [[COPY:%[0-9]+]]:sgpr(s32) = COPY $sgpr0
     ; CHECK-NEXT: [[COPY1:%[0-9]+]]:sgpr(s32) = COPY $sgpr1
     ; CHECK-NEXT: [[UADDO:%[0-9]+]]:sgpr(s32), [[UADDO1:%[0-9]+]]:sgpr(s32) = G_UADDO [[COPY]], [[COPY1]]
-    ; CHECK-NEXT: [[TRUNC:%[0-9]+]]:sgpr(s1) = G_TRUNC [[UADDO1]](s32)
     %0:_(s32) = COPY $sgpr0
     %1:_(s32) = COPY $sgpr1
     %2:_(s32), %3:_(s1) = G_UADDO %0, %1
@@ -92,7 +91,6 @@ body: |
     ; CHECK-NEXT: [[C:%[0-9]+]]:sgpr(s32) = G_CONSTANT i32 1
     ; CHECK-NEXT: [[AND:%[0-9]+]]:sgpr(s32) = G_AND [[COPY2]], [[C]]
     ; CHECK-NEXT: [[UADDE:%[0-9]+]]:sgpr(s32), [[UADDE1:%[0-9]+]]:sgpr(s32) = G_UADDE [[COPY]], [[COPY1]], [[AND]]
-    ; CHECK-NEXT: [[TRUNC:%[0-9]+]]:sgpr(s1) = G_TRUNC [[UADDE1]](s32)
     %0:_(s32) = COPY $sgpr0
     %1:_(s32) = COPY $sgpr1
     %2:_(s32) = COPY $sgpr2

llvm/test/CodeGen/AMDGPU/GlobalISel/regbankselect-sub-overflow.s32.mir

@@ -14,7 +14,6 @@ body: |
     ; CHECK-NEXT: [[COPY:%[0-9]+]]:sgpr(s32) = COPY $sgpr0
     ; CHECK-NEXT: [[COPY1:%[0-9]+]]:sgpr(s32) = COPY $sgpr1
     ; CHECK-NEXT: [[USUBO:%[0-9]+]]:sgpr(s32), [[USUBO1:%[0-9]+]]:sgpr(s32) = G_USUBO [[COPY]], [[COPY1]]
-    ; CHECK-NEXT: [[TRUNC:%[0-9]+]]:sgpr(s1) = G_TRUNC [[USUBO1]](s32)
     %0:_(s32) = COPY $sgpr0
     %1:_(s32) = COPY $sgpr1
     %2:_(s32), %3:_(s1) = G_USUBO %0, %1
@@ -92,7 +91,6 @@ body: |
     ; CHECK-NEXT: [[C:%[0-9]+]]:sgpr(s32) = G_CONSTANT i32 1
     ; CHECK-NEXT: [[AND:%[0-9]+]]:sgpr(s32) = G_AND [[COPY2]], [[C]]
     ; CHECK-NEXT: [[USUBE:%[0-9]+]]:sgpr(s32), [[USUBE1:%[0-9]+]]:sgpr(s32) = G_USUBE [[COPY]], [[COPY1]], [[AND]]
-    ; CHECK-NEXT: [[TRUNC:%[0-9]+]]:sgpr(s1) = G_TRUNC [[USUBE1]](s32)
     %0:_(s32) = COPY $sgpr0
     %1:_(s32) = COPY $sgpr1
     %2:_(s32) = COPY $sgpr2