[GISel] Combine G_INSERT_VECTOR_ELT to G_SHUFFLE_VECTOR

[Pierre-vh]

Summary

Depends on D136922

[Pierre-vh]

Rebase

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUCombine.td

48–52 ↗

(On Diff #467049)

This is generic

llvm/lib/Target/AMDGPU/AMDGPUPreLegalizerCombiner.cpp

162–164 ↗

(On Diff #467049)

The DAG treats this as an initial canonicalization, so the obvious codegen benefit isn't so important

166–170 ↗

(On Diff #467049)

Don't see why this would restrict the vector type

[Pierre-vh]

llvm/lib/Target/AMDGPU/AMDGPUCombine.td

48–52 ↗

(On Diff #467049)

Do you mean it should go in the generic combiner?
I'm worried that if we put it there, and remove the filtering on small vectors/hasScalarPackInsts that all INSERT_VECTOR_ELT instructions will become SHUFFLE_VECTOR and that some targets won't like it?

If we remove

if (!MI.getMF()->getSubtarget<GCNSubtarget>().hasScalarPackInsts())
  return false;

// TODO: Only on small vectors?
LLT VecTy = MRI.getType(MI.getOperand(0).getReg());
if (VecTy.getElementType() != LLT::scalar(16) ||
    (VecTy.getSizeInBits() % 32) != 0)
  return false;

I would leave it in the AMDGPUCombiner, if we want to make it generic, I would at least add some safeguard so it doesn't turn every INSERT_VECTOR_ELT into a shuffle - maybe only do it for 2-elt vectors?
That or just don't add it to "all_combines" - we put it in the generic helper but it's opt-in and targets have too add the combine to their pipeline.

llvm/lib/Target/AMDGPU/AMDGPUPreLegalizerCombiner.cpp

166–170 ↗

(On Diff #467049)

Shouldn't this just be on 2-element vectors?

[foad]

What is the motivation for this?

Can you add a comment somewhere explaining what the combine does?

[Pierre-vh]

Relax some restrictions on the combine and add comment to describe why the current restrictions are in place

[Pierre-vh]

In D135145#3861774, @foad wrote:

What is the motivation for this?

Can you add a comment somewhere explaining what the combine does?

We want to use SHUFFLE_VECTOR (which is always lowered during legalization anyway) as the canonical form for this kind of operation (INSERT_VECTOR_ELT w/ a constant index on small vectors). It benefits mad_mix codegen.

[Pierre-vh]

Rebase

[foad]

llvm/lib/Target/AMDGPU/AMDGPUPreLegalizerCombiner.cpp

162 ↗

(On Diff #468794)

"G_SHUFFLE_VECTOR"

[Pierre-vh]

Comment
@arsenm please review so D134354 can land?

[Pierre-vh]

ping

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUCombine.td

48–52 ↗

(On Diff #467049)

Yes. This is a generic combine as it is. What the target directly wants isn't necessarily the point. A larger shuffle should be legalizable to what the target does want, and is a better canonical form

48–52 ↗

(On Diff #467049)

By as-is I mean DAGCombiner

[Pierre-vh]

Rebase on D136922, make combine generic

[arsenm]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2688–2694

There are cases where insert_vector_elts combine to form shuffles but you don't seem to be handling those. This looks like you're just handling basic cases that can use build_vector (which is already implemented in matchCombineInsertVecElts). I'm not following what the shuffles are adding here

2697

I'm not really sure why this eraseInst helper exists

llvm/test/CodeGen/AMDGPU/GlobalISel/combine-insertvecelt-to-shufflevector.mir

2

Don't need -global-isel with -run-pass

[Pierre-vh]

Comments + rebase

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2688–2694

I thought we ultimately wanted insert_vector_elt to be lowered to shuffle_vector, as the latter is easier to handle (and is needed for mad_mix selection)? Is that not the case?
It the reason why we're lowering shuffle vectors now, no?

In any case, I'm not sure I understand the issue:

• Is it that the combine is unnecessary? (Then why are working towards this? Why did we lower shuffle vector in the Legalizer?)
• Is it that the combine as-is is fine, but should be handling more (like chained insert_vector_elt) ? But then, what makes it different from matchCombineInsertVecElts?

Note that, IIRC, matchCombineInsertVecElts only handles chains of insert_vector_elt. If there's a single one, it doesn't touch it. This combine is targeted towards single insert_vector_elts.

2697

Not sure either, it's a Combiner helper. I thought it was doing some other things like notifying the observer but it really just calls MI.eraseFromParent().
I've removed this use and will propose a patch to remove it entirely.

[Pierre-vh]

Not sure yet this is the right thing to do, I can resurrect the diff later if we still want to do it.

llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

public:
   bool applyFoldBinOpIntoSelect(MachineInstr &MI, const unsigned &SelectOpNo);
 
   bool matchCombineInsertVecElts(MachineInstr &MI,
                                  SmallVectorImpl<Register> &MatchInfo);
 
   void applyCombineInsertVecElts(MachineInstr &MI,
                              SmallVectorImpl<Register> &MatchInfo);
 
+  bool matchInsertVectorEltToShuffle(MachineInstr &MI, unsigned &Idx);
+  void applyInsertVectorEltToShuffle(MachineInstr &MI, unsigned &Idx);
+
   /// Match expression trees of the form
   ///
   /// \code
   ///  sN *a = ...
   ///  sM val = a[0] | (a[1] << N) | (a[2] << 2N) | (a[3] << 3N) ...
   /// \endcode
   ///
   /// And check if the tree can be replaced with a M-bit load + possibly a

llvm/include/llvm/Target/GlobalISel/Combine.td

 
 def extend_through_phis_matchdata: GIDefMatchData<"MachineInstr*">;
 def extend_through_phis : GICombineRule<
   (defs root:$root, extend_through_phis_matchdata:$matchinfo),
   (match (wip_match_opcode G_PHI):$root,
     [{ return Helper.matchExtendThroughPhis(*${root}, ${matchinfo}); }]),
   (apply [{ Helper.applyExtendThroughPhis(*${root}, ${matchinfo}); }])>;
 
-// Currently only the one combine above.
+// Canonicalizes (insert_vector_elt X, K) into a shuffle_vector.
+def insert_vec_elt_to_shuffle : GICombineRule<
+  (defs root:$insertelt, unsigned_matchinfo:$matchinfo),
+  (match (wip_match_opcode G_INSERT_VECTOR_ELT):$insertelt,
+      [{ return Helper.matchInsertVectorEltToShuffle(*${insertelt}, ${matchinfo}); }]),
+  (apply [{ Helper.applyInsertVectorEltToShuffle(*${insertelt}, ${matchinfo}); }])>;
+
 def insert_vec_elt_combines : GICombineGroup<
-                            [combine_insert_vec_elts_build_vector]>;
+                            [combine_insert_vec_elts_build_vector,
+                            insert_vec_elt_to_shuffle]>;
 
 def extract_vec_elt_build_vec : GICombineRule<
   (defs root:$root, register_matchinfo:$matchinfo),
   (match (wip_match_opcode G_EXTRACT_VECTOR_ELT):$root,
     [{ return Helper.matchExtractVecEltBuildVec(*${root}, ${matchinfo}); }]),
   (apply [{ Helper.applyExtractVecEltBuildVec(*${root}, ${matchinfo}); }])>;
 
 // Fold away full elt extracts from a build_vector.

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

void CombinerHelper::applyCombineInsertVecElts(
   for (unsigned I = 0; I < MatchInfo.size(); ++I) {
     if (!MatchInfo[I])
       MatchInfo[I] = GetUndef();
   }
   Builder.buildBuildVector(MI.getOperand(0).getReg(), MatchInfo);
   MI.eraseFromParent();
 }
 
+bool CombinerHelper::matchInsertVectorEltToShuffle(MachineInstr &MI,
+                                                   unsigned &Idx) {
+  assert(MI.getOpcode() == TargetOpcode::G_INSERT_VECTOR_ELT);
+
+  // Canonicalizes a G_INSERT_VECTOR_ELT w/ a constant index into an equivalent
+  // G_SHUFFLE_VECTOR if it is a legal transformation.
+
+  // If this MI is part of a sequence of insert_vec_elts, then
+  // don't do the combine in the middle of the sequence.
+  Register DstReg = MI.getOperand(0).getReg();
+  if (MRI.hasOneUse(DstReg) && MRI.use_instr_begin(DstReg)->getOpcode() ==
+                                   TargetOpcode::G_INSERT_VECTOR_ELT)
+    return false;
+
+  LLT VecTy = MRI.getType(DstReg);
+  LLT EltTy = MRI.getType(MI.getOperand(2).getReg());
+  LLT IdxTy = MRI.getType(MI.getOperand(3).getReg());
+
+  if (VecTy.isScalable() ||
+      !isLegalOrBeforeLegalizer(
+          {TargetOpcode::G_INSERT_VECTOR_ELT, {VecTy, EltTy, IdxTy}}))
+    return false;
+
+  const auto MaybeIdxVal =
+      getIConstantVRegValWithLookThrough(MI.getOperand(3).getReg(), MRI);
+  if (!MaybeIdxVal)
+    return false;
+
+  Idx = MaybeIdxVal->Value.getZExtValue();
+  return Idx < VecTy.getNumElements();
+}
+
+void CombinerHelper::applyInsertVectorEltToShuffle(MachineInstr &MI,
+                                                   unsigned &Idx) {
+  Builder.setInstrAndDebugLoc(MI);
+
+  Register Ins = MI.getOperand(2).getReg();
+  Register Vec = MI.getOperand(1).getReg();
+  Register Dst = MI.getOperand(0).getReg();
+
+  LLT VecTy = MRI.getType(Dst);
+  LLT EltTy = VecTy.getElementType();
+  const unsigned NumElts = VecTy.getNumElements();
+
+  Register Undef = Builder.buildUndef(EltTy).getReg(0);
+
+  SmallVector<Register, 4> Srcs;
+  Srcs.push_back(Ins);
+  for (unsigned K = 1; K < NumElts; ++K)
+    Srcs.push_back(Undef);
+
+  Register OtherVec = Builder.buildBuildVector(VecTy, Srcs).getReg(0);
+
+  // NumElts == Ins in OtherVec
+  // 0...(NumElts-1) = Original elements
+  SmallVector<int, 4> ShuffleMask;
+  for (unsigned CurIdx = 0; CurIdx < NumElts; ++CurIdx) {
+    if (CurIdx == Idx)
+      ShuffleMask.push_back(NumElts);
+    else
+      ShuffleMask.push_back(CurIdx);
+  }
+
+  Builder.buildShuffleVector(Dst, Vec, OtherVec, ShuffleMask);
+  eraseInst(MI);
+}
+
 void CombinerHelper::applySimplifyAddToSub(
     MachineInstr &MI, std::tuple<Register, Register> &MatchInfo) {
   Builder.setInstr(MI);
   Register SubLHS, SubRHS;
   std::tie(SubLHS, SubRHS) = MatchInfo;
   Builder.buildSub(MI.getOperand(0).getReg(), SubLHS, SubRHS);
   MI.eraseFromParent();
 }

llvm/test/CodeGen/AMDGPU/GlobalISel/combine-insertvecelt-to-shufflevector.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=amdgpu-prelegalizer-combiner -verify-machineinstrs -o - %s | FileCheck %s
+# RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=fiji -run-pass=amdgpu-prelegalizer-combiner -verify-machineinstrs -o - %s | FileCheck %s
+
+---
+name: test_v2s16_idx0
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $vgpr0
+    ; CHECK-LABEL: name: test_v2s16_idx0
+    ; CHECK: liveins: $vgpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %src:_(<2 x s16>) = COPY $vgpr0
+    ; CHECK-NEXT: %elt:_(s16) = G_CONSTANT i16 42
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(s16) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x s16>) = G_BUILD_VECTOR %elt(s16), [[DEF]](s16)
+    ; CHECK-NEXT: %ins:_(<2 x s16>) = G_SHUFFLE_VECTOR %src(<2 x s16>), [[BUILD_VECTOR]], shufflemask(2, 1)
+    ; CHECK-NEXT: $vgpr0 = COPY %ins(<2 x s16>)
+    %src:_(<2 x s16>) = COPY $vgpr0
+    %idx:_(s32) = G_CONSTANT i32 0
+    %elt:_(s16) = G_CONSTANT i16 42
+    %ins:_(<2 x s16>) = G_INSERT_VECTOR_ELT %src, %elt, %idx
+    $vgpr0 = COPY %ins
+...
+
+---
+name: test_v2s16_idx1
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $vgpr0
+    ; CHECK-LABEL: name: test_v2s16_idx1
+    ; CHECK: liveins: $vgpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %src:_(<2 x s16>) = COPY $vgpr0
+    ; CHECK-NEXT: %elt:_(s16) = G_CONSTANT i16 42
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(s16) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x s16>) = G_BUILD_VECTOR %elt(s16), [[DEF]](s16)
+    ; CHECK-NEXT: %ins:_(<2 x s16>) = G_SHUFFLE_VECTOR %src(<2 x s16>), [[BUILD_VECTOR]], shufflemask(0, 2)
+    ; CHECK-NEXT: $vgpr0 = COPY %ins(<2 x s16>)
+    %src:_(<2 x s16>) = COPY $vgpr0
+    %idx:_(s32) = G_CONSTANT i32 1
+    %elt:_(s16) = G_CONSTANT i16 42
+    %ins:_(<2 x s16>) = G_INSERT_VECTOR_ELT %src, %elt, %idx
+    $vgpr0 = COPY %ins
+...
+
+---
+name: test_v2s16_idx2_nofold
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $vgpr0
+    ; CHECK-LABEL: name: test_v2s16_idx2_nofold
+    ; CHECK: liveins: $vgpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %ins:_(<2 x s16>) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: $vgpr0 = COPY %ins(<2 x s16>)
+    %src:_(<2 x s16>) = COPY $vgpr0
+    %idx:_(s32) = G_CONSTANT i32 2
+    %elt:_(s16) = G_CONSTANT i16 42
+    %ins:_(<2 x s16>) = G_INSERT_VECTOR_ELT %src, %elt, %idx
+    $vgpr0 = COPY %ins
+...
+
+---
+name: test_v3s16_idx2
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $vgpr0_vgpr1_vgpr2
+    ; CHECK-LABEL: name: test_v3s16_idx2
+    ; CHECK: liveins: $vgpr0_vgpr1_vgpr2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %src:_(<3 x s32>) = COPY $vgpr0_vgpr1_vgpr2
+    ; CHECK-NEXT: %truncsrc:_(<3 x s16>) = G_TRUNC %src(<3 x s32>)
+    ; CHECK-NEXT: %elt:_(s16) = G_CONSTANT i16 42
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(s16) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<3 x s16>) = G_BUILD_VECTOR %elt(s16), [[DEF]](s16), [[DEF]](s16)
+    ; CHECK-NEXT: %ins:_(<3 x s16>) = G_SHUFFLE_VECTOR %truncsrc(<3 x s16>), [[BUILD_VECTOR]], shufflemask(0, 1, 3)
+    ; CHECK-NEXT: %zextins:_(<3 x s32>) = G_ZEXT %ins(<3 x s16>)
+    ; CHECK-NEXT: $vgpr0_vgpr1_vgpr2 = COPY %zextins(<3 x s32>)
+    %src:_(<3 x s32>) = COPY $vgpr0_vgpr1_vgpr2
+    %truncsrc:_(<3 x s16>) = G_TRUNC %src
+    %idx:_(s32) = G_CONSTANT i32 2
+    %elt:_(s16) = G_CONSTANT i16 42
+    %ins:_(<3 x s16>) = G_INSERT_VECTOR_ELT %truncsrc, %elt, %idx
+    %zextins:_(<3 x s32>) = G_ZEXT %ins
+    $vgpr0_vgpr1_vgpr2 = COPY %zextins
+...
+
+---
+name: test_v2s32_idx1
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $vgpr0_vgpr1
+    ; CHECK-LABEL: name: test_v2s32_idx1
+    ; CHECK: liveins: $vgpr0_vgpr1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %src:_(<2 x s32>) = COPY $vgpr0_vgpr1
+    ; CHECK-NEXT: %elt:_(s32) = G_CONSTANT i32 42
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(s32) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x s32>) = G_BUILD_VECTOR %elt(s32), [[DEF]](s32)
+    ; CHECK-NEXT: %ins:_(<2 x s32>) = G_SHUFFLE_VECTOR %src(<2 x s32>), [[BUILD_VECTOR]], shufflemask(0, 2)
+    ; CHECK-NEXT: $vgpr0_vgpr1 = COPY %ins(<2 x s32>)
+    %src:_(<2 x s32>) = COPY $vgpr0_vgpr1
+    %idx:_(s32) = G_CONSTANT i32 1
+    %elt:_(s32) = G_CONSTANT i32 42
+    %ins:_(<2 x s32>) = G_INSERT_VECTOR_ELT %src, %elt, %idx
+    $vgpr0_vgpr1 = COPY %ins
+...
+
+---
+name: test_v4s16_idx3
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $vgpr0_vgpr1
+    ; CHECK-LABEL: name: test_v4s16_idx3
+    ; CHECK: liveins: $vgpr0_vgpr1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %src:_(<4 x s16>) = COPY $vgpr0_vgpr1
+    ; CHECK-NEXT: %elt:_(s16) = G_CONSTANT i16 42
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(s16) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<4 x s16>) = G_BUILD_VECTOR %elt(s16), [[DEF]](s16), [[DEF]](s16), [[DEF]](s16)
+    ; CHECK-NEXT: %ins:_(<4 x s16>) = G_SHUFFLE_VECTOR %src(<4 x s16>), [[BUILD_VECTOR]], shufflemask(0, 1, 2, 4)
+    ; CHECK-NEXT: $vgpr0_vgpr1 = COPY %ins(<4 x s16>)
+    %src:_(<4 x s16>) = COPY $vgpr0_vgpr1
+    %idx:_(s32) = G_CONSTANT i32 3
+    %elt:_(s16) = G_CONSTANT i16 42
+    %ins:_(<4 x s16>) = G_INSERT_VECTOR_ELT %src, %elt, %idx
+    $vgpr0_vgpr1 = COPY %ins
+...