[DAGCombiner] Stop invalid sign conversion in refineIndexType.

[paulwalker-arm]

Summary

When looking through extends of gather/scatter indices it's safe
to convert a known positive signed index to unsigned, but unsigned
indices must remain unsigned.

Depends On D123318

[paulwalker-arm]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10496

Just wanted to say that I'd rather pass an IndexType directly here but I've resisted that because I have a follow-on patch that makes ISD::IndexType easier to work with and also removes refineIndexType's MSG operand so I figured I'd keep this bug fix minimal.

[paulwalker-arm]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10496

D123347 is the aforementioned follow-on patch.

[paulwalker-arm]

review ping

[kmclaughlin]

LGTM

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

bool refineUniformBase(SDValue &BasePtr, SDValue &Index, bool IndexIsScaled,
 
   BasePtr = SplatVal;
   Index = Index.getOperand(1);
   return true;
 }
 
 // Fold sext/zext of index into index type.
 bool refineIndexType(MaskedGatherScatterSDNode *MGS, SDValue &Index,
-                     bool Scaled, SelectionDAG &DAG) {
+                     bool Scaled, bool Signed, SelectionDAG &DAG) {
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
 
+  // It's always safe to look through zero extends.
   if (Index.getOpcode() == ISD::ZERO_EXTEND) {
     SDValue Op = Index.getOperand(0);
     MGS->setIndexType(Scaled ? ISD::UNSIGNED_SCALED : ISD::UNSIGNED_UNSCALED);
     if (TLI.shouldRemoveExtendFromGSIndex(Op.getValueType())) {
       Index = Op;
       return true;
     }
   }
 
-  if (Index.getOpcode() == ISD::SIGN_EXTEND) {
+  // It's only safe to look through sign extends when Index is signed.
+  if (Index.getOpcode() == ISD::SIGN_EXTEND && Signed) {
     SDValue Op = Index.getOperand(0);
     MGS->setIndexType(Scaled ? ISD::SIGNED_SCALED : ISD::SIGNED_UNSCALED);
     if (TLI.shouldRemoveExtendFromGSIndex(Op.getValueType())) {
       Index = Op;
       return true;
     }
   }
 
SDValue DAGCombiner::visitMSCATTER(SDNode *N) {
 
   if (refineUniformBase(BasePtr, Index, MSC->isIndexScaled(), DAG)) {
     SDValue Ops[] = {Chain, StoreVal, Mask, BasePtr, Index, Scale};
     return DAG.getMaskedScatter(
         DAG.getVTList(MVT::Other), MSC->getMemoryVT(), DL, Ops,
         MSC->getMemOperand(), MSC->getIndexType(), MSC->isTruncatingStore());
   }
 
-  if (refineIndexType(MSC, Index, MSC->isIndexScaled(), DAG)) {
+  if (refineIndexType(MSC, Index, MSC->isIndexScaled(), MSC->isIndexSigned(),
+                      DAG)) {
     SDValue Ops[] = {Chain, StoreVal, Mask, BasePtr, Index, Scale};
     return DAG.getMaskedScatter(
         DAG.getVTList(MVT::Other), MSC->getMemoryVT(), DL, Ops,
         MSC->getMemOperand(), MSC->getIndexType(), MSC->isTruncatingStore());
   }
 
   return SDValue();
 }
SDValue DAGCombiner::visitMGATHER(SDNode *N) {
   if (refineUniformBase(BasePtr, Index, MGT->isIndexScaled(), DAG)) {
     SDValue Ops[] = {Chain, PassThru, Mask, BasePtr, Index, Scale};
     return DAG.getMaskedGather(DAG.getVTList(N->getValueType(0), MVT::Other),
                                MGT->getMemoryVT(), DL, Ops,
                                MGT->getMemOperand(), MGT->getIndexType(),
                                MGT->getExtensionType());
   }
 
-  if (refineIndexType(MGT, Index, MGT->isIndexScaled(), DAG)) {
+  if (refineIndexType(MGT, Index, MGT->isIndexScaled(), MGT->isIndexSigned(),
+                      DAG)) {
     SDValue Ops[] = {Chain, PassThru, Mask, BasePtr, Index, Scale};
     return DAG.getMaskedGather(DAG.getVTList(N->getValueType(0), MVT::Other),
                                MGT->getMemoryVT(), DL, Ops,
                                MGT->getMemOperand(), MGT->getIndexType(),
                                MGT->getExtensionType());
   }
 
   return SDValue();

llvm/test/CodeGen/AArch64/sve-gather-scatter-addr-opts.ll

 ; CHECK-NEXT:    ld1w { z0.s }, p0/z, [x0, z0.s, uxtw #2]
 ; CHECK-NEXT:    ret
   %offsets.zext = zext <vscale x 4 x i8> %offsets to <vscale x 4 x i32>
   %ptrs = getelementptr i32, i32* %base, <vscale x 4 x i32> %offsets.zext
   %data = call <vscale x 4 x i32> @llvm.masked.gather.nxv4i32(<vscale x 4 x i32*> %ptrs, i32 4, <vscale x 4 x i1> %mask, <vscale x 4 x i32> undef)
   ret <vscale x 4 x i32> %data
 }
 
-; TODO: The generated code is wrong because we're replicating offset[31] across
-; offset[32:63] even though the IR has explicitly zero'd those bits.
 define <vscale x 4 x i32> @masked_gather_nxv4i32_u32s8_offsets(i32* %base, <vscale x 4 x i8> %offsets, <vscale x 4 x i1> %mask) #0 {
 ; CHECK-LABEL: masked_gather_nxv4i32_u32s8_offsets:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p1.s
 ; CHECK-NEXT:    sxtb z0.s, p1/m, z0.s
-; CHECK-NEXT:    ld1w { z0.s }, p0/z, [x0, z0.s, sxtw #2]
+; CHECK-NEXT:    ld1w { z0.s }, p0/z, [x0, z0.s, uxtw #2]
 ; CHECK-NEXT:    ret
   %offsets.sext = sext <vscale x 4 x i8> %offsets to <vscale x 4 x i32>
   %offsets.sext.zext = zext <vscale x 4 x i32> %offsets.sext to <vscale x 4 x i64>
   %ptrs = getelementptr i32, i32* %base, <vscale x 4 x i64> %offsets.sext.zext
   %data = call <vscale x 4 x i32> @llvm.masked.gather.nxv4i32(<vscale x 4 x i32*> %ptrs, i32 4, <vscale x 4 x i1> %mask, <vscale x 4 x i32> undef)
   ret <vscale x 4 x i32> %data
 }
 

 ; CHECK-NEXT:    st1w { z1.s }, p0, [x0, z0.s, uxtw #2]
 ; CHECK-NEXT:    ret
   %offsets.zext = zext <vscale x 4 x i8> %offsets to <vscale x 4 x i32>
   %ptrs = getelementptr i32, i32* %base, <vscale x 4 x i32> %offsets.zext
   call void @llvm.masked.scatter.nxv4i32(<vscale x 4 x i32> %data, <vscale x 4 x i32*> %ptrs, i32 4, <vscale x 4 x i1> %mask)
   ret void
 }
 
-; TODO: The generated code is wrong because we're replicating offset[31] across
-; offset[32:63] even though the IR has explicitly zero'd those bits.
 define void @masked_scatter_nxv4i32_u32s8_offsets(i32* %base, <vscale x 4 x i8> %offsets, <vscale x 4 x i1> %mask, <vscale x 4 x i32> %data) #0 {
 ; CHECK-LABEL: masked_scatter_nxv4i32_u32s8_offsets:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p1.s
 ; CHECK-NEXT:    sxtb z0.s, p1/m, z0.s
-; CHECK-NEXT:    st1w { z1.s }, p0, [x0, z0.s, sxtw #2]
+; CHECK-NEXT:    st1w { z1.s }, p0, [x0, z0.s, uxtw #2]
 ; CHECK-NEXT:    ret
   %offsets.sext = sext <vscale x 4 x i8> %offsets to <vscale x 4 x i32>
   %offsets.sext.zext = zext <vscale x 4 x i32> %offsets.sext to <vscale x 4 x i64>
   %ptrs = getelementptr i32, i32* %base, <vscale x 4 x i64> %offsets.sext.zext
   call void @llvm.masked.scatter.nxv4i32(<vscale x 4 x i32> %data, <vscale x 4 x i32*> %ptrs, i32 4, <vscale x 4 x i1> %mask)
   ret void
 }