From d699e54ca291a7e58d997760481f246db32b506a Mon Sep 17 00:00:00 2001
From: Jon Roelofs <jonathan_roelofs@apple.com>
Date: Wed, 18 May 2022 18:34:49 -0700
Subject: [PATCH] Fix an or+and miscompile w/ GlobalISel

Fixes #55284
---
 .../lib/CodeGen/GlobalISel/CombinerHelper.cpp | 15 ++++--
 .../combine-and-or-disjoint-mask.mir          | 48 +++++++++++++++++++
 2 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
index f974fe9859ad..b1798d681fa3 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
@@ -4224,18 +4224,23 @@ bool CombinerHelper::matchAndOrDisjointMask(
     return false;
 
   Register Src;
-  int64_t MaskAnd;
-  int64_t MaskOr;
+  Register AndMaskReg;
+  int64_t AndMaskBits;
+  int64_t OrMaskBits;
   if (!mi_match(MI, MRI,
-                m_GAnd(m_GOr(m_Reg(Src), m_ICst(MaskOr)), m_ICst(MaskAnd))))
+                m_GAnd(m_GOr(m_Reg(Src), m_ICst(OrMaskBits)),
+                       m_all_of(m_ICst(AndMaskBits), m_Reg(AndMaskReg)))))
     return false;
 
-  // Check if MaskOr could turn on any bits in Src.
-  if (MaskAnd & MaskOr)
+  // Check if OrMask could turn on any bits in Src.
+  if (AndMaskBits & OrMaskBits)
     return false;
 
   MatchInfo = [=, &MI](MachineIRBuilder &B) {
     Observer.changingInstr(MI);
+    // Canonicalize the result to have the constant on the RHS.
+    if (MI.getOperand(1).getReg() == AndMaskReg)
+      MI.getOperand(2).setReg(AndMaskReg);
     MI.getOperand(1).setReg(Src);
     Observer.changedInstr(MI);
   };
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/combine-and-or-disjoint-mask.mir b/llvm/test/CodeGen/AArch64/GlobalISel/combine-and-or-disjoint-mask.mir
index 1fb6f2ec3f1e..71adf3f6f474 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/combine-and-or-disjoint-mask.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/combine-and-or-disjoint-mask.mir
@@ -28,6 +28,29 @@ body:             |
     RET_ReallyLR implicit $w0
 ...
 ---
+name:            disjoint_masks_rev
+tracksRegLiveness: true
+machineFunctionInfo: {}
+body:             |
+  bb.0:
+    liveins: $w0
+    ; CHECK-LABEL: name: disjoint_masks_rev
+    ; CHECK: liveins: $w0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %x:_(s32) = COPY $w0
+    ; CHECK-NEXT: %two:_(s32) = G_CONSTANT i32 2
+    ; CHECK-NEXT: %and:_(s32) = G_AND %x, %two
+    ; CHECK-NEXT: $w0 = COPY %and(s32)
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+    %x:_(s32) = COPY $w0
+    %one:_(s32) = G_CONSTANT i32 1
+    %two:_(s32) = G_CONSTANT i32 2
+    %or:_(s32) = G_OR %x, %one
+    %and:_(s32) = G_AND %two, %or
+    $w0 = COPY %and(s32)
+    RET_ReallyLR implicit $w0
+...
+---
 name:            intersecting_masks
 tracksRegLiveness: true
 machineFunctionInfo: {}
@@ -53,6 +76,31 @@ body:             |
     RET_ReallyLR implicit $w0
 ...
 ---
+name:            intersecting_masks_rev
+tracksRegLiveness: true
+machineFunctionInfo: {}
+body:             |
+  bb.0:
+    liveins: $w0
+    ; CHECK-LABEL: name: intersecting_masks_rev
+    ; CHECK: liveins: $w0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %x:_(s32) = COPY $w0
+    ; CHECK-NEXT: %one:_(s32) = G_CONSTANT i32 3
+    ; CHECK-NEXT: %two:_(s32) = G_CONSTANT i32 2
+    ; CHECK-NEXT: %or:_(s32) = G_OR %x, %one
+    ; CHECK-NEXT: %and:_(s32) = G_AND %two, %or
+    ; CHECK-NEXT: $w0 = COPY %and(s32)
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+    %x:_(s32) = COPY $w0
+    %one:_(s32) = G_CONSTANT i32 3
+    %two:_(s32) = G_CONSTANT i32 2
+    %or:_(s32) = G_OR %x, %one
+    %and:_(s32) = G_AND %two, %or
+    $w0 = COPY %and(s32)
+    RET_ReallyLR implicit $w0
+...
+---
 name:            disjoint_masks_v
 tracksRegLiveness: true
 machineFunctionInfo: {}