[AMDGPU] Fix crash legalizing G_EXTRACT_VECTOR_ELT with negative index

[foad]

Summary

Fixes https://github.com/llvm/llvm-project/issues/57408

[Petar.Avramovic]

IdxVal >= 0 check looks fine to me,
but what looks wrong in this test is that i1 1(true in llvm ir) index is translated to %7:_(s32) = G_CONSTANT i32 -1
IRTranslator::translateExtractElement did not expect i1 indices and since indices are (afaik) positive zext would make more sense then sext in:
APInt NewIdx = CI->getValue().sextOrTrunc(PreferredVecIdxWidth);
->
APInt NewIdx = CI->getValue().zextOrTrunc(PreferredVecIdxWidth);

Language ref does not specify required type for index, although we probably expected i32,
vector in test %I has two elements and i1 is enough to hold index, so globalisel should probably treat i1 true and false as 0 and 1 i32 indices .

[foad]

Agreed, indexes should be unsigned, and this should be documented in the LangRef like it is for the shl instruction: ‘op2’ is treated as an unsigned value.

[Peter]

Agree. This case seems unique to i1. Any other index with negative constant value would be ignored by translator.

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

2392

You might want to add the non-negative test here too. However, I haven't been able to find a test case that crashed here yet.

[Peter]

I'll send an issue to github and tag this post.

[foad]

Speculative fix for G_INSERT_VECTOR_ELT.

[arsenm]

I'd somewhat prefer a mir test for this in case some IR optimization decided to start deleting these earlier

In D132938#3758431, @Petar.Avramovic wrote:

IdxVal >= 0 check looks fine to me,
but what looks wrong in this test is that i1 1(true in llvm ir) index is translated to %7:_(s32) = G_CONSTANT i32 -1
IRTranslator::translateExtractElement did not expect i1 indices and since indices are (afaik) positive zext would make more sense then sext in:
APInt NewIdx = CI->getValue().sextOrTrunc(PreferredVecIdxWidth);
->
APInt NewIdx = CI->getValue().zextOrTrunc(PreferredVecIdxWidth);

Language ref does not specify required type for index, although we probably expected i32,
vector in test %I has two elements and i1 is enough to hold index, so globalisel should probably treat i1 true and false as 0 and 1 i32 indices .

The type conversion in the IRTranslator surprises me. I would expect to pass through the IR type and let the legalizer deal with it

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

2346

Why not switch this to getZExtValue?

llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-extractelement-crash.ll

2

-global-isel to the front

[foad]

Switch to MIR test case.

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

bool AMDGPULegalizerInfo::legalizeExtractVectorElt(
 
   // FIXME: Artifact combiner probably should have replaced the truncated
   // constant before this, so we shouldn't need
   // getIConstantVRegValWithLookThrough.
   Optional<ValueAndVReg> MaybeIdxVal =
       getIConstantVRegValWithLookThrough(MI.getOperand(2).getReg(), MRI);
   if (!MaybeIdxVal) // Dynamic case will be selected to register indexing.
     return true;
   const int64_t IdxVal = MaybeIdxVal->Value.getSExtValue();
 
   Register Dst = MI.getOperand(0).getReg();
   Register Vec = MI.getOperand(1).getReg();
 
   LLT VecTy = MRI.getType(Vec);
   LLT EltTy = VecTy.getElementType();
   assert(EltTy == MRI.getType(Dst));
 
-  if (IdxVal < VecTy.getNumElements()) {
+  if (IdxVal >= 0 && IdxVal < VecTy.getNumElements()) {
     auto Unmerge = B.buildUnmerge(EltTy, Vec);
     B.buildCopy(Dst, Unmerge.getReg(IdxVal));
   } else {
     B.buildUndef(Dst);
   }
 
   MI.eraseFromParent();
   return true;
bool AMDGPULegalizerInfo::legalizeInsertVectorElt(
   Register Ins = MI.getOperand(2).getReg();
 
   LLT VecTy = MRI.getType(Vec);
   LLT EltTy = VecTy.getElementType();
   assert(EltTy == MRI.getType(Ins));
   (void)Ins;
 
   unsigned NumElts = VecTy.getNumElements();
-  if (IdxVal < NumElts) {
+  if (IdxVal >= 0 && IdxVal < NumElts) {
     SmallVector<Register, 8> SrcRegs;
     for (unsigned i = 0; i < NumElts; ++i)
       SrcRegs.push_back(MRI.createGenericVirtualRegister(EltTy));
     B.buildUnmerge(SrcRegs, Vec);
 
     SrcRegs[IdxVal] = MI.getOperand(2).getReg();
     B.buildMerge(Dst, SrcRegs);
   } else {

llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-extractelement-crash.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=amdgcn -global-isel < %s | FileCheck %s
+
+define void @f() {
+; CHECK-LABEL: f:
+; CHECK:       ; %bb.0: ; %BB
+; CHECK-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; CHECK-NEXT:    s_setpc_b64 s[30:31]
+BB:
+  br label %BB2
+
+BB2:                                              ; preds = %BB
+  br label %BB1
+
+BB1:                                              ; preds = %BB2
+  %B1 = sub <2 x i8> <i8 16, i8 16>, <i8 16, i8 16>
+  %A = alloca i8, align 4, addrspace(5)
+  %L1 = load i8, i8 addrspace(5)* %A, align 1
+  %I = insertelement <2 x i8> %B1, i8 -128, i8 %L1
+  %E1 = extractelement <2 x i8> %I, i1 true
+  store i8 %E1, i8 addrspace(5)* %A, align 1
+  ret void
+}