DAG: Fix assert on load casted to vector with attached range metadata

[arsenm]

Summary

AMDGPU legalizes i64 loads to loads of <2 x i32>, leaving the
i64 MMO with attached range metadata alone. The known bit width
was using the scalar element type, and asserting on a mismatch.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3463

is this correct for big-endian cases?

[arsenm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3463

Probably not. Ideally ranges would naturally support vectors and the integer could be casted into a vector of ranges, but they don't

[arsenm]

Check endianness for this hack

[RKSimon]

LGTM

[arsenm]

116c894d7296f5995c7d280ead32e4bfc821a18f

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

if (ISD::isNON_EXTLoad(LD) && Cst) {
         }
       }
     } else if (ISD::isZEXTLoad(Op.getNode()) && Op.getResNo() == 0) {
       // If this is a ZEXTLoad and we are looking at the loaded value.
       EVT VT = LD->getMemoryVT();
       unsigned MemBits = VT.getScalarSizeInBits();
       Known.Zero.setBitsFrom(MemBits);
     } else if (const MDNode *Ranges = LD->getRanges()) {
-      if (LD->getExtensionType() == ISD::NON_EXTLOAD)
+      EVT VT = LD->getValueType(0);
+
+      // TODO: Handle for extending loads
+      if (LD->getExtensionType() == ISD::NON_EXTLOAD) {
+        if (VT.isVector()) {
+          // Handle truncation to the first demanded element.
+          // TODO: Figure out which demanded elements are covered
+          if (DemandedElts != 1 || !getDataLayout().isLittleEndian())
+            break;
+
+          // Handle the case where a load has a vector type, but scalar memory
+          // with an attached range.
+          EVT MemVT = LD->getMemoryVT();
+          KnownBits KnownFull(MemVT.getSizeInBits());
+
+          computeKnownBitsFromRangeMetadata(*Ranges, KnownFull);
+          Known = KnownFull.trunc(BitWidth);
+        } else
           computeKnownBitsFromRangeMetadata(*Ranges, Known);
       }
+    }
     break;
   }
   case ISD::ZERO_EXTEND_VECTOR_INREG: {
     EVT InVT = Op.getOperand(0).getValueType();
     APInt InDemandedElts = DemandedElts.zext(InVT.getVectorNumElements());
     Known = computeKnownBits(Op.getOperand(0), InDemandedElts, Depth + 1);
     Known = Known.zext(BitWidth);
     break;

llvm/test/CodeGen/AMDGPU/load-range-metadata-assert.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -verify-machineinstrs < %s | FileCheck -enable-var-scope -check-prefix=GCN %s
+
+define <2 x i32> @range_metata_sext_range_0_i24_i64_bitcast(ptr addrspace(1) %ptr) {
+; GCN-LABEL: range_metata_sext_range_0_i24_i64_bitcast:
+; GCN:       ; %bb.0:
+; GCN-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; GCN-NEXT:    global_load_dwordx2 v[0:1], v[0:1], off glc
+; GCN-NEXT:    s_waitcnt vmcnt(0)
+; GCN-NEXT:    s_setpc_b64 s[30:31]
+  %val = load volatile i64, ptr addrspace(1) %ptr, align 4, !range !0
+  %shl = shl i64 %val, 22
+  %ashr = ashr i64 %shl, 22
+  %cast = bitcast i64 %ashr to <2 x i32>
+  ret <2 x i32> %cast
+}
+
+define <2 x i32> @no_range_sext_range_0_i24_i64_bitcast(ptr addrspace(1) %ptr) {
+; GCN-LABEL: no_range_sext_range_0_i24_i64_bitcast:
+; GCN:       ; %bb.0:
+; GCN-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; GCN-NEXT:    global_load_dwordx2 v[0:1], v[0:1], off glc
+; GCN-NEXT:    s_waitcnt vmcnt(0)
+; GCN-NEXT:    v_lshlrev_b64 v[0:1], 22, v[0:1]
+; GCN-NEXT:    v_ashrrev_i64 v[0:1], 22, v[0:1]
+; GCN-NEXT:    s_setpc_b64 s[30:31]
+  %val = load volatile i64, ptr addrspace(1) %ptr, align 4
+  %shl = shl i64 %val, 22
+  %ashr = ashr i64 %shl, 22
+  %cast = bitcast i64 %ashr to <2 x i32>
+  ret <2 x i32> %cast
+}
+
+!0 = !{i64 0, i64 16777216}