[llvm] [X86] Record the enclosed register in X86DomainReassignment::buildClosure (PR #202534)

[via llvm-commits]

https://github.com/mbhade-amd updated https://github.com/llvm/llvm-project/pull/202534

>From 3ad9da96695715fe3ed25e9584f5cc7b2002ba4c Mon Sep 17 00:00:00 2001
From: mbhade <mbhade at amd.com>
Date: Tue, 9 Jun 2026 13:10:22 +0530
Subject: [PATCH 1/2] [X86] Record the enclosed register in
 X86DomainReassignment::buildClosure

buildClosure recorded the seed register Reg in the function-wide
EnclosedEdges map on every worklist iteration instead of CurReg, the
register actually being added to the closure. EnclosedEdges therefore
only ever contained the seed of each closure.

The driver loop in runOnMachineFunction skips registers already present
in EnclosedEdges before starting a new closure. Because only seeds were
recorded, every non-seed member of an already-built closure looked like a
fresh seed, so a redundant closure was built for it and then immediately
discarded by the EnclosedInstrs cross-closure check. The emitted code is
unchanged; the pass just performed redundant work proportional to closure
size.

Key EnclosedEdges by CurReg so each enclosed register is recorded once.
---
 llvm/lib/Target/X86/X86DomainReassignment.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Target/X86/X86DomainReassignment.cpp b/llvm/lib/Target/X86/X86DomainReassignment.cpp
index a9f68393aa3e2..dae1cd2640378 100644
--- a/llvm/lib/Target/X86/X86DomainReassignment.cpp
+++ b/llvm/lib/Target/X86/X86DomainReassignment.cpp
@@ -553,7 +553,7 @@ void X86DomainReassignmentImpl::buildClosure(Closure &C, Register Reg) {
     // Register already in this closure.
     if (!C.insertEdge(CurReg))
       continue;
-    EnclosedEdges[Reg] = C.getID();
+    EnclosedEdges[CurReg] = C.getID();
 
     MachineInstr *DefMI = MRI->getVRegDef(CurReg);
     if (!encloseInstr(C, DefMI))

>From c066292c38ec3ff76a13bf19f8d5eca4fda161df Mon Sep 17 00:00:00 2001
From: mbhade <mbhade at amd.com>
Date: Wed, 10 Jun 2026 19:13:32 +0530
Subject: [PATCH 2/2] [X86] Add NumClosuresBuilt stat and closure-count
 regression test Count every closure built by X86DomainReassignment and add a
 -stats test. test_8bitops builds 12 closures before the EnclosedEdges fix and
 2 after; converted count is unchanged.

---
 llvm/lib/Target/X86/X86DomainReassignment.cpp |  2 +
 .../X86/domain-reassignment-closure-stats.mir | 92 +++++++++++++++++++
 2 files changed, 94 insertions(+)
 create mode 100644 llvm/test/CodeGen/X86/domain-reassignment-closure-stats.mir

diff --git a/llvm/lib/Target/X86/X86DomainReassignment.cpp b/llvm/lib/Target/X86/X86DomainReassignment.cpp
index dae1cd2640378..8afa1cbc78e03 100644
--- a/llvm/lib/Target/X86/X86DomainReassignment.cpp
+++ b/llvm/lib/Target/X86/X86DomainReassignment.cpp
@@ -33,6 +33,7 @@ using namespace llvm;
 #define DEBUG_TYPE "x86-domain-reassignment"
 
 STATISTIC(NumClosuresConverted, "Number of closures converted by the pass");
+STATISTIC(NumClosuresBuilt, "Number of closures built by the pass");
 
 static cl::opt<bool> DisableX86DomainReassignment(
     "disable-x86-domain-reassignment", cl::Hidden,
@@ -812,6 +813,7 @@ bool X86DomainReassignmentImpl::runOnMachineFunction(MachineFunction &MF) {
     // Calculate closure starting with Reg.
     Closure C(ClosureID++, {MaskDomain});
     buildClosure(C, Reg);
+    ++NumClosuresBuilt;
 
     // Collect all closures that can potentially be converted.
     if (!C.empty() && C.isLegal(MaskDomain))
diff --git a/llvm/test/CodeGen/X86/domain-reassignment-closure-stats.mir b/llvm/test/CodeGen/X86/domain-reassignment-closure-stats.mir
new file mode 100644
index 0000000000000..85c401ce1f807
--- /dev/null
+++ b/llvm/test/CodeGen/X86/domain-reassignment-closure-stats.mir
@@ -0,0 +1,92 @@
+# REQUIRES: asserts
+# RUN: llc -run-pass x86-domain-reassignment -mtriple=x86_64-unknown-unknown \
+# RUN:   -mattr=+avx512f,+avx512bw,+avx512dq -stats %s -o /dev/null 2>&1 \
+# RUN:   | FileCheck %s
+
+# Regression test for buildClosure recording the seed register instead of the
+# register actually added to the closure. Before the fix EnclosedEdges only ever
+# held the seed of each closure, so every other member of an already-built
+# closure looked like a fresh seed and a redundant closure was built for it and
+# immediately discarded. The generated code is identical either way (the
+# converted-closure count is unaffected); only the amount of redundant work
+# changes. For test_8bitops the closure has eight members: buggy builds 12
+# closures, the fix builds 2. The number of converted closures stays at 1.
+
+# The leading {{^ *}} anchor prevents a stale "12" count from matching "2".
+# CHECK-DAG: {{^ *}}2 x86-domain-reassignment - Number of closures built by the pass
+# CHECK-DAG: {{^ *}}1 x86-domain-reassignment - Number of closures converted by the pass
+
+--- |
+  target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+  target triple = "x86_64-unknown-unknown"
+
+  define void @test_8bitops() #0 {
+    ret void
+  }
+
+  attributes #0 = { "target-cpu"="skylake-avx512" }
+...
+---
+name:            test_8bitops
+alignment:       16
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: gr64, preferred-register: '' }
+  - { id: 1, class: vr512, preferred-register: '' }
+  - { id: 2, class: vr512, preferred-register: '' }
+  - { id: 3, class: vr512, preferred-register: '' }
+  - { id: 4, class: vr512, preferred-register: '' }
+  - { id: 5, class: vk8, preferred-register: '' }
+  - { id: 6, class: gr32, preferred-register: '' }
+  - { id: 7, class: gr8, preferred-register: '' }
+  - { id: 8, class: gr32, preferred-register: '' }
+  - { id: 9, class: gr32, preferred-register: '' }
+  - { id: 10, class: vk8wm, preferred-register: '' }
+  - { id: 11, class: vr512, preferred-register: '' }
+  - { id: 12, class: gr8, preferred-register: '' }
+  - { id: 13, class: gr8, preferred-register: '' }
+  - { id: 14, class: gr8, preferred-register: '' }
+  - { id: 15, class: gr8, preferred-register: '' }
+  - { id: 16, class: gr8, preferred-register: '' }
+  - { id: 17, class: gr8, preferred-register: '' }
+  - { id: 18, class: gr8, preferred-register: '' }
+liveins:
+  - { reg: '$rdi', virtual-reg: '%0' }
+  - { reg: '$zmm0', virtual-reg: '%1' }
+  - { reg: '$zmm1', virtual-reg: '%2' }
+  - { reg: '$zmm2', virtual-reg: '%3' }
+  - { reg: '$zmm3', virtual-reg: '%4' }
+body:             |
+  bb.0:
+    liveins: $rdi, $zmm0, $zmm1, $zmm2, $zmm3
+
+    %0 = COPY $rdi
+    %1 = COPY $zmm0
+    %2 = COPY $zmm1
+    %3 = COPY $zmm2
+    %4 = COPY $zmm3
+
+    %5 = VCMPPDZrri %3, %4, 0, implicit $mxcsr
+    %6 = COPY %5
+    %7 = COPY %6.sub_8bit
+
+    %12 = SHR8ri %7, 2, implicit-def dead $eflags
+    %13 = SHL8ri %12, 1, implicit-def dead $eflags
+    %14 = NOT8r %13
+    %15 = OR8rr %14, %12, implicit-def dead $eflags
+    %16 = AND8rr %15, %13, implicit-def dead $eflags
+    %17 = XOR8rr %16, %12, implicit-def dead $eflags
+    %18 = ADD8rr %17, %14, implicit-def dead $eflags
+
+    %8 = IMPLICIT_DEF
+    %9 = INSERT_SUBREG %8, %18, %subreg.sub_8bit_hi
+    %10 = COPY %9
+    %11 = VMOVAPDZrrk %2, killed %10, %1
+    VMOVAPDZmr %0, 1, $noreg, 0, $noreg, killed %11
+
+  bb.1:
+
+  bb.2:
+    RET 0
+
+...