[X86] Machine combine vnni instruction.

[LuoYuanke]

Summary

"vpmaddwd + vpaddd" can be combined to vpdpwssd and the latency is
reduced after combination. However when vpdpwssd is in a critical path
the combination get less ILP. It happens when vpdpwssd is in a loop, the
vpmaddwd can be executed in parallel in multi-iterations while vpdpwssd
has data dependency for each iterations. If vpaddd is in a critical path
while vpmaddwd is not, it is profitable to split vpdpwssd into "vpmaddwd
+ vpaddd ".
This patch is based on the machine combiner framework to acheive decision
on "vpmaddwd + vpaddd" combination. The typical example code is as
below.

__m256i foo(int cnt, __m256i c, __m256i b, __m256i *p) {

    for (int i = 0; i < cnt; ++i) {
        __m256i a = p[i];
        __m256i m = _mm256_madd_epi16 (b, a);
        c = _mm256_add_epi32(m, c);
    }

    return c;
}

[LuoYuanke]

I'll add more patterns if it is the right way to solve https://github.com/llvm/llvm-project/issues/62026.

[goldstein.w.n]

llvm/lib/CodeGen/MachineCombiner.cpp

173

What is this change for?

424

This seems like an unrelated change. Can you split it?

llvm/lib/Target/X86/X86InstrInfo.cpp

9822

Maybe this belongs in X86FixupInstTuning.cpp?

llvm/test/CodeGen/X86/avxvnni-combine.ll

4

Can you precommit the test so we can see the diff?

[LuoYuanke]

Fix the AArch64 test case failure.

[LuoYuanke]

llvm/lib/CodeGen/MachineCombiner.cpp

173

For some reason the new created register for new created instructions is not managed by MRI, so it would miss adding the DefInstr's latency. I guess may be due to the new instructions are not inserted in the machine function. I'll debug more on it.

llvm/lib/Target/X86/X86InstrInfo.cpp

9822

The method is to override the virtual function of TargetInstrInfo.
X86FixupInstTuning.cpp may be integrated to MachineCombine.

llvm/test/CodeGen/X86/avxvnni-combine.ll

4

Sure, I'll precommit the test case first.

[LuoYuanke]

Rebase

[RKSimon]

llvm/test/CodeGen/X86/avxvnni-combine.ll

2

Add AVX512-VNNI support as well?

11

Please can you add test coverage for a case where there isn't the cross-loop dependency?

auto bar(int cnt, __m256i *c, __m256i b, __m256i *p) {
    for (int i = 0; i < cnt; ++i) {
        __m256i a = p[i];
        __m256i m = _mm256_madd_epi16 (b, a);
        c[i] = _mm256_add_epi32(m, c[i]);
    }
}

[LuoYuanke]

Add test case that there isn't the cross-loop dependency.

[LuoYuanke]

llvm/test/CodeGen/X86/avxvnni-combine.ll

160

vpmaddwd and vmovdqa (line 159) can be issued in parallel.

[goldstein.w.n]

llvm/lib/Target/X86/X86InstrInfo.cpp

9822

It seems.you are just querying for an opcode and replacing 1-1? Not that it can only be done in X86Fixup... But seems to fit there with less code.

[craig.topper]

llvm/lib/Target/X86/X86InstrInfo.cpp

9822

It’s replacing 1 instruction with 2 after doing analysis that consults the scheduling model.

[LuoYuanke]

Add (avx512, 128, 256, 512) support.

[LuoYuanke]

Update test case.

[LuoYuanke]

Fix an opcode bug.

[craig.topper]

llvm/lib/CodeGen/MachineCombiner.cpp

423–424

Put curly braces around the if body for consistency with the else body.

llvm/lib/Target/X86/X86InstrInfo.cpp

9790

The Vp prefix on these variable names isn't providing much value. You can probably drop it.

[LuoYuanke]

Address Craig's comments.

[LuoYuanke]

Update description

[LuoYuanke]

Rebase

[LuoYuanke]

ping

[pengfei]

LGTM with some nits.

llvm/lib/Target/X86/X86InstrInfo.cpp

9796–9799

indent.

I think we don't need to duplicate for ymm and zmm. And maybe better to choose rr version as example.

9850

Avoid auto

llvm/lib/Target/X86/X86InstrInfo.h

605

Nit: Should copy+paste the comment here like others.

610

ditto.

[LuoYuanke]

Address Phoebe's comments.

[RKSimon]

@LuoYuanke I'm seeing tests failures on my local machine with expensive checks enabled - one buildbot is reporting something similar: https://lab.llvm.org/buildbot/#/builders/16/builds/47297

[LuoYuanke]

In D148980#4305430, @RKSimon wrote:

@LuoYuanke I'm seeing tests failures on my local machine with expensive checks enabled - one buildbot is reporting something similar: https://lab.llvm.org/buildbot/#/builders/16/builds/47297

Thanks, @RKSimon for notifying. I'll revert the patch first.

[LuoYuanke]

There is conflict when I try to revert the patch. I fixed the issue with rG40222ddcf8f5.

llvm/include/llvm/CodeGen/MachineCombinerPattern.h

enum class MachineCombinerPattern {
   FMULv8i16_indexed_OP1,
   FMULv8i16_indexed_OP2,
 
   // RISCV FMADD, FMSUB, FNMSUB patterns
   FMADD_AX,
   FMADD_XA,
   FMSUB,
   FNMSUB,
+
+  // X86 VNNI
+  DPWSSD,
 };
 
 } // end namespace llvm
 
 #endif

llvm/lib/CodeGen/MachineCombiner.cpp

MachineCombiner() : MachineFunctionPass(ID) {
     initializeMachineCombinerPass(*PassRegistry::getPassRegistry());
   }
   void getAnalysisUsage(AnalysisUsage &AU) const override;
   bool runOnMachineFunction(MachineFunction &MF) override;
   StringRef getPassName() const override { return "Machine InstCombiner"; }
 
 private:
   bool combineInstructions(MachineBasicBlock *);
-  MachineInstr *getOperandDef(const MachineOperand &MO);
+  MachineInstr *getOperandDef(const MachineOperand &MO,
+                              SmallVectorImpl<MachineInstr *> &InsInstrs);
   bool isTransientMI(const MachineInstr *MI);
   unsigned getDepth(SmallVectorImpl<MachineInstr *> &InsInstrs,
                     DenseMap<unsigned, unsigned> &InstrIdxForVirtReg,
                     MachineTraceMetrics::Trace BlockTrace,
                     const MachineBasicBlock &MBB);
   unsigned getLatency(MachineInstr *Root, MachineInstr *NewRoot,
                       MachineTraceMetrics::Trace BlockTrace);
   bool
void MachineCombiner::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addPreserved<MachineLoopInfo>();
   AU.addRequired<MachineTraceMetrics>();
   AU.addPreserved<MachineTraceMetrics>();
   AU.addRequired<LazyMachineBlockFrequencyInfoPass>();
   AU.addRequired<ProfileSummaryInfoWrapperPass>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
-MachineInstr *MachineCombiner::getOperandDef(const MachineOperand &MO) {
+MachineInstr *
+MachineCombiner::getOperandDef(const MachineOperand &MO,
+                               SmallVectorImpl<MachineInstr *> &InsInstrs) {
   MachineInstr *DefInstr = nullptr;
   // We need a virtual register definition.
   if (MO.isReg() && MO.getReg().isVirtual())
     DefInstr = MRI->getUniqueVRegDef(MO.getReg());
+  // It is possible that the register is defined in new instructions.
+  if (!DefInstr) {
+    for (auto *MI : InsInstrs) {
+      for (const MachineOperand &DefMO : MI->operands()) {
+        if (!(DefMO.isReg() && DefMO.getReg().isVirtual()))
+          continue;
+        if (!DefMO.isDef())
+          continue;
+        if (DefMO.getReg() != MO.getReg())
+          continue;
+        DefInstr = MI;
+      }
+    }
+  }
   // PHI's have no depth etc.
   if (DefInstr && DefInstr->isPHI())
     DefInstr = nullptr;
   return DefInstr;
 }
 
 /// Return true if MI is unlikely to generate an actual target instruction.
 bool MachineCombiner::isTransientMI(const MachineInstr *MI) {
for (const MachineOperand &MO : InstrPtr->operands()) {
         assert(DefInstr &&
                "There must be a definition for a new virtual register");
         DepthOp = InstrDepth[II->second];
         int DefIdx = DefInstr->findRegisterDefOperandIdx(MO.getReg());
         int UseIdx = InstrPtr->findRegisterUseOperandIdx(MO.getReg());
         LatencyOp = TSchedModel.computeOperandLatency(DefInstr, DefIdx,
                                                       InstrPtr, UseIdx);
       } else {
-        MachineInstr *DefInstr = getOperandDef(MO);
+        MachineInstr *DefInstr = getOperandDef(MO, InsInstrs);
         if (DefInstr && (TII->getMachineCombinerTraceStrategy() !=
                              MachineTraceStrategy::TS_Local ||
                          DefInstr->getParent() == &MBB)) {
           DepthOp = BlockTrace.getInstrCycles(*DefInstr).Depth;
           if (!isTransientMI(DefInstr))
             LatencyOp = TSchedModel.computeOperandLatency(
                 DefInstr, DefInstr->findRegisterDefOperandIdx(MO.getReg()),
                 InstrPtr, InstrPtr->findRegisterUseOperandIdx(MO.getReg()));
if (getCombinerObjective(Pattern) == CombinerObjective::MustReduceDepth) {
     return NewRootDepth < RootDepth;
   }
 
   // A more flexible cost calculation for the critical path includes the slack
   // of the original code sequence. This may allow the transform to proceed
   // even if the instruction depths (data dependency cycles) become worse.
 
   // Account for the latency of the inserted and deleted instructions by
-  unsigned NewRootLatency, RootLatency;
-  std::tie(NewRootLatency, RootLatency) =
-      getLatenciesForInstrSequences(*Root, InsInstrs, DelInstrs, BlockTrace);
+  unsigned NewRootLatency = TSchedModel.computeInstrLatency(InsInstrs.back());
+  unsigned RootLatency = TSchedModel.computeInstrLatency(Root);
 
   unsigned RootSlack = BlockTrace.getInstrSlack(*Root);
   unsigned NewCycleCount = NewRootDepth + NewRootLatency;
   unsigned OldCycleCount =
       RootDepth + RootLatency + (SlackIsAccurate ? RootSlack : 0);
   LLVM_DEBUG(dbgs() << "\n\tNewRootLatency: " << NewRootLatency
                     << "\tRootLatency: " << RootLatency << "\n\tRootSlack: "
                     << RootSlack << " SlackIsAccurate=" << SlackIsAccurate

llvm/lib/Target/X86/X86InstrInfo.h

MachineInstr *commuteInstructionImpl(MachineInstr &MI, bool NewMI,
                                        unsigned CommuteOpIdx2) const override;
 
   /// If the specific machine instruction is a instruction that moves/copies
   /// value from one register to another register return destination and source
   /// registers as machine operands.
   std::optional<DestSourcePair>
   isCopyInstrImpl(const MachineInstr &MI) const override;
 
+  bool
+  getMachineCombinerPatterns(MachineInstr &Root,
+                             SmallVectorImpl<MachineCombinerPattern> &Patterns,
+                             bool DoRegPressureReduce) const override;
+
+  void genAlternativeCodeSequence(
+      MachineInstr &Root, MachineCombinerPattern Pattern,
+      SmallVectorImpl<MachineInstr *> &InsInstrs,
+      SmallVectorImpl<MachineInstr *> &DelInstrs,
+      DenseMap<unsigned, unsigned> &InstrIdxForVirtReg) const override;
 private:
   /// This is a helper for convertToThreeAddress for 8 and 16-bit instructions.
   /// We use 32-bit LEA to form 3-address code by promoting to a 32-bit
   /// super-register and then truncating back down to a 8/16-bit sub-register.
   MachineInstr *convertToThreeAddressWithLEA(unsigned MIOpc, MachineInstr &MI,
                                              LiveVariables *LV,
                                              LiveIntervals *LIS,
                                              bool Is8BitOp) const;

llvm/lib/Target/X86/X86InstrInfo.cpp

 #include "X86MachineFunctionInfo.h"
 #include "X86Subtarget.h"
 #include "X86TargetMachine.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/LivePhysRegs.h"
 #include "llvm/CodeGen/LiveVariables.h"
+#include "llvm/CodeGen/MachineCombinerPattern.h"
 #include "llvm/CodeGen/MachineConstantPool.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
if (C.CallConstructionID == MachineOutlinerTailCall) {
     It = MBB.insert(It,
                   BuildMI(MF, DebugLoc(), get(X86::CALL64pcrel32))
                       .addGlobalAddress(M.getNamedValue(MF.getName())));
   }
 
   return It;
 }
 
+bool X86InstrInfo::getMachineCombinerPatterns(
+    MachineInstr &Root, SmallVectorImpl<MachineCombinerPattern> &Patterns,
+    bool DoRegPressureReduce) const {
+  unsigned Opc = Root.getOpcode();
+  switch (Opc) {
+  default:
+    return TargetInstrInfo::getMachineCombinerPatterns(Root, Patterns,
+                                                       DoRegPressureReduce);
+  case X86::VPDPWSSDYrr:
+  case X86::VPDPWSSDYrm:
+    Patterns.push_back(MachineCombinerPattern::DPWSSD);
+    return true;
+  }
+}
+
+void X86InstrInfo::genAlternativeCodeSequence(
+    MachineInstr &Root, MachineCombinerPattern Pattern,
+    SmallVectorImpl<MachineInstr *> &InsInstrs,
+    SmallVectorImpl<MachineInstr *> &DelInstrs,
+    DenseMap<unsigned, unsigned> &InstrIdxForVirtReg) const {
+  MachineFunction *MF = Root.getMF();
+  // MachineRegisterInfo &MRI = MF->getRegInfo();
+  MachineRegisterInfo &RegInfo = MF->getRegInfo();
+
+  switch (Pattern) {
+  default:
+    // Reassociate instructions.
+    TargetInstrInfo::genAlternativeCodeSequence(Root, Pattern, InsInstrs,
+                                                DelInstrs, InstrIdxForVirtReg);
+    return;
+  case MachineCombinerPattern::DPWSSD:
+    unsigned Opc = Root.getOpcode();
+    switch (Opc) {
+    default:
+      break;
+    // vpdpwssd ymm2,ymm3,YMMWORD PTR [r8+0x20]
+    // -->
+    // vpmaddwd ymm3,ymm3,YMMWORD PTR [r8+0x20]
+    // vpaddd ymm2,ymm2,ymm3
+    case X86::VPDPWSSDYrr:
+    case X86::VPDPWSSDYrm:
+      unsigned NewOpc;
+      switch (Opc) {
+      case X86::VPDPWSSDYrr:
+            NewOpc = X86::VPMADDWDYrr;
+            break;
+      case X86::VPDPWSSDYrm:
+            NewOpc = X86::VPMADDWDYrm;
+            break;
+      }
+      // Create vpmaddwd.
+      auto *RC = RegInfo.getRegClass(Root.getOperand(0).getReg());
+      Register NewReg = RegInfo.createVirtualRegister(RC);
+      MachineInstr *VpMadd = Root.getMF()->CloneMachineInstr(&Root);
+      VpMadd->setDesc(get(NewOpc));
+      VpMadd->untieRegOperand(1);
+      VpMadd->removeOperand(1);
+      VpMadd->getOperand(0).setReg(NewReg);
+      // Create vpaddd.
+      Register DstReg = Root.getOperand(0).getReg();
+      bool IsKill = Root.getOperand(1).isKill();
+      MachineInstr *VpAdd =
+          BuildMI(*MF, MIMetadata(Root), get(X86::VPADDDYrr), DstReg)
+              .addReg(Root.getOperand(1).getReg(), getKillRegState(IsKill))
+              .addReg(VpMadd->getOperand(0).getReg(), getKillRegState(true));
+      InstrIdxForVirtReg.insert(std::make_pair(DstReg, 0));
+      InsInstrs.push_back(VpMadd);
+      InsInstrs.push_back(VpAdd);
+      DelInstrs.push_back(&Root);
+      break;
+    }
+  }
+}
+
 #define GET_INSTRINFO_HELPERS
 #include "X86GenInstrInfo.inc"

llvm/test/CodeGen/X86/avxvnni-combine.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mcpu=alderlake | FileCheck %s
+
+; __m256i foo(int cnt, __m256i c, __m256i b, __m256i *p) {
+;     for (int i = 0; i < cnt; ++i) {
+;         __m256i a = p[i];
+;         __m256i m = _mm256_madd_epi16 (b, a);
+;         c = _mm256_add_epi32(m, c);
+;     }
+;     return c;
+; }
+
+define dso_local <4 x i64> @foo(i32 %0, <4 x i64> %1, <4 x i64> %2, ptr %3) {
+; CHECK-LABEL: foo:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    testl %edi, %edi
+; CHECK-NEXT:    jle .LBB0_6
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    movl %edi, %edx
+; CHECK-NEXT:    movl %edx, %eax
+; CHECK-NEXT:    andl $3, %eax
+; CHECK-NEXT:    cmpl $4, %edi
+; CHECK-NEXT:    jae .LBB0_7
+; CHECK-NEXT:  # %bb.2:
+; CHECK-NEXT:    xorl %ecx, %ecx
+; CHECK-NEXT:    jmp .LBB0_3
+; CHECK-NEXT:  .LBB0_7:
+; CHECK-NEXT:    andl $-4, %edx
+; CHECK-NEXT:    leaq 96(%rsi), %rdi
+; CHECK-NEXT:    xorl %ecx, %ecx
+; CHECK-NEXT:    .p2align 4, 0x90
+; CHECK-NEXT:  .LBB0_8: # =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    {vex} vpdpwssd -96(%rdi), %ymm1, %ymm0
+; CHECK-NEXT:    vpmaddwd -64(%rdi), %ymm1, %ymm2
+; CHECK-NEXT:    vpmaddwd -32(%rdi), %ymm1, %ymm3
+; CHECK-NEXT:    vpaddd %ymm2, %ymm0, %ymm0
+; CHECK-NEXT:    vpaddd %ymm3, %ymm0, %ymm0
+; CHECK-NEXT:    vpmaddwd (%rdi), %ymm1, %ymm2
+; CHECK-NEXT:    vpaddd %ymm2, %ymm0, %ymm0
+; CHECK-NEXT:    addq $4, %rcx
+; CHECK-NEXT:    subq $-128, %rdi
+; CHECK-NEXT:    cmpq %rcx, %rdx
+; CHECK-NEXT:    jne .LBB0_8
+; CHECK-NEXT:  .LBB0_3:
+; CHECK-NEXT:    testq %rax, %rax
+; CHECK-NEXT:    je .LBB0_6
+; CHECK-NEXT:  # %bb.4: # %.preheader
+; CHECK-NEXT:    shlq $5, %rcx
+; CHECK-NEXT:    addq %rcx, %rsi
+; CHECK-NEXT:    shlq $5, %rax
+; CHECK-NEXT:    xorl %ecx, %ecx
+; CHECK-NEXT:    .p2align 4, 0x90
+; CHECK-NEXT:  .LBB0_5: # =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    {vex} vpdpwssd (%rsi,%rcx), %ymm1, %ymm0
+; CHECK-NEXT:    addq $32, %rcx
+; CHECK-NEXT:    cmpq %rcx, %rax
+; CHECK-NEXT:    jne .LBB0_5
+; CHECK-NEXT:  .LBB0_6:
+; CHECK-NEXT:    retq
+  %5 = icmp sgt i32 %0, 0
+  br i1 %5, label %6, label %33
+
+6:                                                ; preds = %4
+  %7 = bitcast <4 x i64> %2 to <16 x i16>
+  %8 = bitcast <4 x i64> %1 to <8 x i32>
+  %9 = zext i32 %0 to i64
+  %10 = and i64 %9, 3
+  %11 = icmp ult i32 %0, 4
+  br i1 %11, label %14, label %12
+
+12:                                               ; preds = %6
+  %13 = and i64 %9, 4294967292
+  br label %35
+
+14:                                               ; preds = %35, %6
+  %15 = phi <8 x i32> [ undef, %6 ], [ %57, %35 ]
+  %16 = phi i64 [ 0, %6 ], [ %58, %35 ]
+  %17 = phi <8 x i32> [ %8, %6 ], [ %57, %35 ]
+  %18 = icmp eq i64 %10, 0
+  br i1 %18, label %30, label %19
+
+19:                                               ; preds = %14, %19
+  %20 = phi i64 [ %27, %19 ], [ %16, %14 ]
+  %21 = phi <8 x i32> [ %26, %19 ], [ %17, %14 ]
+  %22 = phi i64 [ %28, %19 ], [ 0, %14 ]
+  %23 = getelementptr inbounds <4 x i64>, ptr %3, i64 %20
+  %24 = load <16 x i16>, ptr %23, align 32
+  %25 = tail call <8 x i32> @llvm.x86.avx2.pmadd.wd(<16 x i16> %7, <16 x i16> %24)
+  %26 = add <8 x i32> %25, %21
+  %27 = add nuw nsw i64 %20, 1
+  %28 = add i64 %22, 1
+  %29 = icmp eq i64 %28, %10
+  br i1 %29, label %30, label %19
+
+30:                                               ; preds = %19, %14
+  %31 = phi <8 x i32> [ %15, %14 ], [ %26, %19 ]
+  %32 = bitcast <8 x i32> %31 to <4 x i64>
+  br label %33
+
+33:                                               ; preds = %30, %4
+  %34 = phi <4 x i64> [ %32, %30 ], [ %1, %4 ]
+  ret <4 x i64> %34
+
+35:                                               ; preds = %35, %12
+  %36 = phi i64 [ 0, %12 ], [ %58, %35 ]
+  %37 = phi <8 x i32> [ %8, %12 ], [ %57, %35 ]
+  %38 = phi i64 [ 0, %12 ], [ %59, %35 ]
+  %39 = getelementptr inbounds <4 x i64>, ptr %3, i64 %36
+  %40 = load <16 x i16>, ptr %39, align 32
+  %41 = tail call <8 x i32> @llvm.x86.avx2.pmadd.wd(<16 x i16> %7, <16 x i16> %40)
+  %42 = add <8 x i32> %41, %37
+  %43 = or i64 %36, 1
+  %44 = getelementptr inbounds <4 x i64>, ptr %3, i64 %43
+  %45 = load <16 x i16>, ptr %44, align 32
+  %46 = tail call <8 x i32> @llvm.x86.avx2.pmadd.wd(<16 x i16> %7, <16 x i16> %45)
+  %47 = add <8 x i32> %46, %42
+  %48 = or i64 %36, 2
+  %49 = getelementptr inbounds <4 x i64>, ptr %3, i64 %48
+  %50 = load <16 x i16>, ptr %49, align 32
+  %51 = tail call <8 x i32> @llvm.x86.avx2.pmadd.wd(<16 x i16> %7, <16 x i16> %50)
+  %52 = add <8 x i32> %51, %47
+  %53 = or i64 %36, 3
+  %54 = getelementptr inbounds <4 x i64>, ptr %3, i64 %53
+  %55 = load <16 x i16>, ptr %54, align 32
+  %56 = tail call <8 x i32> @llvm.x86.avx2.pmadd.wd(<16 x i16> %7, <16 x i16> %55)
+  %57 = add <8 x i32> %56, %52
+  %58 = add nuw nsw i64 %36, 4
+  %59 = add i64 %38, 4
+  %60 = icmp eq i64 %59, %13
+  br i1 %60, label %14, label %35
+}
+
+declare <8 x i32> @llvm.x86.avx2.pmadd.wd(<16 x i16>, <16 x i16>)