From 885a87033c15cb91c91e70f77cb4b5764d4f3399 Mon Sep 17 00:00:00 2001
From: Amara Emerson <amara@apple.com>
Date: Thu, 22 Sep 2022 15:41:27 +0100
Subject: [PATCH] [GlobalISel] Enforce G_ASSERT_ALIGN to have a valid alignment
 > 0.

---
 llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp      |  7 ++-----
 llvm/lib/CodeGen/MachineVerifier.cpp                |  5 +++++
 llvm/test/MachineVerifier/test_g_assert_align.mir   | 17 +++++++++++++++++
 llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp | 19 +++++--------------
 4 files changed, 29 insertions(+), 19 deletions(-)
 create mode 100644 llvm/test/MachineVerifier/test_g_assert_align.mir

diff --git a/llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp b/llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp
index 02a12a0..2281f67 100644
--- a/llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp
@@ -39,7 +39,7 @@ Align GISelKnownBits::computeKnownAlignment(Register R, unsigned Depth) {
     return computeKnownAlignment(MI->getOperand(1).getReg(), Depth);
   case TargetOpcode::G_ASSERT_ALIGN: {
     // TODO: Min with source
-    return Align(std::max(int64_t(1), MI->getOperand(2).getImm()));
+    return Align(MI->getOperand(2).getImm());
   }
   case TargetOpcode::G_FRAME_INDEX: {
     int FrameIdx = MI->getOperand(1).getIndex();
@@ -471,10 +471,7 @@ void GISelKnownBits::computeKnownBitsImpl(Register R, KnownBits &Known,
     break;
   }
   case TargetOpcode::G_ASSERT_ALIGN: {
-    int64_t Align = MI.getOperand(2).getImm();
-    if (Align == 0)
-      break;
-    int64_t LogOfAlign = Log2_64(Align);
+    int64_t LogOfAlign = Log2_64(MI.getOperand(2).getImm());
 
     // TODO: Should use maximum with source
     // If a node is guaranteed to be aligned, set low zero bits accordingly as
diff --git a/llvm/lib/CodeGen/MachineVerifier.cpp b/llvm/lib/CodeGen/MachineVerifier.cpp
index 3cb2049..011d60d 100644
--- a/llvm/lib/CodeGen/MachineVerifier.cpp
+++ b/llvm/lib/CodeGen/MachineVerifier.cpp
@@ -1725,6 +1725,11 @@ void MachineVerifier::verifyPreISelGenericInstruction(const MachineInstr *MI) {
     }
     break;
   }
+  case TargetOpcode::G_ASSERT_ALIGN: {
+    if (MI->getOperand(2).getImm() < 1)
+      report("alignment immediate must be >= 1", MI);
+    break;
+  }
   default:
     break;
   }
diff --git a/llvm/test/MachineVerifier/test_g_assert_align.mir b/llvm/test/MachineVerifier/test_g_assert_align.mir
new file mode 100644
index 0000000..6404050
--- /dev/null
+++ b/llvm/test/MachineVerifier/test_g_assert_align.mir
@@ -0,0 +1,17 @@
+#RUN: not --crash llc -march=aarch64 -o - -global-isel -run-pass=none -verify-machineinstrs %s 2>&1 | FileCheck %s
+# REQUIRES: aarch64-registered-target
+
+---
+name:            test_assert_align
+legalized:       true
+regBankSelected: false
+selected:        false
+tracksRegLiveness: true
+liveins:
+body:             |
+  bb.0:
+    liveins: $x0, $q0
+    %ptr:_(p0) = COPY $x0
+
+    ; CHECK: Bad machine code: alignment immediate must be >= 1
+    %v:_(s32) = G_ASSERT_ALIGN %ptr:_(p0), 0
diff --git a/llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp b/llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp
index a679c46..6fd299a 100644
--- a/llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp
+++ b/llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp
@@ -1923,9 +1923,6 @@ TEST_F(AMDGPUGISelMITest, TestKnownBitsAssertAlign) {
    %val:_(s64) = COPY $vgpr0_vgpr1
    %ptrval:_(p1) = COPY $vgpr0_vgpr1
 
-   %assert_align0:_(s64) = G_ASSERT_ALIGN %val, 0
-   %copy_assert_align0:_(s64) = COPY %assert_align0
-
    %assert_align1:_(s64) = G_ASSERT_ALIGN %val, 1
    %copy_assert_align1:_(s64) = COPY %assert_align1
 
@@ -1962,17 +1959,11 @@ TEST_F(AMDGPUGISelMITest, TestKnownBitsAssertAlign) {
   };
 
   const unsigned NumSetupCopies = 5;
-  // Check zero align specially.
-  Res = GetKB(NumSetupCopies);
-  EXPECT_EQ(0u, Res.Zero.countTrailingOnes());
-  EXPECT_EQ(64u, Res.One.countTrailingZeros());
-  EXPECT_EQ(Align(1), Info.computeKnownAlignment(Copies[5]));
-
-  CheckBits(1, NumSetupCopies + 1);
-  CheckBits(2, NumSetupCopies + 2);
-  CheckBits(3, NumSetupCopies + 3);
-  CheckBits(4, NumSetupCopies + 4);
-  CheckBits(5, NumSetupCopies + 5);
+  CheckBits(1, NumSetupCopies);
+  CheckBits(2, NumSetupCopies + 1);
+  CheckBits(3, NumSetupCopies + 2);
+  CheckBits(4, NumSetupCopies + 3);
+  CheckBits(5, NumSetupCopies + 4);
 }
 
 TEST_F(AArch64GISelMITest, TestKnownBitsUADDO) {
-- 
2.7.4