[llvm] eabf11f - [CodeGenPrepare] Avoid a scalable-vector crash in ctlz/cttz

[Fraser Cormack via llvm-commits]

Author: Fraser Cormack
Date: 2021-10-20T16:45:55+01:00
New Revision: eabf11f9ea8d647d78e455664658b4faf8880029

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

LOG: [CodeGenPrepare] Avoid a scalable-vector crash in ctlz/cttz

This patch fixes a crash when despeculating ctlz/cttz intrinsics with
scalable-vector types. It is not safe to speculatively get the size of
the vector type in bits in case the vector type is not a fixed-length type. As
it happens this isn't required as vector types are skipped anyway.

Reviewed By: RKSimon

Differential Revision: https://reviews.llvm.org/D112141

Added: 
    llvm/test/Transforms/CodeGenPrepare/RISCV/cttz-ctlz.ll
    llvm/test/Transforms/CodeGenPrepare/RISCV/lit.local.cfg

Modified: 
    llvm/lib/CodeGen/CodeGenPrepare.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index 85def4f17f039..9fdee8f5e4ffd 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -2036,7 +2036,7 @@ static bool despeculateCountZeros(IntrinsicInst *CountZeros,
 
   // Only handle legal scalar cases. Anything else requires too much work.
   Type *Ty = CountZeros->getType();
-  unsigned SizeInBits = Ty->getPrimitiveSizeInBits();
+  unsigned SizeInBits = Ty->getScalarSizeInBits();
   if (Ty->isVectorTy() || SizeInBits > DL->getLargestLegalIntTypeSizeInBits())
     return false;
 

diff  --git a/llvm/test/Transforms/CodeGenPrepare/RISCV/cttz-ctlz.ll b/llvm/test/Transforms/CodeGenPrepare/RISCV/cttz-ctlz.ll
new file mode 100644
index 0000000000000..c70112e91ebd9
--- /dev/null
+++ b/llvm/test/Transforms/CodeGenPrepare/RISCV/cttz-ctlz.ll
@@ -0,0 +1,28 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -codegenprepare < %s | FileCheck %s
+
+target triple = "riscv64-unknown-unknown"
+
+; Check that despeculating count-zeros intrinsics doesn't crash when those
+; intrinsics use scalable types.
+
+define <vscale x 4 x i64> @cttz_nxv4i64(<vscale x 4 x i64> %x) {
+; CHECK-LABEL: @cttz_nxv4i64(
+; CHECK-NEXT:    [[Z:%.*]] = call <vscale x 4 x i64> @llvm.cttz.nxv4i64(<vscale x 4 x i64> [[X:%.*]], i1 false)
+; CHECK-NEXT:    ret <vscale x 4 x i64> [[Z]]
+;
+  %z = call <vscale x 4 x i64> @llvm.cttz.nxv4i64(<vscale x 4 x i64> %x, i1 false)
+  ret <vscale x 4 x i64> %z
+}
+
+define <vscale x 4 x i64> @ctlz_nxv4i64(<vscale x 4 x i64> %x) {
+; CHECK-LABEL: @ctlz_nxv4i64(
+; CHECK-NEXT:    [[Z:%.*]] = call <vscale x 4 x i64> @llvm.ctlz.nxv4i64(<vscale x 4 x i64> [[X:%.*]], i1 false)
+; CHECK-NEXT:    ret <vscale x 4 x i64> [[Z]]
+;
+  %z = call <vscale x 4 x i64> @llvm.ctlz.nxv4i64(<vscale x 4 x i64> %x, i1 false)
+  ret <vscale x 4 x i64> %z
+}
+
+declare <vscale x 4 x i64> @llvm.cttz.nxv4i64(<vscale x 4 x i64>, i1)
+declare <vscale x 4 x i64> @llvm.ctlz.nxv4i64(<vscale x 4 x i64>, i1)

diff  --git a/llvm/test/Transforms/CodeGenPrepare/RISCV/lit.local.cfg b/llvm/test/Transforms/CodeGenPrepare/RISCV/lit.local.cfg
new file mode 100644
index 0000000000000..c63820126f8ca
--- /dev/null
+++ b/llvm/test/Transforms/CodeGenPrepare/RISCV/lit.local.cfg
@@ -0,0 +1,2 @@
+if not 'RISCV' in config.root.targets:
+    config.unsupported = True