[llvm] [LLVM][SVE] Add MachineInst pass to coalesce PTRUE instructions. (PR #204820)

Paul Walker via llvm-commits llvm-commits at lists.llvm.org
Fri Jun 26 09:05:18 PDT 2026


https://github.com/paulwalker-arm updated https://github.com/llvm/llvm-project/pull/204820

>From 8efd4fce706118c15d6a61f0643ef52e62289d3d Mon Sep 17 00:00:00 2001
From: Paul Walker <paul.walker at arm.com>
Date: Thu, 18 Jun 2026 10:20:59 +0000
Subject: [PATCH 1/3] [LLVM][SVE] Add MachineInst pass to coalesce PTRUE
 instructions.

SVE predicate registers contain a bit for each byte of a data
register. When operating on bigger element types the overlaping
predicate bits are grouped together but typically only the
least-significant-bit is read as part of the operation. This means
the value of the other bits does not affect the operation, making
it possible to use the result of PTRUE for a larger element typed
operation, assuming the predicate patterns are equivalent. For
example:

  ADD_S (PTRUE_H 31), Z0, Z1 == ADD_S (PTRUE_S 31), Z0, Z1

AArch64PTrueCoalescingPass uses this fact to reduce the number of
PTRUE instructions in a function, with the general trend towards
needing a single PTRUE based on the smallest element type in use.
---
 llvm/lib/Target/AArch64/AArch64.h             |  10 +-
 .../Target/AArch64/AArch64PTrueCoalescing.cpp | 226 ++++++++++++
 .../Target/AArch64/AArch64PassRegistry.def    |   1 +
 .../Target/AArch64/AArch64TargetMachine.cpp   |   5 +-
 llvm/lib/Target/AArch64/CMakeLists.txt        |   1 +
 llvm/test/CodeGen/AArch64/O3-pipeline.ll      |   3 +-
 .../CodeGen/AArch64/sve-ptrue-coalesce.mir    | 332 ++++++++++++++++++
 7 files changed, 575 insertions(+), 3 deletions(-)
 create mode 100644 llvm/lib/Target/AArch64/AArch64PTrueCoalescing.cpp
 create mode 100644 llvm/test/CodeGen/AArch64/sve-ptrue-coalesce.mir

diff --git a/llvm/lib/Target/AArch64/AArch64.h b/llvm/lib/Target/AArch64/AArch64.h
index e84c2bab20207..77762da13889d 100644
--- a/llvm/lib/Target/AArch64/AArch64.h
+++ b/llvm/lib/Target/AArch64/AArch64.h
@@ -69,7 +69,7 @@ FunctionPass *createAArch64BranchTargetsPass();
 FunctionPass *createAArch64CodeLayoutOptPass();
 FunctionPass *createAArch64MIPeepholeOptLegacyPass();
 FunctionPass *createAArch64PostCoalescerPass();
-
+FunctionPass *createAArch64PTrueCoalescingLegacyPass();
 FunctionPass *createAArch64CleanupLocalDynamicTLSPass();
 
 FunctionPass *createAArch64CollectLOHPass();
@@ -178,6 +178,7 @@ void initializeAArch64LoadStoreOptLegacyPass(PassRegistry &);
 void initializeAArch64LowerHomogeneousPrologEpilogLegacyPass(PassRegistry &);
 void initializeAArch64CodeLayoutOptPass(PassRegistry &);
 void initializeAArch64MIPeepholeOptLegacyPass(PassRegistry &);
+void initializeAArch64PTrueCoalescingLegacyPass(PassRegistry &);
 void initializeAArch64O0PreLegalizerCombinerLegacyPass(PassRegistry &);
 void initializeAArch64PostCoalescerLegacyPass(PassRegistry &);
 void initializeAArch64PostLegalizerCombinerLegacyPass(PassRegistry &);
@@ -292,6 +293,13 @@ class AArch64MIPeepholeOptPass
                         MachineFunctionAnalysisManager &MFAM);
 };
 
+class AArch64PTrueCoalescingPass
+    : public OptionalPassInfoMixin<AArch64PTrueCoalescingPass> {
+public:
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+};
+
 class AArch64ConditionOptimizerPass
     : public OptionalPassInfoMixin<AArch64ConditionOptimizerPass> {
 public:
diff --git a/llvm/lib/Target/AArch64/AArch64PTrueCoalescing.cpp b/llvm/lib/Target/AArch64/AArch64PTrueCoalescing.cpp
new file mode 100644
index 0000000000000..11eb7e874a684
--- /dev/null
+++ b/llvm/lib/Target/AArch64/AArch64PTrueCoalescing.cpp
@@ -0,0 +1,226 @@
+//===- AArch64PTrueCoalescing.cpp - Coalesce SVE PTRUEs ---------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass coalesces compatible all-active SVE PTRUE instructions.
+//
+// Consider two all-active PTRUE instructions X and Y with element sizes XSize
+// and YSize. If X dominates Y and XSize <= YSize, then every predicate bit that
+// Y sets is also set by X. In that case, uses of Y can be redirected to X as
+// long as each user of Y only reads predicate bits at YSize granularity or
+// larger.
+//
+// If the dominating PTRUE has a larger element size, we can coalesce the pair
+// by changing the dominating PTRUE to the smaller element size, provided that
+// all of its existing users are also safe with that granularity.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AArch64.h"
+#include "AArch64InstrInfo.h"
+#include "AArch64Subtarget.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/CodeGen/MachineDominators.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "aarch64-ptrue-coalesce"
+
+static cl::opt<bool> EnablePTrueCoalescing(
+    "aarch64-enable-ptrue-coalescing", cl::init(false), cl::Hidden,
+    cl::desc("Enable coalescing of compatible AArch64 SVE PTRUE instructions"));
+
+namespace {
+
+static bool isAllActivePTrue(const MachineInstr &MI) {
+  switch (MI.getOpcode()) {
+  default:
+    return false;
+  case AArch64::PTRUE_B:
+  case AArch64::PTRUE_H:
+  case AArch64::PTRUE_S:
+  case AArch64::PTRUE_D:
+    return MI.getOperand(1).getImm() == 31;
+  }
+}
+
+class AArch64PTrueCoalescingImpl {
+  const AArch64InstrInfo *TII = nullptr;
+  MachineRegisterInfo *MRI = nullptr;
+  MachineDominatorTree *MDT = nullptr;
+
+public:
+  explicit AArch64PTrueCoalescingImpl(MachineDominatorTree &MDT) : MDT(&MDT) {}
+
+  bool run(MachineFunction &MF);
+
+private:
+  bool allUsersSafeForElementSize(Register Reg, uint64_t ElementSize) const;
+  bool tryCoalesce(MachineInstr &DomPTrue, MachineInstr &PTrue) const;
+};
+
+class AArch64PTrueCoalescingLegacy : public MachineFunctionPass {
+public:
+  static char ID;
+
+  AArch64PTrueCoalescingLegacy() : MachineFunctionPass(ID) {}
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+  StringRef getPassName() const override { return "AArch64 PTRUE Coalescing"; }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesCFG();
+    AU.addRequired<MachineDominatorTreeWrapperPass>();
+    AU.addPreserved<MachineDominatorTreeWrapperPass>();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+};
+
+char AArch64PTrueCoalescingLegacy::ID = 0;
+
+} // end anonymous namespace
+
+INITIALIZE_PASS_BEGIN(AArch64PTrueCoalescingLegacy, DEBUG_TYPE,
+                      "AArch64 PTRUE Coalescing", false, false)
+INITIALIZE_PASS_DEPENDENCY(MachineDominatorTreeWrapperPass)
+INITIALIZE_PASS_END(AArch64PTrueCoalescingLegacy, DEBUG_TYPE,
+                    "AArch64 PTRUE Coalescing", false, false)
+
+bool AArch64PTrueCoalescingImpl::allUsersSafeForElementSize(
+    Register Reg, uint64_t ElementSize) const {
+  for (MachineOperand &UseMO : MRI->use_nodbg_operands(Reg)) {
+    if (UseMO.getSubReg())
+      return false;
+
+    MachineInstr *UseMI = UseMO.getParent();
+    uint64_t UseElementSize = TII->getElementSizeForOpcode(UseMI->getOpcode());
+    if (UseElementSize == AArch64::ElementSizeNone ||
+        UseElementSize < ElementSize)
+      return false;
+  }
+
+  return true;
+}
+
+bool AArch64PTrueCoalescingImpl::tryCoalesce(MachineInstr &DomPTrue,
+                                             MachineInstr &PTrue) const {
+  assert(isAllActivePTrue(DomPTrue) && "Expected all-active PTRUE");
+  assert(isAllActivePTrue(PTrue) && "Expected all-active PTRUE");
+
+  if (&DomPTrue == &PTrue || !MDT->dominates(&DomPTrue, &PTrue))
+    return false;
+
+  Register DomReg = DomPTrue.getOperand(0).getReg();
+  Register Reg = PTrue.getOperand(0).getReg();
+
+  uint64_t DomElementSize = TII->getElementSizeForOpcode(DomPTrue.getOpcode());
+  uint64_t ElementSize = TII->getElementSizeForOpcode(PTrue.getOpcode());
+  assert(DomElementSize != AArch64::ElementSizeNone &&
+         "PTRUE should have an element size");
+  assert(ElementSize != AArch64::ElementSizeNone &&
+         "PTRUE should have an element size");
+
+  if (!MRI->constrainRegClass(DomReg, MRI->getRegClass(Reg)))
+    return false;
+
+  bool MutateDomPTrue = false;
+  if (DomElementSize < ElementSize) {
+    // DomPTrue sets all lanes set by PTrue, plus extra lanes. Prefer to reuse
+    // DomPTrue as-is when PTrue's users do not observe those extra lanes.
+    if (!allUsersSafeForElementSize(Reg, ElementSize)) {
+      if (!allUsersSafeForElementSize(DomReg, ElementSize))
+        return false;
+      MutateDomPTrue = true;
+    }
+  } else if (DomElementSize > ElementSize) {
+    if (!allUsersSafeForElementSize(DomReg, ElementSize))
+      return false;
+    MutateDomPTrue = true;
+  }
+
+  LLVM_DEBUG(dbgs() << "Coalescing PTRUE: " << PTrue);
+  LLVM_DEBUG(dbgs() << "            with: " << DomPTrue);
+
+  if (MutateDomPTrue) {
+    LLVM_DEBUG(dbgs() << "        updated: " << DomPTrue);
+    DomPTrue.setDesc(TII->get(PTrue.getOpcode()));
+    LLVM_DEBUG(dbgs() << "             to: " << DomPTrue);
+  }
+
+  MRI->replaceRegWith(Reg, DomReg);
+  MRI->clearKillFlags(DomReg);
+  PTrue.eraseFromParent();
+  return true;
+}
+
+bool AArch64PTrueCoalescingImpl::run(MachineFunction &MF) {
+  if (!EnablePTrueCoalescing ||
+      !MF.getSubtarget<AArch64Subtarget>().isSVEorStreamingSVEAvailable())
+    return false;
+
+  TII = static_cast<const AArch64InstrInfo *>(MF.getSubtarget().getInstrInfo());
+  MRI = &MF.getRegInfo();
+
+  assert(MRI->isSSA() && "Expected to be run on SSA form!");
+
+  SmallVector<MachineInstr *, 8> PTrues;
+  for (MachineBasicBlock &MBB : MF)
+    for (MachineInstr &MI : MBB)
+      if (isAllActivePTrue(MI))
+        PTrues.push_back(&MI);
+
+  bool Changed = false;
+  auto tryCoalescePTrue = [&](MachineInstr *PTrue) {
+    for (MachineInstr *Candidate : PTrues)
+      if (tryCoalesce(*Candidate, *PTrue))
+        return true;
+    return false;
+  };
+
+  for (auto I = PTrues.begin(); I != PTrues.end();) {
+    if (tryCoalescePTrue(*I)) {
+      I = PTrues.erase(I);
+      Changed = true;
+    } else {
+      ++I;
+    }
+  }
+
+  return Changed;
+}
+
+bool AArch64PTrueCoalescingLegacy::runOnMachineFunction(MachineFunction &MF) {
+  MachineDominatorTree &MDT =
+      getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
+  return AArch64PTrueCoalescingImpl(MDT).run(MF);
+}
+
+FunctionPass *llvm::createAArch64PTrueCoalescingLegacyPass() {
+  return new AArch64PTrueCoalescingLegacy();
+}
+
+PreservedAnalyses
+AArch64PTrueCoalescingPass::run(MachineFunction &MF,
+                                MachineFunctionAnalysisManager &MFAM) {
+  MachineDominatorTree &MDT = MFAM.getResult<MachineDominatorTreeAnalysis>(MF);
+  const bool Changed = AArch64PTrueCoalescingImpl(MDT).run(MF);
+  if (!Changed)
+    return PreservedAnalyses::all();
+
+  PreservedAnalyses PA;
+  PA.preserveSet<CFGAnalyses>();
+  PA.preserve<MachineDominatorTreeAnalysis>();
+  return PA;
+}
diff --git a/llvm/lib/Target/AArch64/AArch64PassRegistry.def b/llvm/lib/Target/AArch64/AArch64PassRegistry.def
index 6f9eb76930d18..fd42f99b49a37 100644
--- a/llvm/lib/Target/AArch64/AArch64PassRegistry.def
+++ b/llvm/lib/Target/AArch64/AArch64PassRegistry.def
@@ -40,6 +40,7 @@ MACHINE_FUNCTION_PASS("aarch64-isel", AArch64DAGToDAGISelPass(*this))
 MACHINE_FUNCTION_PASS("aarch64-jump-tables", AArch64CompressJumpTablesPass())
 MACHINE_FUNCTION_PASS("aarch64-ldst-opt", AArch64LoadStoreOptPass())
 MACHINE_FUNCTION_PASS("aarch64-mi-peephole-opt", AArch64MIPeepholeOptPass())
+MACHINE_FUNCTION_PASS("aarch64-ptrue-coalesce", AArch64PTrueCoalescingPass())
 MACHINE_FUNCTION_PASS("aarch64-post-coalescer", AArch64PostCoalescerPass())
 MACHINE_FUNCTION_PASS("aarch64-post-select-optimize",
                       AArch64PostSelectOptimizePass())
diff --git a/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp b/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
index 24224a5b194ee..ee8fcddd5f68f 100644
--- a/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
+++ b/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
@@ -253,6 +253,7 @@ LLVMInitializeAArch64Target() {
   initializeAArch64ExpandPseudoLegacyPass(PR);
   initializeAArch64LoadStoreOptLegacyPass(PR);
   initializeAArch64MIPeepholeOptLegacyPass(PR);
+  initializeAArch64PTrueCoalescingLegacyPass(PR);
   initializeAArch64SIMDInstrOptLegacyPass(PR);
   initializeAArch64O0PreLegalizerCombinerLegacyPass(PR);
   initializeAArch64PreLegalizerCombinerLegacyPass(PR);
@@ -801,8 +802,10 @@ void AArch64PassConfig::addMachineSSAOptimization() {
   // Run default MachineSSAOptimization first.
   TargetPassConfig::addMachineSSAOptimization();
 
-  if (TM->getOptLevel() != CodeGenOptLevel::None)
+  if (TM->getOptLevel() != CodeGenOptLevel::None) {
     addPass(createAArch64MIPeepholeOptLegacyPass());
+    addPass(createAArch64PTrueCoalescingLegacyPass());
+  }
 }
 
 bool AArch64PassConfig::addILPOpts() {
diff --git a/llvm/lib/Target/AArch64/CMakeLists.txt b/llvm/lib/Target/AArch64/CMakeLists.txt
index 36be664e6eeb4..d953efea257e2 100644
--- a/llvm/lib/Target/AArch64/CMakeLists.txt
+++ b/llvm/lib/Target/AArch64/CMakeLists.txt
@@ -77,6 +77,7 @@ add_llvm_target(AArch64CodeGen
   AArch64PointerAuth.cpp
   AArch64PostCoalescerPass.cpp
   AArch64PromoteConstant.cpp
+  AArch64PTrueCoalescing.cpp
   AArch64PBQPRegAlloc.cpp
   AArch64RegisterInfo.cpp
   AArch64SMEAttributes.cpp
diff --git a/llvm/test/CodeGen/AArch64/O3-pipeline.ll b/llvm/test/CodeGen/AArch64/O3-pipeline.ll
index 1a0ffe234a236..534e18bdc22e2 100644
--- a/llvm/test/CodeGen/AArch64/O3-pipeline.ll
+++ b/llvm/test/CodeGen/AArch64/O3-pipeline.ll
@@ -169,6 +169,8 @@
 ; CHECK-NEXT:       Peephole Optimizations
 ; CHECK-NEXT:       Remove dead machine instructions
 ; CHECK-NEXT:       AArch64 MI Peephole Optimization pass
+; CHECK-NEXT:       MachineDominator Tree Construction
+; CHECK-NEXT:       AArch64 PTRUE Coalescing
 ; CHECK-NEXT:       AArch64 Dead register definitions
 ; CHECK-NEXT:       Detect Dead Lanes
 ; CHECK-NEXT:       Init Undef Pass
@@ -177,7 +179,6 @@
 ; CHECK-NEXT:       Live Variable Analysis
 ; CHECK-NEXT:       Eliminate PHI nodes for register allocation
 ; CHECK-NEXT:       Two-Address instruction pass
-; CHECK-NEXT:       MachineDominator Tree Construction
 ; CHECK-NEXT:       Slot index numbering
 ; CHECK-NEXT:       Live Interval Analysis
 ; CHECK-NEXT:       Register Coalescer
diff --git a/llvm/test/CodeGen/AArch64/sve-ptrue-coalesce.mir b/llvm/test/CodeGen/AArch64/sve-ptrue-coalesce.mir
new file mode 100644
index 0000000000000..26e09ba8b651e
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/sve-ptrue-coalesce.mir
@@ -0,0 +1,332 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 6
+# RUN: llc -mtriple=aarch64 -mattr=+sve -run-pass=aarch64-ptrue-coalesce -aarch64-enable-ptrue-coalescing -verify-machineinstrs -o - %s | FileCheck %s --check-prefixes=CHECK,ENABLED
+# RUN: llc -mtriple=aarch64 -mattr=+sve -passes=aarch64-ptrue-coalesce -aarch64-enable-ptrue-coalescing -o - %s | FileCheck %s --check-prefixes=CHECK,ENABLED
+# RUN: llc -mtriple=aarch64 -mattr=+sve -passes=aarch64-ptrue-coalesce -aarch64-enable-ptrue-coalescing=false -verify-machineinstrs -o - %s | FileCheck %s --check-prefixes=CHECK,DISABLED
+
+---
+name:            ptrue_b_replaces_ptrue_h_for_h_user
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: ppr_3b }
+  - { id: 1, class: ppr_3b }
+  - { id: 2, class: zpr }
+  - { id: 3, class: zpr }
+  - { id: 4, class: ppr }
+liveins:
+  - { reg: '$z0', virtual-reg: '%2' }
+  - { reg: '$z1', virtual-reg: '%3' }
+body:             |
+  bb.0:
+    liveins: $z0, $z1
+
+    ; ENABLED-LABEL: name: ptrue_b_replaces_ptrue_h_for_h_user
+    ; ENABLED: liveins: $z0, $z1
+    ; ENABLED-NEXT: {{  $}}
+    ; ENABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; ENABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; ENABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; ENABLED-NEXT: RET_ReallyLR
+    ;
+    ; DISABLED-LABEL: name: ptrue_b_replaces_ptrue_h_for_h_user
+    ; DISABLED: liveins: $z0, $z1
+    ; DISABLED-NEXT: {{  $}}
+    ; DISABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; DISABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; DISABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+    ; DISABLED-NEXT: [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
+    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; DISABLED-NEXT: RET_ReallyLR
+    %2:zpr = COPY $z0
+    %3:zpr = COPY $z1
+    %0:ppr_3b = PTRUE_B 31, implicit $vg
+    %1:ppr_3b = PTRUE_H 31, implicit $vg
+    %4:ppr = CMPEQ_PPzZZ_H %1, %2, %3, implicit-def dead $nzcv
+    RET_ReallyLR
+
+...
+---
+name:            keep_ptrue_h_for_b_user
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: ppr_3b }
+  - { id: 1, class: ppr_3b }
+  - { id: 2, class: zpr }
+  - { id: 3, class: zpr }
+  - { id: 4, class: ppr }
+liveins:
+  - { reg: '$z0', virtual-reg: '%2' }
+  - { reg: '$z1', virtual-reg: '%3' }
+body:             |
+  bb.0:
+    liveins: $z0, $z1
+
+    ; ENABLED-LABEL: name: keep_ptrue_h_for_b_user
+    ; ENABLED: liveins: $z0, $z1
+    ; ENABLED-NEXT: {{  $}}
+    ; ENABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; ENABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; ENABLED-NEXT: [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
+    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; ENABLED-NEXT: RET_ReallyLR
+    ;
+    ; DISABLED-LABEL: name: keep_ptrue_h_for_b_user
+    ; DISABLED: liveins: $z0, $z1
+    ; DISABLED-NEXT: {{  $}}
+    ; DISABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; DISABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; DISABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+    ; DISABLED-NEXT: [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
+    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; DISABLED-NEXT: RET_ReallyLR
+    %2:zpr = COPY $z0
+    %3:zpr = COPY $z1
+    %0:ppr_3b = PTRUE_B 31, implicit $vg
+    %1:ppr_3b = PTRUE_H 31, implicit $vg
+    %4:ppr = CMPEQ_PPzZZ_B %1, %2, %3, implicit-def dead $nzcv
+    RET_ReallyLR
+
+...
+---
+name:            ptrue_b_replaces_and_updates_ptrue_h
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: ppr_3b }
+  - { id: 1, class: ppr_3b }
+  - { id: 2, class: zpr }
+  - { id: 3, class: zpr }
+  - { id: 4, class: ppr }
+  - { id: 5, class: ppr }
+liveins:
+  - { reg: '$z0', virtual-reg: '%2' }
+  - { reg: '$z1', virtual-reg: '%3' }
+body:             |
+  bb.0:
+    liveins: $z0, $z1
+
+    ; ENABLED-LABEL: name: ptrue_b_replaces_and_updates_ptrue_h
+    ; ENABLED: liveins: $z0, $z1
+    ; ENABLED-NEXT: {{  $}}
+    ; ENABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; ENABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; ENABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; ENABLED-NEXT: RET_ReallyLR
+    ;
+    ; DISABLED-LABEL: name: ptrue_b_replaces_and_updates_ptrue_h
+    ; DISABLED: liveins: $z0, $z1
+    ; DISABLED-NEXT: {{  $}}
+    ; DISABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; DISABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; DISABLED-NEXT: [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
+    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; DISABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; DISABLED-NEXT: RET_ReallyLR
+    %2:zpr = COPY $z0
+    %3:zpr = COPY $z1
+    %0:ppr_3b = PTRUE_H 31, implicit $vg
+    %4:ppr = CMPEQ_PPzZZ_H %0, %2, %3, implicit-def dead $nzcv
+    %1:ppr_3b = PTRUE_B 31, implicit $vg
+    %5:ppr = CMPEQ_PPzZZ_B %1, %2, %3, implicit-def dead $nzcv
+    RET_ReallyLR
+
+...
+---
+name:            ptrue_b_replaces_ptrue_h_in_dominated_block
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: ppr_3b }
+  - { id: 1, class: ppr_3b }
+  - { id: 2, class: zpr }
+  - { id: 3, class: zpr }
+  - { id: 4, class: ppr }
+liveins:
+  - { reg: '$z0', virtual-reg: '%2' }
+  - { reg: '$z1', virtual-reg: '%3' }
+body:             |
+  ; ENABLED-LABEL: name: ptrue_b_replaces_ptrue_h_in_dominated_block
+  ; ENABLED: bb.0:
+  ; ENABLED-NEXT:   successors: %bb.1(0x80000000)
+  ; ENABLED-NEXT:   liveins: $z0, $z1
+  ; ENABLED-NEXT: {{  $}}
+  ; ENABLED-NEXT:   [[COPY:%[0-9]+]]:zpr = COPY $z0
+  ; ENABLED-NEXT:   [[COPY1:%[0-9]+]]:zpr = COPY $z1
+  ; ENABLED-NEXT:   [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+  ; ENABLED-NEXT:   B %bb.1
+  ; ENABLED-NEXT: {{  $}}
+  ; ENABLED-NEXT: bb.1:
+  ; ENABLED-NEXT:   [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+  ; ENABLED-NEXT:   RET_ReallyLR
+  ;
+  ; DISABLED-LABEL: name: ptrue_b_replaces_ptrue_h_in_dominated_block
+  ; DISABLED: bb.0:
+  ; DISABLED-NEXT:   successors: %bb.1(0x80000000)
+  ; DISABLED-NEXT:   liveins: $z0, $z1
+  ; DISABLED-NEXT: {{  $}}
+  ; DISABLED-NEXT:   [[COPY:%[0-9]+]]:zpr = COPY $z0
+  ; DISABLED-NEXT:   [[COPY1:%[0-9]+]]:zpr = COPY $z1
+  ; DISABLED-NEXT:   [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+  ; DISABLED-NEXT:   B %bb.1
+  ; DISABLED-NEXT: {{  $}}
+  ; DISABLED-NEXT: bb.1:
+  ; DISABLED-NEXT:   [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
+  ; DISABLED-NEXT:   [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+  ; DISABLED-NEXT:   RET_ReallyLR
+  bb.0:
+    successors: %bb.1
+    liveins: $z0, $z1
+
+    %2:zpr = COPY $z0
+    %3:zpr = COPY $z1
+    %0:ppr_3b = PTRUE_B 31, implicit $vg
+    B %bb.1
+
+  bb.1:
+    %1:ppr_3b = PTRUE_H 31, implicit $vg
+    %4:ppr = CMPEQ_PPzZZ_H %1, %2, %3, implicit-def dead $nzcv
+    RET_ReallyLR
+
+...
+
+---
+name:            keep_ptrue_h_when_b_user_would_observe_extra_lanes
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: ppr_3b }
+  - { id: 1, class: ppr_3b }
+  - { id: 2, class: zpr }
+  - { id: 3, class: zpr }
+  - { id: 4, class: ppr }
+  - { id: 5, class: ppr }
+liveins:
+  - { reg: '$z0', virtual-reg: '%2' }
+  - { reg: '$z1', virtual-reg: '%3' }
+body:             |
+  bb.0:
+    liveins: $z0, $z1
+
+    ; CHECK-LABEL: name: keep_ptrue_h_when_b_user_would_observe_extra_lanes
+    ; CHECK: liveins: $z0, $z1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; CHECK-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+    ; CHECK-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; CHECK-NEXT: [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
+    ; CHECK-NEXT: [[CMPEQ_PPzZZ_B1:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; CHECK-NEXT: RET_ReallyLR
+    %2:zpr = COPY $z0
+    %3:zpr = COPY $z1
+    %0:ppr_3b = PTRUE_B 31, implicit $vg
+    %4:ppr = CMPEQ_PPzZZ_B %0, %2, %3, implicit-def dead $nzcv
+    %1:ppr_3b = PTRUE_H 31, implicit $vg
+    %5:ppr = CMPEQ_PPzZZ_B %1, %2, %3, implicit-def dead $nzcv
+    RET_ReallyLR
+
+...
+---
+name:            keep_ptrue_s_when_d_user_cannot_be_narrowed
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: ppr_3b }
+  - { id: 1, class: ppr_3b }
+  - { id: 2, class: zpr }
+  - { id: 3, class: zpr }
+  - { id: 4, class: ppr }
+  - { id: 5, class: zpr }
+  - { id: 6, class: zpr }
+liveins:
+  - { reg: '$z0', virtual-reg: '%2' }
+  - { reg: '$z1', virtual-reg: '%3' }
+body:             |
+  bb.0:
+    liveins: $z0, $z1
+
+    ; CHECK-LABEL: name: keep_ptrue_s_when_d_user_cannot_be_narrowed
+    ; CHECK: liveins: $z0, $z1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; CHECK-NEXT: [[PTRUE_D:%[0-9]+]]:ppr_3b = PTRUE_D 31, implicit $vg
+    ; CHECK-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_D]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; CHECK-NEXT: [[PTRUE_S:%[0-9]+]]:ppr_3b = PTRUE_S 31, implicit $vg
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:zpr = IMPLICIT_DEF
+    ; CHECK-NEXT: [[SDIV_ZPZZ_S_UNDEF:%[0-9]+]]:zpr = SDIV_ZPZZ_S_UNDEF [[PTRUE_S]], [[DEF]], [[COPY]]
+    ; CHECK-NEXT: RET_ReallyLR
+    %2:zpr = COPY $z0
+    %3:zpr = COPY $z1
+    %0:ppr_3b = PTRUE_D 31, implicit $vg
+    %4:ppr = CMPEQ_PPzZZ_H %0, %2, %3, implicit-def dead $nzcv
+    %1:ppr_3b = PTRUE_S 31, implicit $vg
+    %6:zpr = IMPLICIT_DEF
+    %5:zpr = SDIV_ZPZZ_S_UNDEF %1, %6, %2
+    RET_ReallyLR
+
+...
+---
+name:            keep_ptrue_h_when_candidate_does_not_dominate
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: ppr_3b }
+  - { id: 1, class: ppr_3b }
+  - { id: 2, class: zpr }
+  - { id: 3, class: zpr }
+  - { id: 4, class: ppr }
+  - { id: 5, class: gpr32 }
+liveins:
+  - { reg: '$z0', virtual-reg: '%2' }
+  - { reg: '$z1', virtual-reg: '%3' }
+  - { reg: '$w0', virtual-reg: '%5' }
+body:             |
+  ; CHECK-LABEL: name: keep_ptrue_h_when_candidate_does_not_dominate
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.2(0x40000000), %bb.1(0x40000000)
+  ; CHECK-NEXT:   liveins: $z0, $z1, $w0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:zpr = COPY $z0
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:zpr = COPY $z1
+  ; CHECK-NEXT:   [[COPY2:%[0-9]+]]:gpr32 = COPY $w0
+  ; CHECK-NEXT:   CBZW [[COPY2]], %bb.2
+  ; CHECK-NEXT:   B %bb.1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   successors: %bb.2(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+  ; CHECK-NEXT:   B %bb.2
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
+  ; CHECK-NEXT:   [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+  ; CHECK-NEXT:   RET_ReallyLR
+  bb.0:
+    successors: %bb.2, %bb.1
+    liveins: $z0, $z1, $w0
+
+    %2:zpr = COPY $z0
+    %3:zpr = COPY $z1
+    %5:gpr32 = COPY $w0
+    CBZW %5, %bb.2
+    B %bb.1
+
+  bb.1:
+    successors: %bb.2
+
+    %0:ppr_3b = PTRUE_B 31, implicit $vg
+    B %bb.2
+
+  bb.2:
+    %1:ppr_3b = PTRUE_H 31, implicit $vg
+    %4:ppr = CMPEQ_PPzZZ_H %1, %2, %3, implicit-def dead $nzcv
+    RET_ReallyLR
+
+...

>From 25097aebe2116e9e62231e2b7f4fa0b05bf92060 Mon Sep 17 00:00:00 2001
From: Paul Walker <paul.walker at arm.com>
Date: Wed, 24 Jun 2026 17:30:46 +0100
Subject: [PATCH 2/3] Improve test coverage.

---
 .../CodeGen/AArch64/sve-ptrue-coalesce.mir    | 216 +++++++++++++++++-
 1 file changed, 206 insertions(+), 10 deletions(-)

diff --git a/llvm/test/CodeGen/AArch64/sve-ptrue-coalesce.mir b/llvm/test/CodeGen/AArch64/sve-ptrue-coalesce.mir
index 26e09ba8b651e..9fe54de67302c 100644
--- a/llvm/test/CodeGen/AArch64/sve-ptrue-coalesce.mir
+++ b/llvm/test/CodeGen/AArch64/sve-ptrue-coalesce.mir
@@ -1,7 +1,7 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 6
-# RUN: llc -mtriple=aarch64 -mattr=+sve -run-pass=aarch64-ptrue-coalesce -aarch64-enable-ptrue-coalescing -verify-machineinstrs -o - %s | FileCheck %s --check-prefixes=CHECK,ENABLED
-# RUN: llc -mtriple=aarch64 -mattr=+sve -passes=aarch64-ptrue-coalesce -aarch64-enable-ptrue-coalescing -o - %s | FileCheck %s --check-prefixes=CHECK,ENABLED
-# RUN: llc -mtriple=aarch64 -mattr=+sve -passes=aarch64-ptrue-coalesce -aarch64-enable-ptrue-coalescing=false -verify-machineinstrs -o - %s | FileCheck %s --check-prefixes=CHECK,DISABLED
+# RUN: llc -mtriple=aarch64 -mattr=+sve,+sme -run-pass=aarch64-ptrue-coalesce -aarch64-enable-ptrue-coalescing -verify-machineinstrs -o - %s | FileCheck %s --check-prefixes=CHECK,ENABLED
+# RUN: llc -mtriple=aarch64 -mattr=+sve,+sme -passes=aarch64-ptrue-coalesce -aarch64-enable-ptrue-coalescing -o - %s | FileCheck %s --check-prefixes=CHECK,ENABLED
+# RUN: llc -mtriple=aarch64 -mattr=+sve,+sme -passes=aarch64-ptrue-coalesce -aarch64-enable-ptrue-coalescing=false -verify-machineinstrs -o - %s | FileCheck %s --check-prefixes=CHECK,DISABLED
 
 ---
 name:            ptrue_b_replaces_ptrue_h_for_h_user
@@ -13,6 +13,7 @@ registers:
   - { id: 2, class: zpr }
   - { id: 3, class: zpr }
   - { id: 4, class: ppr }
+  - { id: 5, class: ppr }
 liveins:
   - { reg: '$z0', virtual-reg: '%2' }
   - { reg: '$z1', virtual-reg: '%3' }
@@ -27,6 +28,7 @@ body:             |
     ; ENABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
     ; ENABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
     ; ENABLED-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
     ; ENABLED-NEXT: RET_ReallyLR
     ;
     ; DISABLED-LABEL: name: ptrue_b_replaces_ptrue_h_for_h_user
@@ -37,12 +39,14 @@ body:             |
     ; DISABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
     ; DISABLED-NEXT: [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
     ; DISABLED-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
     ; DISABLED-NEXT: RET_ReallyLR
     %2:zpr = COPY $z0
     %3:zpr = COPY $z1
     %0:ppr_3b = PTRUE_B 31, implicit $vg
     %1:ppr_3b = PTRUE_H 31, implicit $vg
     %4:ppr = CMPEQ_PPzZZ_H %1, %2, %3, implicit-def dead $nzcv
+    %5:ppr = CMPEQ_PPzZZ_B %0, %2, %3, implicit-def dead $nzcv
     RET_ReallyLR
 
 ...
@@ -56,6 +60,7 @@ registers:
   - { id: 2, class: zpr }
   - { id: 3, class: zpr }
   - { id: 4, class: ppr }
+  - { id: 5, class: ppr }
 liveins:
   - { reg: '$z0', virtual-reg: '%2' }
   - { reg: '$z1', virtual-reg: '%3' }
@@ -70,6 +75,7 @@ body:             |
     ; ENABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
     ; ENABLED-NEXT: [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
     ; ENABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
     ; ENABLED-NEXT: RET_ReallyLR
     ;
     ; DISABLED-LABEL: name: keep_ptrue_h_for_b_user
@@ -80,12 +86,14 @@ body:             |
     ; DISABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
     ; DISABLED-NEXT: [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
     ; DISABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
     ; DISABLED-NEXT: RET_ReallyLR
     %2:zpr = COPY $z0
     %3:zpr = COPY $z1
     %0:ppr_3b = PTRUE_B 31, implicit $vg
     %1:ppr_3b = PTRUE_H 31, implicit $vg
     %4:ppr = CMPEQ_PPzZZ_B %1, %2, %3, implicit-def dead $nzcv
+    %5:ppr = CMPEQ_PPzZZ_H %0, %2, %3, implicit-def dead $nzcv
     RET_ReallyLR
 
 ...
@@ -146,6 +154,7 @@ registers:
   - { id: 2, class: zpr }
   - { id: 3, class: zpr }
   - { id: 4, class: ppr }
+  - { id: 5, class: ppr }
 liveins:
   - { reg: '$z0', virtual-reg: '%2' }
   - { reg: '$z1', virtual-reg: '%3' }
@@ -158,6 +167,7 @@ body:             |
   ; ENABLED-NEXT:   [[COPY:%[0-9]+]]:zpr = COPY $z0
   ; ENABLED-NEXT:   [[COPY1:%[0-9]+]]:zpr = COPY $z1
   ; ENABLED-NEXT:   [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+  ; ENABLED-NEXT:   [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
   ; ENABLED-NEXT:   B %bb.1
   ; ENABLED-NEXT: {{  $}}
   ; ENABLED-NEXT: bb.1:
@@ -172,6 +182,7 @@ body:             |
   ; DISABLED-NEXT:   [[COPY:%[0-9]+]]:zpr = COPY $z0
   ; DISABLED-NEXT:   [[COPY1:%[0-9]+]]:zpr = COPY $z1
   ; DISABLED-NEXT:   [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+  ; DISABLED-NEXT:   [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
   ; DISABLED-NEXT:   B %bb.1
   ; DISABLED-NEXT: {{  $}}
   ; DISABLED-NEXT: bb.1:
@@ -185,6 +196,7 @@ body:             |
     %2:zpr = COPY $z0
     %3:zpr = COPY $z1
     %0:ppr_3b = PTRUE_B 31, implicit $vg
+    %5:ppr = CMPEQ_PPzZZ_B %0, %2, %3, implicit-def dead $nzcv
     B %bb.1
 
   bb.1:
@@ -241,7 +253,7 @@ registers:
   - { id: 2, class: zpr }
   - { id: 3, class: zpr }
   - { id: 4, class: ppr }
-  - { id: 5, class: zpr }
+  - { id: 5, class: ppr }
   - { id: 6, class: zpr }
 liveins:
   - { reg: '$z0', virtual-reg: '%2' }
@@ -259,7 +271,7 @@ body:             |
     ; CHECK-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_D]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
     ; CHECK-NEXT: [[PTRUE_S:%[0-9]+]]:ppr_3b = PTRUE_S 31, implicit $vg
     ; CHECK-NEXT: [[DEF:%[0-9]+]]:zpr = IMPLICIT_DEF
-    ; CHECK-NEXT: [[SDIV_ZPZZ_S_UNDEF:%[0-9]+]]:zpr = SDIV_ZPZZ_S_UNDEF [[PTRUE_S]], [[DEF]], [[COPY]]
+    ; CHECK-NEXT: [[CMPEQ_PPzZZ_S:%[0-9]+]]:ppr = CMPEQ_PPzZZ_S [[PTRUE_S]], [[DEF]], [[COPY]], implicit-def dead $nzcv
     ; CHECK-NEXT: RET_ReallyLR
     %2:zpr = COPY $z0
     %3:zpr = COPY $z1
@@ -267,7 +279,7 @@ body:             |
     %4:ppr = CMPEQ_PPzZZ_H %0, %2, %3, implicit-def dead $nzcv
     %1:ppr_3b = PTRUE_S 31, implicit $vg
     %6:zpr = IMPLICIT_DEF
-    %5:zpr = SDIV_ZPZZ_S_UNDEF %1, %6, %2
+    %5:ppr = CMPEQ_PPzZZ_S %1, %6, %2, implicit-def dead $nzcv
     RET_ReallyLR
 
 ...
@@ -281,11 +293,12 @@ registers:
   - { id: 2, class: zpr }
   - { id: 3, class: zpr }
   - { id: 4, class: ppr }
-  - { id: 5, class: gpr32 }
+  - { id: 5, class: ppr }
+  - { id: 6, class: gpr32 }
 liveins:
   - { reg: '$z0', virtual-reg: '%2' }
   - { reg: '$z1', virtual-reg: '%3' }
-  - { reg: '$w0', virtual-reg: '%5' }
+  - { reg: '$w0', virtual-reg: '%6' }
 body:             |
   ; CHECK-LABEL: name: keep_ptrue_h_when_candidate_does_not_dominate
   ; CHECK: bb.0:
@@ -302,6 +315,7 @@ body:             |
   ; CHECK-NEXT:   successors: %bb.2(0x80000000)
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT:   [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+  ; CHECK-NEXT:   [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
   ; CHECK-NEXT:   B %bb.2
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT: bb.2:
@@ -314,14 +328,15 @@ body:             |
 
     %2:zpr = COPY $z0
     %3:zpr = COPY $z1
-    %5:gpr32 = COPY $w0
-    CBZW %5, %bb.2
+    %6:gpr32 = COPY $w0
+    CBZW %6, %bb.2
     B %bb.1
 
   bb.1:
     successors: %bb.2
 
     %0:ppr_3b = PTRUE_B 31, implicit $vg
+    %5:ppr = CMPEQ_PPzZZ_B %0, %2, %3, implicit-def dead $nzcv
     B %bb.2
 
   bb.2:
@@ -330,3 +345,184 @@ body:             |
     RET_ReallyLR
 
 ...
+---
+name:            ptrue_s_replaces_ptrue_d_for_d_user
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: ppr_3b }
+  - { id: 1, class: ppr_3b }
+  - { id: 2, class: zpr }
+  - { id: 3, class: zpr }
+  - { id: 4, class: ppr }
+  - { id: 5, class: ppr }
+liveins:
+  - { reg: '$z0', virtual-reg: '%2' }
+  - { reg: '$z1', virtual-reg: '%3' }
+body:             |
+  bb.0:
+    liveins: $z0, $z1
+
+    ; ENABLED-LABEL: name: ptrue_s_replaces_ptrue_d_for_d_user
+    ; ENABLED: liveins: $z0, $z1
+    ; ENABLED-NEXT: {{  $}}
+    ; ENABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; ENABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; ENABLED-NEXT: [[PTRUE_S:%[0-9]+]]:ppr_3b = PTRUE_S 31, implicit $vg
+    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_D:%[0-9]+]]:ppr = CMPEQ_PPzZZ_D [[PTRUE_S]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_S:%[0-9]+]]:ppr = CMPEQ_PPzZZ_S [[PTRUE_S]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; ENABLED-NEXT: RET_ReallyLR
+    ;
+    ; DISABLED-LABEL: name: ptrue_s_replaces_ptrue_d_for_d_user
+    ; DISABLED: liveins: $z0, $z1
+    ; DISABLED-NEXT: {{  $}}
+    ; DISABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; DISABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; DISABLED-NEXT: [[PTRUE_S:%[0-9]+]]:ppr_3b = PTRUE_S 31, implicit $vg
+    ; DISABLED-NEXT: [[PTRUE_D:%[0-9]+]]:ppr_3b = PTRUE_D 31, implicit $vg
+    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_D:%[0-9]+]]:ppr = CMPEQ_PPzZZ_D [[PTRUE_D]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_S:%[0-9]+]]:ppr = CMPEQ_PPzZZ_S [[PTRUE_S]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; DISABLED-NEXT: RET_ReallyLR
+    %2:zpr = COPY $z0
+    %3:zpr = COPY $z1
+    %0:ppr_3b = PTRUE_S 31, implicit $vg
+    %1:ppr_3b = PTRUE_D 31, implicit $vg
+    %4:ppr = CMPEQ_PPzZZ_D %1, %2, %3, implicit-def dead $nzcv
+    %5:ppr = CMPEQ_PPzZZ_S %0, %2, %3, implicit-def dead $nzcv
+    RET_ReallyLR
+
+...
+---
+name:            keep_non_all_ptrue
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: ppr_3b }
+  - { id: 1, class: ppr_3b }
+  - { id: 2, class: zpr }
+  - { id: 3, class: zpr }
+  - { id: 4, class: ppr }
+  - { id: 5, class: ppr }
+liveins:
+  - { reg: '$z0', virtual-reg: '%2' }
+  - { reg: '$z1', virtual-reg: '%3' }
+body:             |
+  bb.0:
+    liveins: $z0, $z1
+
+    ; CHECK-LABEL: name: keep_non_all_ptrue
+    ; CHECK: liveins: $z0, $z1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; CHECK-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 1, implicit $vg
+    ; CHECK-NEXT: [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
+    ; CHECK-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; CHECK-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; CHECK-NEXT: RET_ReallyLR
+    %2:zpr = COPY $z0
+    %3:zpr = COPY $z1
+    %0:ppr_3b = PTRUE_B 1, implicit $vg
+    %1:ppr_3b = PTRUE_H 31, implicit $vg
+    %4:ppr = CMPEQ_PPzZZ_B %0, %2, %3, implicit-def dead $nzcv
+    %5:ppr = CMPEQ_PPzZZ_H %1, %2, %3, implicit-def dead $nzcv
+    RET_ReallyLR
+
+...
+---
+name:            keep_ptrue_s_when_existing_b_user_would_observe_extra_lanes
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: ppr_3b }
+  - { id: 1, class: ppr_3b }
+  - { id: 2, class: zpr }
+  - { id: 3, class: zpr }
+  - { id: 4, class: zpr }
+  - { id: 5, class: ppr }
+liveins:
+  - { reg: '$z0', virtual-reg: '%2' }
+  - { reg: '$z1', virtual-reg: '%3' }
+body:             |
+  bb.0:
+    liveins: $z0, $z1
+
+    ; ENABLED-LABEL: name: keep_ptrue_s_when_existing_b_user_would_observe_extra_lanes
+    ; ENABLED: liveins: $z0, $z1
+    ; ENABLED-NEXT: {{  $}}
+    ; ENABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; ENABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; ENABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+    ; ENABLED-NEXT: [[ADD_ZPmZ_B:%[0-9]+]]:zpr = ADD_ZPmZ_B [[PTRUE_B]], [[COPY]], [[COPY1]]
+    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[ADD_ZPmZ_B]], [[COPY1]], implicit-def dead $nzcv
+    ; ENABLED-NEXT: RET_ReallyLR
+    ;
+    ; DISABLED-LABEL: name: keep_ptrue_s_when_existing_b_user_would_observe_extra_lanes
+    ; DISABLED: liveins: $z0, $z1
+    ; DISABLED-NEXT: {{  $}}
+    ; DISABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; DISABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; DISABLED-NEXT: [[PTRUE_S:%[0-9]+]]:ppr_3b = PTRUE_S 31, implicit $vg
+    ; DISABLED-NEXT: [[ADD_ZPmZ_B:%[0-9]+]]:zpr = ADD_ZPmZ_B [[PTRUE_S]], [[COPY]], [[COPY1]]
+    ; DISABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[ADD_ZPmZ_B]], [[COPY1]], implicit-def dead $nzcv
+    ; DISABLED-NEXT: RET_ReallyLR
+    %2:zpr = COPY $z0
+    %3:zpr = COPY $z1
+    %0:ppr_3b = PTRUE_S 31, implicit $vg
+    %4:zpr = ADD_ZPmZ_B %0, %2, %3
+    %1:ppr_3b = PTRUE_B 31, implicit $vg
+    %5:ppr = CMPEQ_PPzZZ_B %1, %4, %3, implicit-def dead $nzcv
+    RET_ReallyLR
+
+...
+---
+name:            do_notcoalesce_across_vg_definitions
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: ppr_3b }
+  - { id: 1, class: ppr_3b }
+  - { id: 2, class: zpr }
+  - { id: 3, class: zpr }
+  - { id: 4, class: ppr }
+  - { id: 5, class: ppr }
+liveins:
+  - { reg: '$z0', virtual-reg: '%2' }
+  - { reg: '$z1', virtual-reg: '%3' }
+body:             |
+  bb.0:
+    liveins: $z0, $z1
+
+    ; ENABLED-LABEL: name: do_notcoalesce_across_vg_definitions
+    ; ENABLED: liveins: $z0, $z1
+    ; ENABLED-NEXT: {{  $}}
+    ; ENABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; ENABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; ENABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; ENABLED-NEXT: MSRpstatesvcrImm1 1, 1, csr_aarch64_smstartstop, implicit-def dead $nzcv, implicit $vg, implicit-def $vg
+    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; ENABLED-NEXT: RET_ReallyLR
+    ;
+    ; DISABLED-LABEL: name: do_notcoalesce_across_vg_definitions
+    ; DISABLED: liveins: $z0, $z1
+    ; DISABLED-NEXT: {{  $}}
+    ; DISABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; DISABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; DISABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; DISABLED-NEXT: MSRpstatesvcrImm1 1, 1, csr_aarch64_smstartstop, implicit-def dead $nzcv, implicit $vg, implicit-def $vg
+    ; DISABLED-NEXT: [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
+    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; DISABLED-NEXT: RET_ReallyLR
+    %2:zpr = COPY $z0
+    %3:zpr = COPY $z1
+    %0:ppr_3b = PTRUE_B 31, implicit $vg
+    %5:ppr = CMPEQ_PPzZZ_B %0, %2, %3, implicit-def dead $nzcv
+    MSRpstatesvcrImm1 1, 1, csr_aarch64_smstartstop, implicit-def dead $nzcv, implicit $vg, implicit-def $vg
+    %1:ppr_3b = PTRUE_H 31, implicit $vg
+    %4:ppr = CMPEQ_PPzZZ_H %1, %2, %3, implicit-def dead $nzcv
+    RET_ReallyLR
+
+...

>From 747ea406c3c20f66ed9c29aea76c5d21d1aebe82 Mon Sep 17 00:00:00 2001
From: Paul Walker <paul.walker at arm.com>
Date: Wed, 24 Jun 2026 17:30:52 +0100
Subject: [PATCH 3/3] [AArch64][SVE] Rework PTRUE coalescing candidate tracking

Cache the relevant predicate information while collecting PTRUE
candidates, rather than repeatedly recomputing it during coalescing.
This also simplifies iteration: candidates are invalidated after
coalescing instead of being erased from the worklist.

Rework the core safety check to compare predicate element sizes
separately from the smallest element size observed by each predicate's
users. This makes the coalescing decision explicit for cases where a
PTRUE definition may be reused by users that interpret the predicate
at a different granularity.
---
 .../Target/AArch64/AArch64PTrueCoalescing.cpp | 186 +++++++++++-------
 .../CodeGen/AArch64/sve-ptrue-coalesce.mir    |  63 ++----
 2 files changed, 136 insertions(+), 113 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64PTrueCoalescing.cpp b/llvm/lib/Target/AArch64/AArch64PTrueCoalescing.cpp
index 11eb7e874a684..8811b7e8d35b7 100644
--- a/llvm/lib/Target/AArch64/AArch64PTrueCoalescing.cpp
+++ b/llvm/lib/Target/AArch64/AArch64PTrueCoalescing.cpp
@@ -43,18 +43,6 @@ static cl::opt<bool> EnablePTrueCoalescing(
 
 namespace {
 
-static bool isAllActivePTrue(const MachineInstr &MI) {
-  switch (MI.getOpcode()) {
-  default:
-    return false;
-  case AArch64::PTRUE_B:
-  case AArch64::PTRUE_H:
-  case AArch64::PTRUE_S:
-  case AArch64::PTRUE_D:
-    return MI.getOperand(1).getImm() == 31;
-  }
-}
-
 class AArch64PTrueCoalescingImpl {
   const AArch64InstrInfo *TII = nullptr;
   MachineRegisterInfo *MRI = nullptr;
@@ -66,8 +54,48 @@ class AArch64PTrueCoalescingImpl {
   bool run(MachineFunction &MF);
 
 private:
-  bool allUsersSafeForElementSize(Register Reg, uint64_t ElementSize) const;
-  bool tryCoalesce(MachineInstr &DomPTrue, MachineInstr &PTrue) const;
+  struct PredicateInfo {
+    // Instruction that created the predicate.
+    MachineInstr *MI = nullptr;
+    // Element size of the MI.
+    unsigned ElementSize = AArch64::ElementSizeNone;
+    // Smallest element size of all instructions that use the predicate.
+    unsigned SmallestUsedElementSize = AArch64::ElementSizeNone;
+
+    bool isValid() const {
+      return MI && ElementSize != AArch64::ElementSizeNone &&
+             SmallestUsedElementSize != AArch64::ElementSizeNone;
+    }
+
+    void invalidate() {
+      assert(isValid());
+      MI = nullptr;
+    }
+  };
+
+  std::optional<PredicateInfo> createPredicateInfo(MachineInstr &MI) const {
+    // TODO: Extend support beyond "PTRUE all"?
+    if (!isPTrueOpcode(MI.getOpcode()) || MI.getOperand(1).getImm() != 31)
+      return std::nullopt;
+
+    Register Pred = MI.getOperand(0).getReg();
+    unsigned SmallestUsedElementSize = getSmallestElementSizeInUse(Pred);
+    unsigned ElementSize = TII->getElementSizeForOpcode(MI.getOpcode());
+
+    if (ElementSize == AArch64::ElementSizeNone ||
+        SmallestUsedElementSize == AArch64::ElementSizeNone)
+      return std::nullopt;
+
+    return PredicateInfo{&MI, ElementSize, SmallestUsedElementSize};
+  }
+
+  // Return the smallest element size of all instructions that use Reg, or
+  // AArch64::ElementSizeNone when unknown.
+  unsigned getSmallestElementSizeInUse(Register Reg) const;
+
+  // Try to replace uses of CanPred with DomPred. In some cases that means
+  // modifying DomPred to support smaller element types.
+  bool tryCoalesce(PredicateInfo &DomPred, PredicateInfo &CanPred) const;
 };
 
 class AArch64PTrueCoalescingLegacy : public MachineFunctionPass {
@@ -98,70 +126,80 @@ INITIALIZE_PASS_DEPENDENCY(MachineDominatorTreeWrapperPass)
 INITIALIZE_PASS_END(AArch64PTrueCoalescingLegacy, DEBUG_TYPE,
                     "AArch64 PTRUE Coalescing", false, false)
 
-bool AArch64PTrueCoalescingImpl::allUsersSafeForElementSize(
-    Register Reg, uint64_t ElementSize) const {
+unsigned
+AArch64PTrueCoalescingImpl::getSmallestElementSizeInUse(Register Reg) const {
+  // SSA form only applies to virtual registers.
+  if (!Reg.isVirtual())
+    return AArch64::ElementSizeNone;
+
+  unsigned SmallestElementSize = AArch64::ElementSizeNone;
+
   for (MachineOperand &UseMO : MRI->use_nodbg_operands(Reg)) {
     if (UseMO.getSubReg())
-      return false;
+      return AArch64::ElementSizeNone;
 
     MachineInstr *UseMI = UseMO.getParent();
-    uint64_t UseElementSize = TII->getElementSizeForOpcode(UseMI->getOpcode());
-    if (UseElementSize == AArch64::ElementSizeNone ||
-        UseElementSize < ElementSize)
-      return false;
+
+    unsigned ElementSize = TII->getElementSizeForOpcode(UseMI->getOpcode());
+    if (ElementSize == AArch64::ElementSizeNone)
+      return AArch64::ElementSizeNone;
+
+    if (SmallestElementSize == AArch64::ElementSizeNone ||
+        SmallestElementSize > ElementSize)
+      SmallestElementSize = ElementSize;
   }
 
-  return true;
+  return SmallestElementSize;
 }
 
-bool AArch64PTrueCoalescingImpl::tryCoalesce(MachineInstr &DomPTrue,
-                                             MachineInstr &PTrue) const {
-  assert(isAllActivePTrue(DomPTrue) && "Expected all-active PTRUE");
-  assert(isAllActivePTrue(PTrue) && "Expected all-active PTRUE");
+bool AArch64PTrueCoalescingImpl::tryCoalesce(PredicateInfo &DomPI,
+                                             PredicateInfo &CanPI) const {
+  assert(DomPI.isValid() && CanPI.isValid());
+  MachineInstr *DomMI = DomPI.MI;
+  MachineInstr *CanMI = CanPI.MI;
 
-  if (&DomPTrue == &PTrue || !MDT->dominates(&DomPTrue, &PTrue))
+  if (DomMI == CanMI || !MDT->dominates(DomMI, CanMI))
     return false;
 
-  Register DomReg = DomPTrue.getOperand(0).getReg();
-  Register Reg = PTrue.getOperand(0).getReg();
-
-  uint64_t DomElementSize = TII->getElementSizeForOpcode(DomPTrue.getOpcode());
-  uint64_t ElementSize = TII->getElementSizeForOpcode(PTrue.getOpcode());
-  assert(DomElementSize != AArch64::ElementSizeNone &&
-         "PTRUE should have an element size");
-  assert(ElementSize != AArch64::ElementSizeNone &&
-         "PTRUE should have an element size");
-
-  if (!MRI->constrainRegClass(DomReg, MRI->getRegClass(Reg)))
-    return false;
+  // A predicate's observable shape is the larger of the element size of the
+  // instruction writing the predicate and the one reading it. First check if
+  // DomPI can replace CanPI as-is for CanPI's users. If not, try changing DomPI
+  // to CanPI's element size, but only if DomPI's existing users would observe
+  // the same shape after that change.
 
   bool MutateDomPTrue = false;
-  if (DomElementSize < ElementSize) {
-    // DomPTrue sets all lanes set by PTrue, plus extra lanes. Prefer to reuse
-    // DomPTrue as-is when PTrue's users do not observe those extra lanes.
-    if (!allUsersSafeForElementSize(Reg, ElementSize)) {
-      if (!allUsersSafeForElementSize(DomReg, ElementSize))
-        return false;
-      MutateDomPTrue = true;
-    }
-  } else if (DomElementSize > ElementSize) {
-    if (!allUsersSafeForElementSize(DomReg, ElementSize))
+  if (std::max(CanPI.ElementSize, CanPI.SmallestUsedElementSize) !=
+      std::max(DomPI.ElementSize, CanPI.SmallestUsedElementSize)) {
+    if (std::max(CanPI.ElementSize, DomPI.SmallestUsedElementSize) !=
+        std::max(DomPI.ElementSize, DomPI.SmallestUsedElementSize))
       return false;
+
     MutateDomPTrue = true;
   }
 
-  LLVM_DEBUG(dbgs() << "Coalescing PTRUE: " << PTrue);
-  LLVM_DEBUG(dbgs() << "            with: " << DomPTrue);
+  Register DomReg = DomMI->getOperand(0).getReg();
+  Register CanReg = CanMI->getOperand(0).getReg();
+  if (!MRI->constrainRegClass(DomReg, MRI->getRegClass(CanReg)))
+    return false;
+
+  LLVM_DEBUG(dbgs() << "Coalescing PTRUE: " << CanMI);
+  LLVM_DEBUG(dbgs() << "            with: " << DomMI);
 
   if (MutateDomPTrue) {
-    LLVM_DEBUG(dbgs() << "        updated: " << DomPTrue);
-    DomPTrue.setDesc(TII->get(PTrue.getOpcode()));
-    LLVM_DEBUG(dbgs() << "             to: " << DomPTrue);
+    LLVM_DEBUG(dbgs() << "        updated: " << DomMI);
+    DomMI->setDesc(TII->get(CanMI->getOpcode()));
+    DomPI.ElementSize = CanPI.ElementSize;
+    LLVM_DEBUG(dbgs() << "             to: " << DomMI);
   }
 
-  MRI->replaceRegWith(Reg, DomReg);
+  MRI->replaceRegWith(CanReg, DomReg);
   MRI->clearKillFlags(DomReg);
-  PTrue.eraseFromParent();
+  CanMI->eraseFromParent();
+
+  // Update DomPI based on uses inherited from CanPI.
+  if (CanPI.SmallestUsedElementSize < DomPI.SmallestUsedElementSize)
+    DomPI.SmallestUsedElementSize = CanPI.SmallestUsedElementSize;
+  CanPI.invalidate();
   return true;
 }
 
@@ -175,26 +213,32 @@ bool AArch64PTrueCoalescingImpl::run(MachineFunction &MF) {
 
   assert(MRI->isSSA() && "Expected to be run on SSA form!");
 
-  SmallVector<MachineInstr *, 8> PTrues;
+  // TODO: Until we prove candidates share the same VG definition, do not
+  // coalesce in functions that define VG.
+  if (!MRI->def_empty(AArch64::VG))
+    return false;
+
+  // A list of predicate setting instructions with some usage information.
+  SmallVector<PredicateInfo, 8> PIs;
+
+  // Build a list of predicates whose uses all have a known size.
   for (MachineBasicBlock &MBB : MF)
     for (MachineInstr &MI : MBB)
-      if (isAllActivePTrue(MI))
-        PTrues.push_back(&MI);
+      if (auto PI = createPredicateInfo(MI))
+        PIs.push_back(*PI);
 
+  LLVM_DEBUG(dbgs() << "Coalescable PTRUE candidates: " << PIs.size() << "\n");
   bool Changed = false;
-  auto tryCoalescePTrue = [&](MachineInstr *PTrue) {
-    for (MachineInstr *Candidate : PTrues)
-      if (tryCoalesce(*Candidate, *PTrue))
-        return true;
-    return false;
-  };
 
-  for (auto I = PTrues.begin(); I != PTrues.end();) {
-    if (tryCoalescePTrue(*I)) {
-      I = PTrues.erase(I);
-      Changed = true;
-    } else {
-      ++I;
+  for (PredicateInfo &DominantPI : PIs) {
+    if (!DominantPI.isValid())
+      continue;
+
+    for (PredicateInfo &CandidatePI : PIs) {
+      if (!CandidatePI.isValid())
+        continue;
+
+      Changed |= tryCoalesce(DominantPI, CandidatePI);
     }
   }
 
diff --git a/llvm/test/CodeGen/AArch64/sve-ptrue-coalesce.mir b/llvm/test/CodeGen/AArch64/sve-ptrue-coalesce.mir
index 9fe54de67302c..c9c4284cf40d4 100644
--- a/llvm/test/CodeGen/AArch64/sve-ptrue-coalesce.mir
+++ b/llvm/test/CodeGen/AArch64/sve-ptrue-coalesce.mir
@@ -447,26 +447,16 @@ body:             |
   bb.0:
     liveins: $z0, $z1
 
-    ; ENABLED-LABEL: name: keep_ptrue_s_when_existing_b_user_would_observe_extra_lanes
-    ; ENABLED: liveins: $z0, $z1
-    ; ENABLED-NEXT: {{  $}}
-    ; ENABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
-    ; ENABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
-    ; ENABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
-    ; ENABLED-NEXT: [[ADD_ZPmZ_B:%[0-9]+]]:zpr = ADD_ZPmZ_B [[PTRUE_B]], [[COPY]], [[COPY1]]
-    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[ADD_ZPmZ_B]], [[COPY1]], implicit-def dead $nzcv
-    ; ENABLED-NEXT: RET_ReallyLR
-    ;
-    ; DISABLED-LABEL: name: keep_ptrue_s_when_existing_b_user_would_observe_extra_lanes
-    ; DISABLED: liveins: $z0, $z1
-    ; DISABLED-NEXT: {{  $}}
-    ; DISABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
-    ; DISABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
-    ; DISABLED-NEXT: [[PTRUE_S:%[0-9]+]]:ppr_3b = PTRUE_S 31, implicit $vg
-    ; DISABLED-NEXT: [[ADD_ZPmZ_B:%[0-9]+]]:zpr = ADD_ZPmZ_B [[PTRUE_S]], [[COPY]], [[COPY1]]
-    ; DISABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
-    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[ADD_ZPmZ_B]], [[COPY1]], implicit-def dead $nzcv
-    ; DISABLED-NEXT: RET_ReallyLR
+    ; CHECK-LABEL: name: keep_ptrue_s_when_existing_b_user_would_observe_extra_lanes
+    ; CHECK: liveins: $z0, $z1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; CHECK-NEXT: [[PTRUE_S:%[0-9]+]]:ppr_3b = PTRUE_S 31, implicit $vg
+    ; CHECK-NEXT: [[ADD_ZPmZ_B:%[0-9]+]]:zpr = ADD_ZPmZ_B [[PTRUE_S]], [[COPY]], [[COPY1]]
+    ; CHECK-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+    ; CHECK-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[ADD_ZPmZ_B]], [[COPY1]], implicit-def dead $nzcv
+    ; CHECK-NEXT: RET_ReallyLR
     %2:zpr = COPY $z0
     %3:zpr = COPY $z1
     %0:ppr_3b = PTRUE_S 31, implicit $vg
@@ -494,28 +484,17 @@ body:             |
   bb.0:
     liveins: $z0, $z1
 
-    ; ENABLED-LABEL: name: do_notcoalesce_across_vg_definitions
-    ; ENABLED: liveins: $z0, $z1
-    ; ENABLED-NEXT: {{  $}}
-    ; ENABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
-    ; ENABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
-    ; ENABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
-    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
-    ; ENABLED-NEXT: MSRpstatesvcrImm1 1, 1, csr_aarch64_smstartstop, implicit-def dead $nzcv, implicit $vg, implicit-def $vg
-    ; ENABLED-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
-    ; ENABLED-NEXT: RET_ReallyLR
-    ;
-    ; DISABLED-LABEL: name: do_notcoalesce_across_vg_definitions
-    ; DISABLED: liveins: $z0, $z1
-    ; DISABLED-NEXT: {{  $}}
-    ; DISABLED-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
-    ; DISABLED-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
-    ; DISABLED-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
-    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
-    ; DISABLED-NEXT: MSRpstatesvcrImm1 1, 1, csr_aarch64_smstartstop, implicit-def dead $nzcv, implicit $vg, implicit-def $vg
-    ; DISABLED-NEXT: [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
-    ; DISABLED-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
-    ; DISABLED-NEXT: RET_ReallyLR
+    ; CHECK-LABEL: name: do_notcoalesce_across_vg_definitions
+    ; CHECK: liveins: $z0, $z1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:zpr = COPY $z1
+    ; CHECK-NEXT: [[PTRUE_B:%[0-9]+]]:ppr_3b = PTRUE_B 31, implicit $vg
+    ; CHECK-NEXT: [[CMPEQ_PPzZZ_B:%[0-9]+]]:ppr = CMPEQ_PPzZZ_B [[PTRUE_B]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; CHECK-NEXT: MSRpstatesvcrImm1 1, 1, csr_aarch64_smstartstop, implicit-def dead $nzcv, implicit $vg, implicit-def $vg
+    ; CHECK-NEXT: [[PTRUE_H:%[0-9]+]]:ppr_3b = PTRUE_H 31, implicit $vg
+    ; CHECK-NEXT: [[CMPEQ_PPzZZ_H:%[0-9]+]]:ppr = CMPEQ_PPzZZ_H [[PTRUE_H]], [[COPY]], [[COPY1]], implicit-def dead $nzcv
+    ; CHECK-NEXT: RET_ReallyLR
     %2:zpr = COPY $z0
     %3:zpr = COPY $z1
     %0:ppr_3b = PTRUE_B 31, implicit $vg



More information about the llvm-commits mailing list