[AArch64][GlobalISel] Check unmergeSrc is a vector in matchCombineBuildUnmerge

[HolyMolyCowMan]

This aims to fix the crash in #168495, my combine rule was missing a check that the source vector was in fact a vector. This then caused the legality check to fail in this example as the concat was trying to concat a non vector.

I have also gated the bitcast of the concat to only work on non-scalable vectors as the mutation calls getNumElements which crashes when called on a scalable vector.

Fixes #168495

[llvmbot]

@llvm/pr-subscribers-llvm-globalisel

Author: Ryan Cowan (HolyMolyCowMan)

Changes

This aims to fix the crash in #168495, my combine rule was missing a check that the source vector was in fact a vector. This then caused the legality check to fail in this example as the concat was trying to concat a non vector.

I have also gated the lowering of the concat to only work on non-scalable vectors as the mutation calls getNumElements which crashed when called on a scalable vector.

---

Full diff: https://github.com/llvm/llvm-project/pull/168692.diff

2 Files Affected:

• (modified) llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp (+3)
• (modified) llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp (+3-1)

diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
index 45a08347b1ec2..f0fbe0135353f 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
@@ -3499,6 +3499,9 @@ bool CombinerHelper::matchCombineBuildUnmerge(MachineInstr &MI,
   LLT DstTy = MRI.getType(MI.getOperand(0).getReg());
   LLT UnmergeSrcTy = MRI.getType(UnmergeSrc);
 
+  if (!UnmergeSrcTy.isVector())
+    return false;
+
   // Ensure we only generate legal instructions post-legalizer
   if (!IsPreLegalize &&
       !isLegal({TargetOpcode::G_CONCAT_VECTORS, {DstTy, UnmergeSrcTy}}))
diff --git a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
index fdf69b04bf676..2ffdae7d51704 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
@@ -1243,7 +1243,9 @@ AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
       .bitcastIf(
           [=](const LegalityQuery &Query) {
             return Query.Types[0].getSizeInBits() <= 128 &&
-                   Query.Types[1].getSizeInBits() <= 64;
+                   Query.Types[1].getSizeInBits() <= 64 &&
+                   !Query.Types[0].isScalableVector() &&
+                   !Query.Types[1].isScalableVector();
           },
           [=](const LegalityQuery &Query) {
             const LLT DstTy = Query.Types[0];

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Ryan Cowan (HolyMolyCowMan)

Changes

This aims to fix the crash in #168495, my combine rule was missing a check that the source vector was in fact a vector. This then caused the legality check to fail in this example as the concat was trying to concat a non vector.

I have also gated the lowering of the concat to only work on non-scalable vectors as the mutation calls getNumElements which crashed when called on a scalable vector.

---

Full diff: https://github.com/llvm/llvm-project/pull/168692.diff

2 Files Affected:

• (modified) llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp (+3)
• (modified) llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp (+3-1)

diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
index 45a08347b1ec2..f0fbe0135353f 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
@@ -3499,6 +3499,9 @@ bool CombinerHelper::matchCombineBuildUnmerge(MachineInstr &MI,
   LLT DstTy = MRI.getType(MI.getOperand(0).getReg());
   LLT UnmergeSrcTy = MRI.getType(UnmergeSrc);
 
+  if (!UnmergeSrcTy.isVector())
+    return false;
+
   // Ensure we only generate legal instructions post-legalizer
   if (!IsPreLegalize &&
       !isLegal({TargetOpcode::G_CONCAT_VECTORS, {DstTy, UnmergeSrcTy}}))
diff --git a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
index fdf69b04bf676..2ffdae7d51704 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
@@ -1243,7 +1243,9 @@ AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
       .bitcastIf(
           [=](const LegalityQuery &Query) {
             return Query.Types[0].getSizeInBits() <= 128 &&
-                   Query.Types[1].getSizeInBits() <= 64;
+                   Query.Types[1].getSizeInBits() <= 64 &&
+                   !Query.Types[0].isScalableVector() &&
+                   !Query.Types[1].isScalableVector();
           },
           [=](const LegalityQuery &Query) {
             const LLT DstTy = Query.Types[0];

[cofibrant]

Could you add the reproducer from #168495 as a test case?

[HolyMolyCowMan]

I could add the MIR & only run the aarch64-postlegalizer-combiner pass to show the crash no longer happens, but there are other issues with selection for that reproducer.

[cofibrant]

I could add the MIR & only run the aarch64-postlegalizer-combiner pass to show the crash no longer happens, but there are other issues with selection for that reproducer.

Sounds good!

[github-actions]

🐧 Linux x64 Test Results

• 186368 tests passed
• 4859 tests skipped

[davemgreen]

An llc test that falls-back to SDAG would be fine as well, so long as it doesn't crash. The mir test has advantages that it can be specific to this issue though, they are good at testing the edge cases.

[HolyMolyCowMan]

I've added the MIR for the reproducer to a test, as well as some basic tests of functionality for completeness' sake.

[cofibrant]

Assuming CI passes, LGTM. Let's see what @davemgreen says before merging, though.

[davemgreen]

Thanks. LGTM