[GISel]: Allow G_EXTRACT_VEC_ELT's result to be larger the source element type

[aditya_nandakumar]

Summary

It's quite inconvenient when dealing with G_EXTRACT_VEC_ELT of <4xs8>,0 where s8 is not legal in the target - typically one ends up with G_UNMERGE_VALUES + G_LSHR by 8 which expands to three instructions.
Instead allow the destination type to be larger than source scalar type and it's higher bits are any extended. This is similar to DAG's opcode.
Note - I couldn't really find a use in tree - so I added a unit test.

[arsenm]

I always thought this behavior in SelectionDAG was error prone. I think I broke this on every combine I ever wrote for extract_vector_elt. Should there be a separate extending opcode instead? Although that would probably make some combines more painful. In GlobalISel we already have an explicit G_BUILD_VECTOR_TRUNC instead of the old ISD::BUILD_VECTOR truncate behavior, so that would match. However I don't think G_BUILD_VECTOR_TRUNC is well developed enough to declare that was definitely a good decision yet.

lib/CodeGen/MachineVerifier.cpp

1276

braces

1281

braces

[aditya_nandakumar]

In D58419#1404103, @arsenm wrote:

I always thought this behavior in SelectionDAG was error prone.

Personally in an out of tree backend, I find BUILD_VEC_TRUNC to be very elegant when trying to legalize/optimize vector types whose scalar types are illegal in the architecture otherwise needing several unmerges + shifts + ands.
I'm fine with adding a separate opcode (say G_EXTRACT_VECTOR_ELT_ANYEXT - suggestions accepted).
We need to do something similar for INSERT_VEC_ELT which can implicitly truncate.

[arsenm]

In D58419#1404416, @aditya_nandakumar wrote:

In D58419#1404103, @arsenm wrote:

I always thought this behavior in SelectionDAG was error prone.

Personally in an out of tree backend, I find BUILD_VEC_TRUNC to be very elegant when trying to legalize/optimize vector types whose scalar types are illegal in the architecture otherwise needing several unmerges + shifts + ands.
I'm fine with adding a separate opcode (say G_EXTRACT_VECTOR_ELT_ANYEXT - suggestions accepted).
We need to do something similar for INSERT_VEC_ELT which can implicitly truncate.

One question I've been thinking about is whether implicit extension is needed for G_INSERT/G_EXTRACT. I figured extract_vector_elt/insert_vector_elt with a constant could always be implemented as a G_EXTRACT/G_INSERT.

[aditya_nandakumar]

In D58419#1404440, @arsenm wrote:

In D58419#1404416, @aditya_nandakumar wrote:

In D58419#1404103, @arsenm wrote:

I always thought this behavior in SelectionDAG was error prone.

Personally in an out of tree backend, I find BUILD_VEC_TRUNC to be very elegant when trying to legalize/optimize vector types whose scalar types are illegal in the architecture otherwise needing several unmerges + shifts + ands.
I'm fine with adding a separate opcode (say G_EXTRACT_VECTOR_ELT_ANYEXT - suggestions accepted).
We need to do something similar for INSERT_VEC_ELT which can implicitly truncate.

One question I've been thinking about is whether implicit extension is needed for G_INSERT/G_EXTRACT. I figured extract_vector_elt/insert_vector_elt with a constant could always be implemented as a G_EXTRACT/G_INSERT.

G_EXTRACT and G_INSERT are way more powerful instructions IMO and going from EXTRACT_VEC_ELT to EXTRACT seems to adding complexity (where as I'd prefer the other way if possible as sort of a canonicalization). I'm curious to hear others thoughts on this as well.

[aditya_nandakumar]

Missing braces.

[aditya_nandakumar]

ping?

[arsenm]

I think this should be symmetric with G_BUILD_VECTOR_TRUNC, so either a separate opcode or re-merge BUILD_VECTOR_TRUNC with BUILD_VECTOR

[aemerson]

In D58419#1419010, @arsenm wrote:

I think this should be symmetric with G_BUILD_VECTOR_TRUNC, so either a separate opcode or re-merge BUILD_VECTOR_TRUNC with BUILD_VECTOR

Seems reasonable. Let's go ahead with this.

include/llvm/Target/GenericOpcodes.td

 
 // Generic insertelement.
 def G_INSERT_VECTOR_ELT : GenericInstruction {
   let OutOperandList = (outs type0:$dst);
   let InOperandList = (ins type0:$src, type1:$elt, type2:$idx);
   let hasSideEffects = 0;
 }
 
-// Generic extractelement.
+// Generic extractelement. If the return type is larger than the vector's scalar element,
+// then the higher bits are undefined and the result is any_ext'd.
 def G_EXTRACT_VECTOR_ELT : GenericInstruction {
   let OutOperandList = (outs type0:$dst);
   let InOperandList = (ins type1:$src, type2:$idx);
   let hasSideEffects = 0;
 }
 
 // Generic shufflevector.
 def G_SHUFFLE_VECTOR: GenericInstruction {
   let OutOperandList = (outs type0:$dst);
   let InOperandList = (ins type1:$v1, type1:$v2, type2:$mask);
   let hasSideEffects = 0;
 }
 
 // TODO: Add the other generic opcodes.

lib/CodeGen/GlobalISel/MachineIRBuilder.cpp

MachineInstrBuilder MachineIRBuilder::buildInstr(unsigned Opc,
   case TargetOpcode::G_EXTRACT_VECTOR_ELT: {
     assert(DstOps.size() == 1 && "Invalid Dst size");
     assert(SrcOps.size() == 2 && "Invalid Src size");
     assert(SrcOps[0].getLLTTy(*getMRI()).isVector() && "Invalid operand type");
     assert((DstOps[0].getLLTTy(*getMRI()).isScalar() ||
             DstOps[0].getLLTTy(*getMRI()).isPointer()) &&
            "Invalid operand type");
     assert(SrcOps[1].getLLTTy(*getMRI()).isScalar() && "Invalid operand type");
-    assert(SrcOps[0].getLLTTy(*getMRI()).getElementType() ==
-               DstOps[0].getLLTTy(*getMRI()) &&
+    assert(SrcOps[0].getLLTTy(*getMRI()).getElementType().getSizeInBits() <=
+               DstOps[0].getLLTTy(*getMRI()).getSizeInBits() &&
            "Type mismatch");
     break;
   }
   case TargetOpcode::G_INSERT_VECTOR_ELT: {
     assert(DstOps.size() == 1 && "Invalid dst size");
     assert(SrcOps.size() == 3 && "Invalid src size");
     assert(DstOps[0].getLLTTy(*getMRI()).isVector() &&
            SrcOps[0].getLLTTy(*getMRI()).isVector() && "Invalid operand type");

lib/CodeGen/MachineVerifier.cpp

case TargetOpcode::G_FCMP: {
     LLT SrcTy = MRI->getType(MI->getOperand(2).getReg());
 
     if ((DstTy.isVector() != SrcTy.isVector()) ||
         (DstTy.isVector() && DstTy.getNumElements() != SrcTy.getNumElements()))
       report("Generic vector icmp/fcmp must preserve number of lanes", MI);
 
     break;
   }
+  case TargetOpcode::G_EXTRACT_VECTOR_ELT: {
+    LLT DstTy = MRI->getType(MI->getOperand(0).getReg());
+    LLT SrcTy = MRI->getType(MI->getOperand(1).getReg());
+    if (!SrcTy.isVector())
+      report("Generic G_EXTRACT_VECTOR_ELT's source should be vector", MI);
+    if (!DstTy.isScalar() && !DstTy.isPointer())
+      report("Generic G_EXTRACT_VECTOR_ELT's dst should be scalar/pointer", MI);
+    if (DstTy.isPointer() && DstTy != SrcTy.getElementType()) {
+      report("Generic G_EXTRACT_VECTOR_ELT's dst and src should match when "
+             "pointer",
+             MI);
+    }
+    if (DstTy.isScalar() && SrcTy.isVector() &&
+        DstTy.getSizeInBits() < SrcTy.getElementType().getSizeInBits()) {
+      report("Generic G_EXTRACT_VECTOR_ELT's dst can't be smaller than src "
+             "element type",
+             MI);
+    }
+    break;
+  }
   case TargetOpcode::G_EXTRACT: {
     const MachineOperand &SrcOp = MI->getOperand(1);
     if (!SrcOp.isReg()) {
       report("extract source must be a register", MI);
       break;
     }
 
     const MachineOperand &OffsetOp = MI->getOperand(2);

test/MachineVerifier/test_g_extract_vec_elt.mir

+# RUN: not llc -march=aarch64 -o /dev/null -run-pass=none -verify-machineinstrs %s 2>&1 | FileCheck %s
+# REQUIRES: global-isel, aarch64-registered-target
+
+---
+name:            test_extract
+legalized:       true
+regBankSelected: false
+selected:        false
+tracksRegLiveness: true
+liveins:
+body:             |
+  bb.0:
+
+    %0:_(<4 x s16>) = G_IMPLICIT_DEF
+    %1:_(s32) = G_IMPLICIT_DEF
+    %2:_(s32) = G_CONSTANT i32 0
+
+    ; CHECK: Generic G_EXTRACT_VECTOR_ELT's source should be vector
+    %3:_(s32) = G_EXTRACT_VECTOR_ELT %1, %2
+
+    ; CHECK: Generic G_EXTRACT_VECTOR_ELT's dst can't be smaller than src element type
+    %4:_(s8) = G_EXTRACT_VECTOR_ELT %0, %2
+
+    ; This should be okay.
+    %5:_(s32) = G_EXTRACT_VECTOR_ELT %0, %2
+
+    ; CHECK: Generic G_EXTRACT_VECTOR_ELT's dst should be scalar/pointer
+    %6:_(<4 x s8>) = G_EXTRACT_VECTOR_ELT %0, %2
+
+    %7:_(<4 x p0>) = G_IMPLICIT_DEF
+    ; CHECK: Generic G_EXTRACT_VECTOR_ELT's dst and src should match when pointer
+    %8:_(p1) = G_EXTRACT_VECTOR_ELT %7, %2
+...

unittests/CodeGen/GlobalISel/MachineIRBuilderTest.cpp

TEST_F(GISelMITest, DstOpSrcOp) {
   ; CHECK: [[COPY0:%[0-9]+]]:_(s64) = COPY $x0
   ; CHECK: [[COPY1:%[0-9]+]]:_(s64) = COPY $x1
   ; CHECK: [[ADD:%[0-9]+]]:_(s64) = G_ADD [[COPY0]]:_, [[COPY1]]:_
   ; CHECK: [[ADD]]:_(s64) = G_ADD [[COPY0]]:_, [[COPY1]]:_
   )";
 
   EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
 }
+
+TEST_F(GISelMITest, ExtractVecEltAnyExt) {
+  if (!TM)
+    return;
+
+  SmallVector<unsigned, 4> Copies;
+  collectCopies(Copies, MF);
+
+  LLT s32 = LLT::scalar(32);
+  LLT s16 = LLT::scalar(16);
+  auto MIB = B.buildExtractVectorElement(s32, B.buildUndef(LLT::vector(4, s16)),
+                                         B.buildConstant(s32, 0));
+
+  EXPECT_TRUE(MIB->getOpcode() == TargetOpcode::G_EXTRACT_VECTOR_ELT);
+}