[llvm] af6eb1c - [AArch64][GlobalISel] Fix a crash during unsuccessful G_CTPOP <2 x s64> legalization.

[Amara Emerson via llvm-commits]

Author: Amara Emerson
Date: 2021-05-13T17:28:11-07:00
New Revision: af6eb1c710ca17f2f9bce667b9fb593b6541f260

URL: https://github.com/llvm/llvm-project/commit/af6eb1c710ca17f2f9bce667b9fb593b6541f260
DIFF: https://github.com/llvm/llvm-project/commit/af6eb1c710ca17f2f9bce667b9fb593b6541f260.diff

LOG: [AArch64][GlobalISel] Fix a crash during unsuccessful G_CTPOP <2 x s64> legalization.

The legalization rule for scalar-same-as doesn't handle vectors. Until we
implement custom legalization for this, at least fall back properly.

Added: 
    llvm/test/CodeGen/AArch64/GlobalISel/legalize-vector-ctpop.mir

Modified: 
    llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h
    llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h b/llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h
index 2d26e8306b73d..37144be8dc79e 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h
@@ -956,6 +956,23 @@ class LegalizeRuleSet {
         });
   }
 
+  /// Conditionally narrow the scalar or elt to match the size of another.
+  LegalizeRuleSet &maxScalarEltSameAsIf(LegalityPredicate Predicate,
+                                        unsigned TypeIdx,
+                                        unsigned SmallTypeIdx) {
+    typeIdx(TypeIdx);
+    return narrowScalarIf(
+        [=](const LegalityQuery &Query) {
+          return Query.Types[SmallTypeIdx].getScalarSizeInBits() <
+                     Query.Types[TypeIdx].getScalarSizeInBits() &&
+                 Predicate(Query);
+        },
+        [=](const LegalityQuery &Query) {
+          LLT T = Query.Types[SmallTypeIdx];
+          return std::make_pair(TypeIdx, T);
+        });
+  }
+
   /// Add more elements to the vector to reach the next power of two.
   /// No effect if the type is not a vector or the element count is a power of
   /// two.

diff  --git a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
index 234fe533c3a59..edaf1277ad820 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
@@ -726,11 +726,13 @@ AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
   // TODO: Custom legalization for s128
   // TODO: v2s64, v2s32, v4s32, v4s16, v8s16
   // TODO: Use generic lowering when custom lowering is not possible.
+  auto always = [=](const LegalityQuery &Q) { return true; };
   getActionDefinitionsBuilder(G_CTPOP)
       .legalFor({{v8s8, v8s8}, {v16s8, v16s8}})
       .clampScalar(0, s32, s128)
       .widenScalarToNextPow2(0)
-      .scalarSameSizeAs(1, 0)
+      .minScalarEltSameAsIf(always, 1, 0)
+      .maxScalarEltSameAsIf(always, 1, 0)
       .customFor({{s32, s32}, {s64, s64}});
 
   computeTables();

diff  --git a/llvm/test/CodeGen/AArch64/GlobalISel/legalize-vector-ctpop.mir b/llvm/test/CodeGen/AArch64/GlobalISel/legalize-vector-ctpop.mir
new file mode 100644
index 0000000000000..84d3096527c46
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/legalize-vector-ctpop.mir
@@ -0,0 +1,22 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64-unknown-unknown -verify-machineinstrs -run-pass=legalizer %s -o - -global-isel-abort=2 | FileCheck %s
+...
+# This test doesn't currently legalize but should at least not crash.
+---
+name:            v2s64_dont_crash
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $q0
+    ; CHECK-LABEL: name: v2s64_dont_crash
+    ; CHECK: liveins: $q0
+    ; CHECK: %copy:_(<2 x s64>) = COPY $q0
+    ; CHECK: %ctpop:_(<2 x s64>) = G_CTPOP %copy(<2 x s64>)
+    ; CHECK: $q0 = COPY %ctpop(<2 x s64>)
+    ; CHECK: RET_ReallyLR implicit $q0
+    %copy:_(<2 x s64>) = COPY $q0
+    %ctpop:_(<2 x s64>) = G_CTPOP %copy(<2 x s64>)
+    $q0 = COPY %ctpop(<2 x s64>)
+    RET_ReallyLR implicit $q0
+
+...