[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 ↗

(On Diff #457204)

-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();
+  const uint64_t IdxVal = MaybeIdxVal->Value.getZExtValue();
 
   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));
 
bool AMDGPULegalizerInfo::legalizeInsertVectorElt(
   // FIXME: Artifact combiner probably should have replaced the truncated
   // constant before this, so we shouldn't need
   // getIConstantVRegValWithLookThrough.
   Optional<ValueAndVReg> MaybeIdxVal =
       getIConstantVRegValWithLookThrough(MI.getOperand(3).getReg(), MRI);
   if (!MaybeIdxVal) // Dynamic case will be selected to register indexing.
     return true;
 
-  int64_t IdxVal = MaybeIdxVal->Value.getSExtValue();
+  const uint64_t IdxVal = MaybeIdxVal->Value.getZExtValue();
   Register Dst = MI.getOperand(0).getReg();
   Register Vec = MI.getOperand(1).getReg();
   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) {
     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.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=amdgcn -run-pass=legalizer -verify-machineinstrs -o - %s | FileCheck %s
+
+---
+name: f
+body: |
+  bb.0:
+    ; CHECK-LABEL: name: f
+    ; CHECK: SI_RETURN
+    %0:_(<2 x s32>) = COPY $vgpr0_vgpr1
+
+    %2:_(s32) = G_CONSTANT i32 -1
+    %3:_(s32) = G_EXTRACT_VECTOR_ELT %0(<2 x s32>), %2(s32)
+    %4:_(<2 x s32>) = G_INSERT_VECTOR_ELT %0(<2 x s32>), %3(s32), %2(s32)
+
+    %5:_(s32) = G_CONSTANT i32 2
+    %6:_(s32) = G_EXTRACT_VECTOR_ELT %0(<2 x s32>), %5(s32)
+    %7:_(<2 x s32>) = G_INSERT_VECTOR_ELT %0(<2 x s32>), %6(s32), %5(s32)
+
+    %8:_(s1) = G_CONSTANT i1 1
+    %9:_(s32) = G_EXTRACT_VECTOR_ELT %0(<2 x s32>), %8(s1)
+    %10:_(<2 x s32>) = G_INSERT_VECTOR_ELT %0(<2 x s32>), %9(s32), %8(s1)
+
+    SI_RETURN
+...