[X86] Allow X86ISD::Wrapper to be folded into the base of gather/scatter address

[craig.topper]

Summary

If the base of our gather corresponds to something contained in X86ISD::Wrapper we should be able to fold it into the address.

This patch refactors some of the address matching to more fully use the X86ISelAddressMode struct and the getAddressOperands helper. A new helper function matchVectorAddress is added to call matchWrapper or fall back to matchAddressBase.

We should also be able to support constant offsets from a wrapper, but I'll look into that in a future patch. We may even be able to completely reuse matchAddress here, but I wanted to start simple and work up to it.

[delena]

lib/Target/X86/X86ISelDAGToDAG.cpp

1513

Do you have an example when the Wrapper is not matched?

1548

{} format

1549

Could you, please, add a comment that explains what does "matchVectorAddress" mean?

[craig.topper]

lib/Target/X86/X86ISelDAGToDAG.cpp

1513

I think right now the only way that can happen is if the codemodel is not small or kernel. Do you want me to add a test?

1548

What's wrong with the formatting? Do you want the curly braces on the if removed? Or did you want me to merge the else and the "if (matchVectorAddress"

[craig.topper]

Add a large code model test. Merge 'else' and 'if (matchVectorAddress". Add a comment.

lib/Target/X86/X86ISelDAGToDAG.cpp

#include "X86GenDAGISel.inc"
 
   private:
     void Select(SDNode *N) override;
 
     bool foldOffsetIntoAddress(uint64_t Offset, X86ISelAddressMode &AM);
     bool matchLoadInAddress(LoadSDNode *N, X86ISelAddressMode &AM);
     bool matchWrapper(SDValue N, X86ISelAddressMode &AM);
     bool matchAddress(SDValue N, X86ISelAddressMode &AM);
+    bool matchVectorAddress(SDValue N, X86ISelAddressMode &AM);
     bool matchAdd(SDValue N, X86ISelAddressMode &AM, unsigned Depth);
     bool matchAddressRecursively(SDValue N, X86ISelAddressMode &AM,
                                  unsigned Depth);
     bool matchAddressBase(SDValue N, X86ISelAddressMode &AM);
     bool selectAddr(SDNode *Parent, SDValue N, SDValue &Base,
                     SDValue &Scale, SDValue &Index, SDValue &Disp,
                     SDValue &Segment);
     bool selectVectorAddr(SDNode *Parent, SDValue N, SDValue &Base,
bool X86DAGToDAGISel::matchAddressBase(SDValue N, X86ISelAddressMode &AM) {
   }
 
   // Default, generate it as a register.
   AM.BaseType = X86ISelAddressMode::RegBase;
   AM.Base_Reg = N;
   return false;
 }
 
+/// Helper for selectVectorAddr. Handles things that can be folded into a
+/// gather scatter address. The index register and scale should have already
+/// been handled.
+bool X86DAGToDAGISel::matchVectorAddress(SDValue N, X86ISelAddressMode &AM) {
+  // TODO: Support other operations.
+  switch (N.getOpcode()) {
+  case X86ISD::Wrapper:
+    if (!matchWrapper(N, AM))
+      return false;
+    break;
+  }
+
+  return matchAddressBase(N, AM);
+}
+
 bool X86DAGToDAGISel::selectVectorAddr(SDNode *Parent, SDValue N, SDValue &Base,
                                        SDValue &Scale, SDValue &Index,
                                        SDValue &Disp, SDValue &Segment) {
-  unsigned ScalarSize;
+  X86ISelAddressMode AM;
   if (auto Mgs = dyn_cast<MaskedGatherScatterSDNode>(Parent)) {
-    Base = Mgs->getBasePtr();
-    Index = Mgs->getIndex();
-    ScalarSize = Mgs->getValue().getScalarValueSizeInBits();
+    AM.IndexReg = Mgs->getIndex();
+    AM.Scale = Mgs->getValue().getScalarValueSizeInBits() / 8;
   } else {
     auto X86Gather = cast<X86MaskedGatherSDNode>(Parent);
-    Base = X86Gather->getBasePtr();
-    Index = X86Gather->getIndex();
-    ScalarSize = X86Gather->getValue().getScalarValueSizeInBits();
+    AM.IndexReg = X86Gather->getIndex();
+    AM.Scale = X86Gather->getValue().getScalarValueSizeInBits() / 8;
   }
 
-  X86ISelAddressMode AM;
   unsigned AddrSpace = cast<MemSDNode>(Parent)->getPointerInfo().getAddrSpace();
   // AddrSpace 256 -> GS, 257 -> FS, 258 -> SS.
   if (AddrSpace == 256)
     AM.Segment = CurDAG->getRegister(X86::GS, MVT::i16);
   if (AddrSpace == 257)
     AM.Segment = CurDAG->getRegister(X86::FS, MVT::i16);
   if (AddrSpace == 258)
     AM.Segment = CurDAG->getRegister(X86::SS, MVT::i16);
 
-  SDLoc DL(N);
-  Scale = getI8Imm(ScalarSize/8, DL);
-
   // If Base is 0, the whole address is in index and the Scale is 1
-  if (isa<ConstantSDNode>(Base)) {
-    assert(cast<ConstantSDNode>(Base)->isNullValue() &&
+  if (isa<ConstantSDNode>(N)) {
+    assert(cast<ConstantSDNode>(N)->isNullValue() &&
            "Unexpected base in gather/scatter");
-    Scale = getI8Imm(1, DL);
-    Base = CurDAG->getRegister(0, MVT::i32);
+    AM.Scale = 1;
+  } else {
+    if (matchVectorAddress(N, AM))
+      return false;
   }
-  if (AM.Segment.getNode())
-    Segment = AM.Segment;
-  else
-    Segment = CurDAG->getRegister(0, MVT::i32);
-  Disp = CurDAG->getTargetConstant(0, DL, MVT::i32);
+
+  MVT VT = N.getSimpleValueType();
+  if (AM.BaseType == X86ISelAddressMode::RegBase) {
+    if (!AM.Base_Reg.getNode())
+      AM.Base_Reg = CurDAG->getRegister(0, VT);
+  }
+
+  getAddressOperands(AM, SDLoc(N), Base, Scale, Index, Disp, Segment);
   return true;
 }
 
 /// Returns true if it is able to pattern match an addressing mode.
 /// It returns the operands which make up the maximal addressing mode it can
 /// match by reference.
 ///
 /// Parent is the parent node of the addr operand that is being matched.  It

test/CodeGen/X86/masked_gather_scatter.ll

; SKX_32-NEXT:    retl
   %b = add <4 x i64> %a, %g3
   ret <4 x i64> %b
 }
 declare <4 x i64> @llvm.masked.gather.v4i64.v4p0i64(<4 x i64*>, i32, <4 x i1>, <4 x i64>)
 
 define <8 x i32> @test_global_array(<8 x i64> %indxs) {
 ; KNL_64-LABEL: test_global_array:
 ; KNL_64:       # BB#0:
-; KNL_64-NEXT:    movl $glob_array, %eax
 ; KNL_64-NEXT:    kxnorw %k0, %k0, %k1
-; KNL_64-NEXT:    vpgatherqd (%rax,%zmm0,4), %ymm1 {%k1}
+; KNL_64-NEXT:    vpgatherqd glob_array(,%zmm0,4), %ymm1 {%k1}
 ; KNL_64-NEXT:    vmovdqa %ymm1, %ymm0
 ; KNL_64-NEXT:    retq
 ;
 ; KNL_32-LABEL: test_global_array:
 ; KNL_32:       # BB#0:
-; KNL_32-NEXT:    movl $glob_array, %eax
 ; KNL_32-NEXT:    kxnorw %k0, %k0, %k1
-; KNL_32-NEXT:    vpgatherqd (%eax,%zmm0,4), %ymm1 {%k1}
+; KNL_32-NEXT:    vpgatherqd glob_array(,%zmm0,4), %ymm1 {%k1}
 ; KNL_32-NEXT:    vmovdqa %ymm1, %ymm0
 ; KNL_32-NEXT:    retl
 ;
 ; SKX-LABEL: test_global_array:
 ; SKX:       # BB#0:
-; SKX-NEXT:    movl $glob_array, %eax
 ; SKX-NEXT:    kxnorw %k0, %k0, %k1
-; SKX-NEXT:    vpgatherqd (%rax,%zmm0,4), %ymm1 {%k1}
+; SKX-NEXT:    vpgatherqd glob_array(,%zmm0,4), %ymm1 {%k1}
 ; SKX-NEXT:    vmovdqa %ymm1, %ymm0
 ; SKX-NEXT:    retq
 ;
 ; SKX_32-LABEL: test_global_array:
 ; SKX_32:       # BB#0:
-; SKX_32-NEXT:    movl $glob_array, %eax
 ; SKX_32-NEXT:    kxnorw %k0, %k0, %k1
-; SKX_32-NEXT:    vpgatherqd (%eax,%zmm0,4), %ymm1 {%k1}
+; SKX_32-NEXT:    vpgatherqd glob_array(,%zmm0,4), %ymm1 {%k1}
 ; SKX_32-NEXT:    vmovdqa %ymm1, %ymm0
 ; SKX_32-NEXT:    retl
   %p = getelementptr inbounds [16 x i32], [16 x i32]* @glob_array, i64 0, <8 x i64> %indxs
   %g = call <8 x i32> @llvm.masked.gather.v8i32.v8p0i32(<8 x i32*> %p, i32 8, <8 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true>, <8 x i32> undef)
   ret <8 x i32> %g
 }