[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

This is generic

llvm/lib/Target/AMDGPU/AMDGPUPreLegalizerCombiner.cpp

162–164

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

166–170

Don't see why this would restrict the vector type

[Pierre-vh]

llvm/lib/Target/AMDGPU/AMDGPUCombine.td

48–52

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

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

"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

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

By as-is I mean DAGCombiner

[Pierre-vh]

Rebase on D136922, make combine generic

[arsenm]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2688–2694 ↗

(On Diff #471430)

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 ↗

(On Diff #471430)

I'm not really sure why this eraseInst helper exists

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

2 ↗

(On Diff #471430)

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

[Pierre-vh]

Comments + rebase

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2688–2694 ↗

(On Diff #471430)

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 ↗

(On Diff #471430)

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/lib/Target/AMDGPU/AMDGPUCombine.td

def cvt_f32_ubyteN : GICombineRule<
   (defs root:$cvt_f32_ubyteN, cvt_f32_ubyteN_matchdata:$matchinfo),
   (match (wip_match_opcode G_AMDGPU_CVT_F32_UBYTE0,
                            G_AMDGPU_CVT_F32_UBYTE1,
                            G_AMDGPU_CVT_F32_UBYTE2,
                            G_AMDGPU_CVT_F32_UBYTE3):$cvt_f32_ubyteN,
          [{ return PostLegalizerHelper.matchCvtF32UByteN(*${cvt_f32_ubyteN}, ${matchinfo}); }]),
   (apply [{ PostLegalizerHelper.applyCvtF32UByteN(*${cvt_f32_ubyteN}, ${matchinfo}); }])>;
 
+def insert_vec_elt_to_shuffle : GICombineRule<
+  (defs root:$insertelt, unsigned_matchinfo:$matchinfo),
+  (match (wip_match_opcode G_INSERT_VECTOR_ELT):$insertelt,
+      [{ return PreLegalizerHelper.matchInsertVectorEltToShuffle(*${insertelt}, ${matchinfo}); }]),
+  (apply [{ PreLegalizerHelper.applyInsertVectorEltToShuffle(*${insertelt}, ${matchinfo}); }])>;
+
 def clamp_i64_to_i16_matchdata : GIDefMatchData<"AMDGPUPreLegalizerCombinerHelper::ClampI64ToI16MatchInfo">;
 
 def clamp_i64_to_i16 : GICombineRule<
   (defs root:$clamp_i64_to_i16, clamp_i64_to_i16_matchdata:$matchinfo),
   (match (wip_match_opcode G_TRUNC):$clamp_i64_to_i16,
       [{ return PreLegalizerHelper.matchClampI64ToI16(*${clamp_i64_to_i16}, MRI, *MF, ${matchinfo}); }]),
   (apply [{ PreLegalizerHelper.applyClampI64ToI16(*${clamp_i64_to_i16}, ${matchinfo}); }])>;
 
(match (wip_match_opcode G_FNEG):$ffn,
          [{ return Helper.matchFoldableFneg(*${ffn}, ${matchinfo}); }]),
   (apply [{ Helper.applyFoldableFneg(*${ffn}, ${matchinfo}); }])>;
 
 // Combines which should only apply on SI/VI
 def gfx6gfx7_combines : GICombineGroup<[fcmp_select_to_fmin_fmax_legacy]>;
 
 def AMDGPUPreLegalizerCombinerHelper: GICombinerHelper<
   "AMDGPUGenPreLegalizerCombinerHelper",
-  [all_combines, clamp_i64_to_i16, foldable_fneg]> {
+  [all_combines, clamp_i64_to_i16, foldable_fneg, insert_vec_elt_to_shuffle]> {
   let DisableRuleOption = "amdgpuprelegalizercombiner-disable-rule";
   let StateClass = "AMDGPUPreLegalizerCombinerHelperState";
   let AdditionalArguments = [];
 }
 
 def AMDGPUPostLegalizerCombinerHelper: GICombinerHelper<
   "AMDGPUGenPostLegalizerCombinerHelper",
   [all_combines, gfx6gfx7_combines,

llvm/lib/Target/AMDGPU/AMDGPUPreLegalizerCombiner.cpp

public:
   };
 
   bool matchClampI64ToI16(MachineInstr &MI, MachineRegisterInfo &MRI,
                           MachineFunction &MF,
                           ClampI64ToI16MatchInfo &MatchInfo);
 
   void applyClampI64ToI16(MachineInstr &MI,
                           const ClampI64ToI16MatchInfo &MatchInfo);
+
+  bool matchInsertVectorEltToShuffle(MachineInstr &MI, unsigned &Idx);
+  void applyInsertVectorEltToShuffle(MachineInstr &MI, unsigned &Idx);
 };
 
 bool AMDGPUPreLegalizerCombinerHelper::matchClampI64ToI16(
     MachineInstr &MI, MachineRegisterInfo &MRI, MachineFunction &MF,
     ClampI64ToI16MatchInfo &MatchInfo) {
   assert(MI.getOpcode() == TargetOpcode::G_TRUNC && "Invalid instruction!");
 
   // Try to find a pattern where an i64 value should get clamped to short.
auto Med3 = B.buildInstr(
       {MinBoundaryDst.getReg(0), Bitcast.getReg(0), MaxBoundaryDst.getReg(0)},
       MI.getFlags());
 
   B.buildTrunc(MI.getOperand(0).getReg(), Med3);
 
   MI.eraseFromParent();
 }
 
+bool AMDGPUPreLegalizerCombinerHelper::matchInsertVectorEltToShuffle(
+    MachineInstr &MI, unsigned &Idx) {
+  // This is only beneficial if scalar pack insts are present.
+  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;
+
+  Optional<ValueAndVReg> MaybeIdxVal =
+      getIConstantVRegValWithLookThrough(MI.getOperand(3).getReg(), MRI);
+  if (!MaybeIdxVal)
+    return false;
+
+  Idx = MaybeIdxVal->Value.getZExtValue();
+  return true;
+}
+
+void AMDGPUPreLegalizerCombinerHelper::applyInsertVectorEltToShuffle(
+    MachineInstr &MI, unsigned &Idx) {
+  B.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();
+
+  const auto Undef = MRI.createGenericVirtualRegister(EltTy);
+  B.buildUndef(Undef);
+
+  const auto OtherVec = MRI.createGenericVirtualRegister(VecTy);
+
+  SmallVector<Register, 4> Srcs;
+  Srcs.push_back(Ins);
+  for (unsigned K = 1; K < NumElts; ++K)
+    Srcs.push_back(Undef);
+
+  B.buildBuildVector(OtherVec, Srcs);
+
+  // 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);
+  }
+
+  B.buildShuffleVector(Dst, Vec, OtherVec, ShuffleMask);
+  Helper.eraseInst(MI);
+}
+
 class AMDGPUPreLegalizerCombinerHelperState {
 protected:
   AMDGPUCombinerHelper &Helper;
   AMDGPUPreLegalizerCombinerHelper &PreLegalizerHelper;
 
 public:
   AMDGPUPreLegalizerCombinerHelperState(
       AMDGPUCombinerHelper &Helper,

llvm/test/CodeGen/AMDGPU/GlobalISel/prelegalizer-combiner-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 --check-prefixes=CHECK,GFX9PLUS
+# RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=fiji -run-pass=amdgpu-prelegalizer-combiner -verify-machineinstrs -o - %s | FileCheck %s --check-prefixes=CHECK,VI
+
+---
+name: test_v2s16_idx0
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $vgpr0
+    ; GFX9PLUS-LABEL: name: test_v2s16_idx0
+    ; GFX9PLUS: liveins: $vgpr0
+    ; GFX9PLUS-NEXT: {{  $}}
+    ; GFX9PLUS-NEXT: %src:_(<2 x s16>) = COPY $vgpr0
+    ; GFX9PLUS-NEXT: %elt:_(s16) = G_CONSTANT i16 42
+    ; GFX9PLUS-NEXT: [[DEF:%[0-9]+]]:_(s16) = G_IMPLICIT_DEF
+    ; GFX9PLUS-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x s16>) = G_BUILD_VECTOR %elt(s16), [[DEF]](s16)
+    ; GFX9PLUS-NEXT: %ins:_(<2 x s16>) = G_SHUFFLE_VECTOR %src(<2 x s16>), [[BUILD_VECTOR]], shufflemask(2, 1)
+    ; GFX9PLUS-NEXT: $vgpr0 = COPY %ins(<2 x s16>)
+    ; VI-LABEL: name: test_v2s16_idx0
+    ; VI: liveins: $vgpr0
+    ; VI-NEXT: {{  $}}
+    ; VI-NEXT: %src:_(<2 x s16>) = COPY $vgpr0
+    ; VI-NEXT: %idx:_(s32) = G_CONSTANT i32 0
+    ; VI-NEXT: %elt:_(s16) = G_CONSTANT i16 42
+    ; VI-NEXT: %ins:_(<2 x s16>) = G_INSERT_VECTOR_ELT %src, %elt(s16), %idx(s32)
+    ; VI-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
+    ; GFX9PLUS-LABEL: name: test_v2s16_idx1
+    ; GFX9PLUS: liveins: $vgpr0
+    ; GFX9PLUS-NEXT: {{  $}}
+    ; GFX9PLUS-NEXT: %src:_(<2 x s16>) = COPY $vgpr0
+    ; GFX9PLUS-NEXT: %elt:_(s16) = G_CONSTANT i16 42
+    ; GFX9PLUS-NEXT: [[DEF:%[0-9]+]]:_(s16) = G_IMPLICIT_DEF
+    ; GFX9PLUS-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x s16>) = G_BUILD_VECTOR %elt(s16), [[DEF]](s16)
+    ; GFX9PLUS-NEXT: %ins:_(<2 x s16>) = G_SHUFFLE_VECTOR %src(<2 x s16>), [[BUILD_VECTOR]], shufflemask(0, 2)
+    ; GFX9PLUS-NEXT: $vgpr0 = COPY %ins(<2 x s16>)
+    ; VI-LABEL: name: test_v2s16_idx1
+    ; VI: liveins: $vgpr0
+    ; VI-NEXT: {{  $}}
+    ; VI-NEXT: %src:_(<2 x s16>) = COPY $vgpr0
+    ; VI-NEXT: %idx:_(s32) = G_CONSTANT i32 1
+    ; VI-NEXT: %elt:_(s16) = G_CONSTANT i16 42
+    ; VI-NEXT: %ins:_(<2 x s16>) = G_INSERT_VECTOR_ELT %src, %elt(s16), %idx(s32)
+    ; VI-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_nofold
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $vgpr0_vgpr1_vgpr2
+    ; CHECK-LABEL: name: test_v3s16_idx2_nofold
+    ; 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: %idx:_(s32) = G_CONSTANT i32 2
+    ; CHECK-NEXT: %elt:_(s16) = G_CONSTANT i16 42
+    ; CHECK-NEXT: %ins:_(<3 x s16>) = G_INSERT_VECTOR_ELT %truncsrc, %elt(s16), %idx(s32)
+    ; 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_nofold
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $vgpr0_vgpr1
+    ; CHECK-LABEL: name: test_v2s32_idx1_nofold
+    ; CHECK: liveins: $vgpr0_vgpr1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %src:_(<2 x s32>) = COPY $vgpr0_vgpr1
+    ; CHECK-NEXT: %idx:_(s32) = G_CONSTANT i32 1
+    ; CHECK-NEXT: %elt:_(s16) = G_CONSTANT i16 42
+    ; CHECK-NEXT: %ins:_(<2 x s32>) = G_INSERT_VECTOR_ELT %src, %elt(s16), %idx(s32)
+    ; CHECK-NEXT: $vgpr0_vgpr1 = COPY %ins(<2 x s32>)
+    %src:_(<2 x s32>) = COPY $vgpr0_vgpr1
+    %idx:_(s32) = G_CONSTANT i32 1
+    %elt:_(s16) = G_CONSTANT i16 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
+    ; GFX9PLUS-LABEL: name: test_v4s16_idx3
+    ; GFX9PLUS: liveins: $vgpr0_vgpr1
+    ; GFX9PLUS-NEXT: {{  $}}
+    ; GFX9PLUS-NEXT: %src:_(<4 x s16>) = COPY $vgpr0_vgpr1
+    ; GFX9PLUS-NEXT: %elt:_(s16) = G_CONSTANT i16 42
+    ; GFX9PLUS-NEXT: [[DEF:%[0-9]+]]:_(s16) = G_IMPLICIT_DEF
+    ; GFX9PLUS-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<4 x s16>) = G_BUILD_VECTOR %elt(s16), [[DEF]](s16), [[DEF]](s16), [[DEF]](s16)
+    ; GFX9PLUS-NEXT: %ins:_(<4 x s16>) = G_SHUFFLE_VECTOR %src(<4 x s16>), [[BUILD_VECTOR]], shufflemask(0, 1, 2, 4)
+    ; GFX9PLUS-NEXT: $vgpr0_vgpr1 = COPY %ins(<4 x s16>)
+    ; VI-LABEL: name: test_v4s16_idx3
+    ; VI: liveins: $vgpr0_vgpr1
+    ; VI-NEXT: {{  $}}
+    ; VI-NEXT: %src:_(<4 x s16>) = COPY $vgpr0_vgpr1
+    ; VI-NEXT: %idx:_(s32) = G_CONSTANT i32 3
+    ; VI-NEXT: %elt:_(s16) = G_CONSTANT i16 42
+    ; VI-NEXT: %ins:_(<4 x s16>) = G_INSERT_VECTOR_ELT %src, %elt(s16), %idx(s32)
+    ; VI-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
+...