[llvm] [LoopVectorize] Check for vector-to-scalar casts in legalizer (PR #106244)

[via llvm-commits]

https://github.com/ErikHogeman created https://github.com/llvm/llvm-project/pull/106244

The code makes assumptions later on the operations and their inputs being scalar in the loops that are processed, so we should make sure this is the case in the legalizer.

>From 11451c6fc56e8c09c81540660f8afbefafe2146c Mon Sep 17 00:00:00 2001
From: Erik Hogeman <erik.hogeman at arm.com>
Date: Tue, 27 Aug 2024 18:20:23 +0200
Subject: [PATCH] [LoopVectorize] Check for vector-to-scalar casts in legalizer

The code makes assumptions later on the operations and
their inputs being scalar in the loops that are processed,
so we should make sure this is the case in the legalizer.
---
 .../Vectorize/LoopVectorizationLegality.cpp   |  3 ++
 .../LoopVectorize/vector-to-scalar-cast.ll    | 32 +++++++++++++++++++
 2 files changed, 35 insertions(+)
 create mode 100644 llvm/test/Transforms/LoopVectorize/vector-to-scalar-cast.ll

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
index 66a779da8c25bc..b1366b484587d6 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
@@ -943,9 +943,12 @@ bool LoopVectorizationLegality::canVectorizeInstrs() {
         VecCallVariantsFound = true;
 
       // Check that the instruction return type is vectorizable.
+      // We can't vectorize casts from vector type to scalar type.
       // Also, we can't vectorize extractelement instructions.
       if ((!VectorType::isValidElementType(I.getType()) &&
            !I.getType()->isVoidTy()) ||
+          (isa<CastInst>(&I) &&
+           !VectorType::isValidElementType(I.getOperand(0)->getType())) ||
           isa<ExtractElementInst>(I)) {
         reportVectorizationFailure("Found unvectorizable type",
             "instruction return type cannot be vectorized",
diff --git a/llvm/test/Transforms/LoopVectorize/vector-to-scalar-cast.ll b/llvm/test/Transforms/LoopVectorize/vector-to-scalar-cast.ll
new file mode 100644
index 00000000000000..1a34c815f5a18d
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/vector-to-scalar-cast.ll
@@ -0,0 +1,32 @@
+; RUN: opt --force-widen-divrem-via-safe-divisor=false -passes=loop-vectorize -S --debug-only=loop-vectorize < %s 2>&1 | FileCheck %s
+
+; CHECK: LV: Not vectorizing: Found unvectorizable type   %s2 = bitcast <2 x i16> %vec1 to i32
+; CHECK-NOT: Assertion
+
+; Function Attrs: willreturn
+define void @__start() #0 {
+entry:
+  %vec0 = insertelement <2 x i16> undef, i16 0, i64 0
+  %vec1 = insertelement <2 x i16> %vec0, i16 0, i64 1
+  br label %bb0
+
+bb0:
+  %s0 = phi i32 [ %s1, %bb1 ], [ 1, %entry ]
+  br i1 0, label %bb2, label %bb1
+
+bb1:
+  %s1 = add nuw nsw i32 %s0, 1
+  %exitcond = icmp ne i32 %s1, 11
+  br i1 %exitcond, label %bb0, label %bb3
+
+bb2:
+  %s2 = bitcast <2 x i16> %vec1 to i32
+  %s3 = srem i32 0, %s2
+  br label %bb1
+
+bb3:
+  ret void
+}
+
+attributes #0 = { willreturn }
+