[CodeGen] Using ZExt for extractelement indices.

[Peter]

Summary

In https://github.com/llvm/llvm-project/issues/57452, we found that IRTranslator is translating i1 true into i32 -1.
This is because IRTranslator uses SExt for indices.
In this fix, we move from SExt to ZExt. We also included a test for AMDGPU, an updated test for AArch64

This patch fixes this issue.

[foad]

Can you add some text to the LangRef to clarify that insert/extractelement indices are treated as unsigned?

By inspection, SelectionDAG seems to have the same bug. See SelectionDAGBuilder::visitInsertElement and SelectionDAGBuilder::visitExtractElement. But perhaps it is harmless there for some reason.

[Peter]

If you can point me where the doc is located, I will make the change.

As of SelectionDAGBuilder::visitInsertElement and SelectionDAGBuilder::visitExtractElement, we will change them so its compliant with the documentation.

Please allow me a few days before I submit another patch, it's been busy days for me.

[foad]

If you can point me where the doc is located, I will make the change.

docs/LangRef.rst in the LLVM repo.

Please allow me a few days before I submit another patch, it's been busy days for me.

No problem!

[Peter]

updated the documentation so it is complient with our code

[foad]

llvm/docs/LangRef.rst

9694–9695 ↗

(On Diff #459203)

From a semantics point of view, I think all we need to say is that the value is "treated as unsigned" or words to that effect. I don't think we should mention extending here (and definitely not truncating which sounds like it could change the value!).

[Peter]

llvm/docs/LangRef.rst

9694–9695 ↗

(On Diff #459203)

https://llvm.godbolt.org/z/Mszb7EvPe

The (unfortunate) fact that we are using zextOrTrunc already implies that we can change the value given a large index, even thou no one is doing it.
I see two ways out of this:

1. Explain in detail in the doc "when" truncation will happen or
2. Further change the doc so the index can only be i64, which would involve even more code change.

[foad]

llvm/docs/LangRef.rst

9694–9695 ↗

(On Diff #459203)

The document should describe the intended semantics of the IR, not the details of the current GlobalISel implementation.

(You're correct that GlobalISel does odd things if the index type is wider than 64 bits, but I don't think it's really *wrong*. The only effect it can have is to change an out-of-range index into an in-range index. And if the index was out of range in the first place then the result was undefined anyway (or poison) so it doesn't really matter what GlobalISel does with it.)

[tschuett]

Currently, you say what to do: zext or truncate. I thought that you should say what it is: an unsigned integer with up to 64 bit.

[Peter]

What about we say: "The index may be a variable of any integer type, and will be treated as an unsigned integer."

If that looks good I think I will update it and fix all the testcases sometime next week.

[foad]

In D132978#3787383, @Peter wrote:

What about we say: "The index may be a variable of any integer type, and will be treated as an unsigned integer."

That works for me!

[Peter]

Update tests. All passing.

[Peter]

update doc

[Peter]

Many archs' tests were failing due to this change. I have updated all the test. Most of them can be done using a script.
Still, do we want to include some arch experts to review the test changes?

[foad]

[IRTranslator] Using ZExt for extractelement indices.

Maybe change this to [CodeGen] now that SelectionDAG changes are included?

In D132978#3816171, @Peter wrote:

Many archs' tests were failing due to this change. I have updated all the test. Most of them can be done using a script.
Still, do we want to include some arch experts to review the test changes?

Yes I think this needs some arch-specific review, especially since the code looks worse (longer) in some cases. @craig.topper @RKSimon @t.p.northover @nemanjai @kaz7 @sunfish @theraven

[Petar.Avramovic]

I checked a few code size regressions and they look correct. Combines would replace extract/insert with undef because index was i32 -1 or i64 -1 (from i1 true). Now they extract/insert element at index 1.
My best guess how this bug was not discovered sooner is that frontends use i32/i64 for index type.
Other type of regression is when i32 (marked as signed argument) has to be zero-extended to i64 (and with mask), before it used to be a copy. Considering langref changes it is probably more correct to mark index arguments as zeroext (index will be copied again no code size regression). But this should be checked with people who know more about specific targets.

[RKSimon]

llvm/test/CodeGen/X86/var-permute-128.ll

132 ↗

(On Diff #463018)

These are mildly annoying, on some targets movd is notably quicker than pextr* - I think we're just missing a SimplifyDemandedBits handler though.

[pengfei]

llvm/test/CodeGen/X86/vec_extract.ll

114 ↗

(On Diff #463018)

This seems redundant.

[RKSimon]

llvm/test/CodeGen/X86/vec_extract.ll

114 ↗

(On Diff #463018)

D123394 /might/ address this - it handles something similar in a couple of places.

[pengfei]

llvm/test/CodeGen/X86/vec_extract.ll

114 ↗

(On Diff #463018)

That's great :)

[Peter]

llvm/test/CodeGen/X86/var-permute-128.ll

132 ↗

(On Diff #463018)

I just checked the documentation of both instructions. It seems you are right.
Is there anything you hoped we could've done differently, or it is another issue?

[frasercrmck]

In D132978#3817541, @Petar.Avramovic wrote:

Other type of regression is when i32 (marked as signed argument) has to be zero-extended to i64 (and with mask), before it used to be a copy. Considering langref changes it is probably more correct to mark index arguments as zeroext (index will be copied again no code size regression). But this should be checked with people who know more about specific targets.

For RISCV I think that the indices should indeed be marked zext. Testing sext doesn't get us a whole lot IMO. @craig.topper, @reames?

[reames]

As a data point in support of the unsigned semantics, InstComine already uses zext to canonicalize index type (see getPreferredVectorIndex). Several of the transforms use unsigned comparison. So, at IR level, we seem to uniformly be expecting unsigned behavior.

On the RISCV changes, several of these look like regressions, but it comes down to the fact that arguments in the tests are already sign extended, so the sext is cheaper than the zext. One thing I note - which may be generally useful - is that anything above i32 (i.e. where the zext vs sext difference is observable in these tests) is definitely out of bounds. Since out of bounds returns poison, I believe we can use an any extend here. This might reduce the impact.

[arsenm]

llvm/test/CodeGen/PowerPC/aix-vec_extract_p9.ll

104 ↗

(On Diff #463018)

Are we missing a SimplifyDemandedBits combine on the extract index?

[efriedma]

Since out of bounds returns poison, I believe we can use an any extend here.

I think any_extend is too aggressive. But we could check isSExtCheaperThanZExt(), sure.

[Peter]

Should we merge this or is there anything else that needs to be addressed?

[arsenm]

In D132978#3828158, @Peter wrote:

Should we merge this or is there anything else that needs to be addressed?

The regressions don't look particularly important, but there may be a missing SimplifyDemandedBits style optimization on the index

[RKSimon]

llvm/test/CodeGen/PowerPC/variable_elem_vec_extracts.ll

1 ↗

(On Diff #463018)

Please can you pre-commit the autogen change and then rebase so the patch shows the codegen diff?

[Peter]

Rebase with upstream and changed tests in llvm/test/CodeGen/PowerPC/variable_elem_vec_extracts.ll

[Peter]

llvm/test/CodeGen/PowerPC/variable_elem_vec_extracts.ll

9 ↗

(On Diff #467255)

I am wondering which part of the code, if any, in the front end is responsible for generating these attributes. I feel like we should generate zeroext for indices instead of signext due to the document change.

[Peter]

rebase to upstream.

[pengfei]

llvm/test/CodeGen/PowerPC/variable_elem_vec_extracts.ll

9 ↗

(On Diff #467255)

I think it's CodeGenModule::ConstructAttributeList

[craig.topper]

llvm/test/CodeGen/PowerPC/variable_elem_vec_extracts.ll

9 ↗

(On Diff #467255)

The attributes tend to follow the types of the arguments in the C source code. So if it was a unsigned type it would be zeroext, if it was a signed type it would be signext. It's not based on how it is used in the function.

[Peter]

llvm/test/CodeGen/PowerPC/variable_elem_vec_extracts.ll

9 ↗

(On Diff #467255)

In that case I don't have to change anything. BYW thanks pengfei for pointing out, I just reviewed that code.

[Peter]

Thanks for approving the diff, I'll let it sit for another day to see if anyone has more comments and merging it tomorrow.

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

bool IRTranslator::translateExtractElement(const User &U,
 
   Register Res = getOrCreateVReg(U);
   Register Val = getOrCreateVReg(*U.getOperand(0));
   const auto &TLI = *MF->getSubtarget().getTargetLowering();
   unsigned PreferredVecIdxWidth = TLI.getVectorIdxTy(*DL).getSizeInBits();
   Register Idx;
   if (auto *CI = dyn_cast<ConstantInt>(U.getOperand(1))) {
     if (CI->getBitWidth() != PreferredVecIdxWidth) {
-      APInt NewIdx = CI->getValue().sextOrTrunc(PreferredVecIdxWidth);
+      APInt NewIdx = CI->getValue().zextOrTrunc(PreferredVecIdxWidth);
       auto *NewIdxCI = ConstantInt::get(CI->getContext(), NewIdx);
       Idx = getOrCreateVReg(*NewIdxCI);
     }
   }
   if (!Idx)
     Idx = getOrCreateVReg(*U.getOperand(1));
   if (MRI->getType(Idx).getSizeInBits() != PreferredVecIdxWidth) {
     const LLT VecIdxTy = LLT::scalar(PreferredVecIdxWidth);
-    Idx = MIRBuilder.buildSExtOrTrunc(VecIdxTy, Idx).getReg(0);
+    Idx = MIRBuilder.buildZExtOrTrunc(VecIdxTy, Idx).getReg(0);
   }
   MIRBuilder.buildExtractVectorElement(Res, Val, Idx);
   return true;
 }
 
 bool IRTranslator::translateShuffleVector(const User &U,
                                           MachineIRBuilder &MIRBuilder) {
   ArrayRef<int> Mask;

llvm/test/CodeGen/AArch64/GlobalISel/arm64-irtranslator.ll

; CHECK: $d0 = COPY [[RES]](<2 x s32>)
   %res = insertelement <2 x i32> %vec, i32 %elt, i32 %idx
   ret <2 x i32> %res
 }
 
 define i32 @test_extractelement(<2 x i32> %vec, i32 %idx) {
 ; CHECK-LABEL: name: test_extractelement
 ; CHECK: [[VEC:%[0-9]+]]:_(<2 x s32>) = COPY $d0
 ; CHECK: [[IDX:%[0-9]+]]:_(s32) = COPY $w0
-; CHECK: [[IDXEXT:%[0-9]+]]:_(s64) = G_SEXT [[IDX]]
+; CHECK: [[IDXEXT:%[0-9]+]]:_(s64) = G_ZEXT [[IDX]]
 ; CHECK: [[RES:%[0-9]+]]:_(s32) = G_EXTRACT_VECTOR_ELT [[VEC]](<2 x s32>), [[IDXEXT]](s64)
 ; CHECK: $w0 = COPY [[RES]](s32)
   %res = extractelement <2 x i32> %vec, i32 %idx
   ret i32 %res
 }
 
 define i32 @test_extractelement_const_idx(<2 x i32> %vec) {
 ; CHECK-LABEL: name: test_extractelement
 ; CHECK: [[VEC:%[0-9]+]]:_(<2 x s32>) = COPY $d0
 ; CHECK: [[IDX:%[0-9]+]]:_(s64) = G_CONSTANT i64 1
 ; CHECK: [[RES:%[0-9]+]]:_(s32) = G_EXTRACT_VECTOR_ELT [[VEC]](<2 x s32>), [[IDX]](s64)
 ; CHECK: $w0 = COPY [[RES]](s32)
   %res = extractelement <2 x i32> %vec, i32 1
   ret i32 %res
 }
 
+define i32 @test_extractelement_const_idx_zext_i1(<2 x i32> %vec) {
+; CHECK-LABEL: name: test_extractelement
+; CHECK: [[VEC:%[0-9]+]]:_(<2 x s32>) = COPY $d0
+; CHECK: [[IDX:%[0-9]+]]:_(s64) = G_CONSTANT i64 1
+; CHECK: [[RES:%[0-9]+]]:_(s32) = G_EXTRACT_VECTOR_ELT [[VEC]](<2 x s32>), [[IDX]](s64)
+; CHECK: $w0 = COPY [[RES]](s32)
+  %res = extractelement <2 x i32> %vec, i1 true
+  ret i32 %res
+}
+
+define i32 @test_extractelement_const_idx_zext_i8(<2 x i32> %vec) {
+; CHECK-LABEL: name: test_extractelement
+; CHECK: [[VEC:%[0-9]+]]:_(<2 x s32>) = COPY $d0
+; CHECK: [[IDX:%[0-9]+]]:_(s64) = G_CONSTANT i64 255
+; CHECK: [[RES:%[0-9]+]]:_(s32) = G_EXTRACT_VECTOR_ELT [[VEC]](<2 x s32>), [[IDX]](s64)
+; CHECK: $w0 = COPY [[RES]](s32)
+  %res = extractelement <2 x i32> %vec, i8 255
+  ret i32 %res
+}
+
+
 define i32 @test_singleelementvector(i32 %elt){
 ; CHECK-LABEL: name: test_singleelementvector
 ; CHECK: [[ELT:%[0-9]+]]:_(s32) = COPY $w0
 ; CHECK-NOT: G_INSERT_VECTOR_ELT
 ; CHECK-NOT: G_EXTRACT_VECTOR_ELT
 ; CHECK: $w0 = COPY [[ELT]](s32)
   %vec = insertelement <1 x i32> undef, i32 %elt, i32 0
   %res = extractelement <1 x i32> %vec, i32 0

llvm/test/CodeGen/AMDGPU/GlobalISel/irtranslator-zext-vec-index.ll

+; NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+; RUN: llc -march=amdgcn -O0 -stop-after=irtranslator -global-isel -verify-machineinstrs %s -o - | FileCheck %s
+
+define i8 @f_i1_1() {
+  ; CHECK-LABEL: name: f_i1_1
+  ; CHECK: bb.1 (%ir-block.0):
+  ; CHECK-NEXT:   [[DEF:%[0-9]+]]:_(<256 x s8>) = G_IMPLICIT_DEF
+  ; CHECK-NEXT:   [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+  ; CHECK-NEXT:   [[EVEC:%[0-9]+]]:_(s8) = G_EXTRACT_VECTOR_ELT [[DEF]](<256 x s8>), [[C]](s32)
+  ; CHECK-NEXT:   [[ANYEXT:%[0-9]+]]:_(s32) = G_ANYEXT [[EVEC]](s8)
+  ; CHECK-NEXT:   $vgpr0 = COPY [[ANYEXT]](s32)
+  ; CHECK-NEXT:   SI_RETURN implicit $vgpr0
+  %E1 = extractelement <256 x i8> undef, i1 true
+  ret i8 %E1
+}
+
+define i8 @f_i8_255() {
+  ; CHECK-LABEL: name: f_i8_255
+  ; CHECK: bb.1 (%ir-block.0):
+  ; CHECK-NEXT:   [[DEF:%[0-9]+]]:_(<256 x s8>) = G_IMPLICIT_DEF
+  ; CHECK-NEXT:   [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 255
+  ; CHECK-NEXT:   [[EVEC:%[0-9]+]]:_(s8) = G_EXTRACT_VECTOR_ELT [[DEF]](<256 x s8>), [[C]](s32)
+  ; CHECK-NEXT:   [[ANYEXT:%[0-9]+]]:_(s32) = G_ANYEXT [[EVEC]](s8)
+  ; CHECK-NEXT:   $vgpr0 = COPY [[ANYEXT]](s32)
+  ; CHECK-NEXT:   SI_RETURN implicit $vgpr0
+  %E1 = extractelement <256 x i8> undef, i8 255
+  ret i8 %E1
+}