AMDGPU/GlobalISel: Lower constant-32-bit zextload/sextload consistently

[arsenm]

Summary

We were accidentally leaning on code in lowerLoad which expands
extending loads which should be removed.

[foad]

Looks reasonable.

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

2416

Couldn't this be an assert?

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

2416

No, sextload and zextload should theoretically be handled here too

[arsenm]

ee359000898c40ada69db3f20a87e6424c23596e

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

auto &ExtLoads = getActionDefinitionsBuilder({G_SEXTLOAD, G_ZEXTLOAD})
                                                   {S32, PrivatePtr, 16, 16},
                                                   {S32, ConstantPtr, 8, 8},
                                                   {S32, ConstantPtr, 16, 2 * 8}});
   if (ST.hasFlatAddressSpace()) {
     ExtLoads.legalForTypesWithMemDesc(
         {{S32, FlatPtr, 8, 8}, {S32, FlatPtr, 16, 16}});
   }
 
+  // Constant 32-bit is handled by addrspacecasting the 32-bit pointer to
+  // 64-bits.
+  //
+  // TODO: Should generalize bitcast action into coerce, which will also cover
+  // inserting addrspacecasts.
+  ExtLoads.customIf(typeIs(1, Constant32Ptr));
+
   ExtLoads.clampScalar(0, S32, S32)
           .widenScalarToNextPow2(0)
           .unsupportedIfMemSizeNotPow2()
           .lower();
 
   auto &Atomics = getActionDefinitionsBuilder(
     {G_ATOMICRMW_XCHG, G_ATOMICRMW_ADD, G_ATOMICRMW_SUB,
      G_ATOMICRMW_AND, G_ATOMICRMW_OR, G_ATOMICRMW_XOR,
bool AMDGPULegalizerInfo::legalizeCustom(LegalizerHelper &Helper,
   case TargetOpcode::G_SHUFFLE_VECTOR:
     return legalizeShuffleVector(MI, MRI, B);
   case TargetOpcode::G_FSIN:
   case TargetOpcode::G_FCOS:
     return legalizeSinCos(MI, MRI, B);
   case TargetOpcode::G_GLOBAL_VALUE:
     return legalizeGlobalValue(MI, MRI, B);
   case TargetOpcode::G_LOAD:
+  case TargetOpcode::G_SEXTLOAD:
+  case TargetOpcode::G_ZEXTLOAD:
     return legalizeLoad(Helper, MI);
   case TargetOpcode::G_FMAD:
     return legalizeFMad(MI, MRI, B);
   case TargetOpcode::G_FDIV:
     return legalizeFDIV(MI, MRI, B);
   case TargetOpcode::G_UDIV:
   case TargetOpcode::G_UREM:
     return legalizeUDIV_UREM(MI, MRI, B);
if (AddrSpace == AMDGPUAS::CONSTANT_ADDRESS_32BIT) {
     LLT ConstPtr = LLT::pointer(AMDGPUAS::CONSTANT_ADDRESS, 64);
     auto Cast = B.buildAddrSpaceCast(ConstPtr, PtrReg);
     Observer.changingInstr(MI);
     MI.getOperand(1).setReg(Cast.getReg(0));
     Observer.changedInstr(MI);
     return true;
   }
 
+  if (MI.getOpcode() != AMDGPU::G_LOAD)
+    return false;
+
   Register ValReg = MI.getOperand(0).getReg();
   LLT ValTy = MRI.getType(ValReg);
 
   MachineMemOperand *MMO = *MI.memoperands_begin();
   const unsigned ValSize = ValTy.getSizeInBits();
   const unsigned MemSize = 8 * MMO->getSize();
   const Align MemAlign = MMO->getAlign();
   const unsigned AlignInBits = 8 * MemAlign.value();

llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-sextload-constant-32bit.mir

 body: |
   bb.0:
     liveins: $sgpr0
 
     ; CI-LABEL: name: test_sextload_constant32bit_s32_s8_align1
     ; CI: [[COPY:%[0-9]+]]:_(p6) = COPY $sgpr0
     ; CI: [[C:%[0-9]+]]:_(p6) = G_CONSTANT i32 0
     ; CI: [[MV:%[0-9]+]]:_(p4) = G_MERGE_VALUES [[COPY]](p6), [[C]](p6)
-    ; CI: [[LOAD:%[0-9]+]]:_(s32) = G_LOAD [[MV]](p4) :: (load 1, addrspace 6)
-    ; CI: [[COPY1:%[0-9]+]]:_(s32) = COPY [[LOAD]](s32)
-    ; CI: [[SEXT_INREG:%[0-9]+]]:_(s32) = G_SEXT_INREG [[COPY1]], 8
-    ; CI: $vgpr0 = COPY [[SEXT_INREG]](s32)
+    ; CI: [[SEXTLOAD:%[0-9]+]]:_(s32) = G_SEXTLOAD [[MV]](p4) :: (load 1, addrspace 6)
+    ; CI: $vgpr0 = COPY [[SEXTLOAD]](s32)
     %0:_(p6) = COPY $sgpr0
     %1:_(s32) = G_SEXTLOAD %0 :: (load 1, align 1, addrspace 6)
     $vgpr0 = COPY %1
 ...
 
 ---
 name: test_sextload_constant32bit_s32_s16_align2
 body: |
   bb.0:
     liveins: $sgpr0
 
     ; CI-LABEL: name: test_sextload_constant32bit_s32_s16_align2
     ; CI: [[COPY:%[0-9]+]]:_(p6) = COPY $sgpr0
     ; CI: [[C:%[0-9]+]]:_(p6) = G_CONSTANT i32 0
     ; CI: [[MV:%[0-9]+]]:_(p4) = G_MERGE_VALUES [[COPY]](p6), [[C]](p6)
-    ; CI: [[LOAD:%[0-9]+]]:_(s32) = G_LOAD [[MV]](p4) :: (load 2, addrspace 6)
-    ; CI: [[COPY1:%[0-9]+]]:_(s32) = COPY [[LOAD]](s32)
-    ; CI: [[SEXT_INREG:%[0-9]+]]:_(s32) = G_SEXT_INREG [[COPY1]], 16
-    ; CI: $vgpr0 = COPY [[SEXT_INREG]](s32)
+    ; CI: [[SEXTLOAD:%[0-9]+]]:_(s32) = G_SEXTLOAD [[MV]](p4) :: (load 2, addrspace 6)
+    ; CI: $vgpr0 = COPY [[SEXTLOAD]](s32)
     %0:_(p6) = COPY $sgpr0
     %1:_(s32) = G_SEXTLOAD %0 :: (load 2, align 2, addrspace 6)
     $vgpr0 = COPY %1
 ...
 
 ---
 name: test_sextload_constant32bit_s32_s16_align1
 body: |

llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-zextload-constant-32bit.mir

 body: |
   bb.0:
     liveins: $sgpr0
 
     ; CI-LABEL: name: test_zextload_constant32bit_s32_s8_align1
     ; CI: [[COPY:%[0-9]+]]:_(p6) = COPY $sgpr0
     ; CI: [[C:%[0-9]+]]:_(p6) = G_CONSTANT i32 0
     ; CI: [[MV:%[0-9]+]]:_(p4) = G_MERGE_VALUES [[COPY]](p6), [[C]](p6)
-    ; CI: [[LOAD:%[0-9]+]]:_(s32) = G_LOAD [[MV]](p4) :: (load 1, addrspace 6)
-    ; CI: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 255
-    ; CI: [[COPY1:%[0-9]+]]:_(s32) = COPY [[LOAD]](s32)
-    ; CI: [[AND:%[0-9]+]]:_(s32) = G_AND [[COPY1]], [[C1]]
-    ; CI: $vgpr0 = COPY [[AND]](s32)
+    ; CI: [[ZEXTLOAD:%[0-9]+]]:_(s32) = G_ZEXTLOAD [[MV]](p4) :: (load 1, addrspace 6)
+    ; CI: $vgpr0 = COPY [[ZEXTLOAD]](s32)
     %0:_(p6) = COPY $sgpr0
     %1:_(s32) = G_ZEXTLOAD %0 :: (load 1, align 1, addrspace 6)
     $vgpr0 = COPY %1
 ...
 
 ---
 name: test_zextload_constant32bit_s32_s16_align2
 body: |
   bb.0:
     liveins: $sgpr0
 
     ; CI-LABEL: name: test_zextload_constant32bit_s32_s16_align2
     ; CI: [[COPY:%[0-9]+]]:_(p6) = COPY $sgpr0
     ; CI: [[C:%[0-9]+]]:_(p6) = G_CONSTANT i32 0
     ; CI: [[MV:%[0-9]+]]:_(p4) = G_MERGE_VALUES [[COPY]](p6), [[C]](p6)
-    ; CI: [[LOAD:%[0-9]+]]:_(s32) = G_LOAD [[MV]](p4) :: (load 2, addrspace 6)
-    ; CI: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 65535
-    ; CI: [[COPY1:%[0-9]+]]:_(s32) = COPY [[LOAD]](s32)
-    ; CI: [[AND:%[0-9]+]]:_(s32) = G_AND [[COPY1]], [[C1]]
-    ; CI: $vgpr0 = COPY [[AND]](s32)
+    ; CI: [[ZEXTLOAD:%[0-9]+]]:_(s32) = G_ZEXTLOAD [[MV]](p4) :: (load 2, addrspace 6)
+    ; CI: $vgpr0 = COPY [[ZEXTLOAD]](s32)
     %0:_(p6) = COPY $sgpr0
     %1:_(s32) = G_ZEXTLOAD %0 :: (load 2, align 2, addrspace 6)
     $vgpr0 = COPY %1
 ...
 
 ---
 name: test_zextload_constant32bit_s32_s16_align1
 body: |