[llvm] [SelectionDAG] Handle VSCALE in isKnownNeverZero (PR #97789)

[Luke Lau via llvm-commits]

https://github.com/lukel97 created https://github.com/llvm/llvm-project/pull/97789

VSCALE is by definition greater than zero, but this checks it via getVScaleRange anyway.

The motivation for this is to be able to check if the EVL for a VP strided load is non-zero in #97394.

I added the tests to the RISC-V backend since the existing X86 known-never-zero.ll test crashed when trying to lower vscale for the +sse2 RUN line.


>From e1ba4fb98f13d792e9198de72921add5122fa96b Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Fri, 5 Jul 2024 13:47:38 +0800
Subject: [PATCH 1/2] Precommit test

---
 .../CodeGen/RISCV/rvv/known-never-zero.ll     | 43 +++++++++++++++++++
 1 file changed, 43 insertions(+)
 create mode 100644 llvm/test/CodeGen/RISCV/rvv/known-never-zero.ll

diff --git a/llvm/test/CodeGen/RISCV/rvv/known-never-zero.ll b/llvm/test/CodeGen/RISCV/rvv/known-never-zero.ll
new file mode 100644
index 0000000000000..0aeb41c4799bc
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/rvv/known-never-zero.ll
@@ -0,0 +1,43 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -mtriple=riscv64 -mattr=+v -verify-machineinstrs | FileCheck %s
+
+; Use cttz to test if we properly prove never-zero. There is a very
+; simple transform from cttz -> cttz_zero_undef if its operand is
+; known never zero.
+
+; Even without vscale_range, vscale is always guaranteed to be non-zero.
+define i32 @vscale_known_nonzero() {
+; CHECK-LABEL: vscale_known_nonzero:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    addi sp, sp, -16
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    sd ra, 8(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s0, 0(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    .cfi_offset ra, -8
+; CHECK-NEXT:    .cfi_offset s0, -16
+; CHECK-NEXT:    csrr a0, vlenb
+; CHECK-NEXT:    srli s0, a0, 3
+; CHECK-NEXT:    neg a0, s0
+; CHECK-NEXT:    and a0, s0, a0
+; CHECK-NEXT:    lui a1, 30667
+; CHECK-NEXT:    addiw a1, a1, 1329
+; CHECK-NEXT:    call __muldi3
+; CHECK-NEXT:    bnez s0, .LBB0_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    li a0, 32
+; CHECK-NEXT:    j .LBB0_3
+; CHECK-NEXT:  .LBB0_2:
+; CHECK-NEXT:    srliw a0, a0, 27
+; CHECK-NEXT:    lui a1, %hi(.LCPI0_0)
+; CHECK-NEXT:    addi a1, a1, %lo(.LCPI0_0)
+; CHECK-NEXT:    add a0, a1, a0
+; CHECK-NEXT:    lbu a0, 0(a0)
+; CHECK-NEXT:  .LBB0_3:
+; CHECK-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s0, 0(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    addi sp, sp, 16
+; CHECK-NEXT:    ret
+  %x = call i32 @llvm.vscale()
+  %r = call i32 @llvm.cttz.i32(i32 %x, i1 false)
+  ret i32 %r
+}

>From 7f2edf3465b7aa0d2d555c40217fe017f37e6aed Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Fri, 5 Jul 2024 13:51:40 +0800
Subject: [PATCH 2/2] [SelectionDAG] Handle VSCALE in isKnownNeverZero

VSCALE is by definition greater than zero, but this checks it via getVScaleRange anyway.

The motivation for this is to be able to check if the EVL for a VP strided load is non-zero in #97394.

I added the tests to the RISC-V backend since the existing X86 known-never-zero.ll test crashed when trying to lower vscale for the +sse2 RUN line.
---
 llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp  |  9 +++++++++
 llvm/test/CodeGen/RISCV/rvv/known-never-zero.ll | 15 +++------------
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 9d2deb4b5f511..139e86a009739 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -5623,6 +5623,15 @@ bool SelectionDAG::isKnownNeverZero(SDValue Op, unsigned Depth) const {
   case ISD::ZERO_EXTEND:
   case ISD::SIGN_EXTEND:
     return isKnownNeverZero(Op.getOperand(0), Depth + 1);
+  case ISD::VSCALE: {
+    const Function &F = getMachineFunction().getFunction();
+    const APInt &Multiplier = Op.getConstantOperandAPInt(0);
+    ConstantRange CR =
+      getVScaleRange(&F, Op.getScalarValueSizeInBits()).multiply(Multiplier);
+    if (!CR.getUnsignedMin().isZero())
+      return true;
+    break;
+  }
   }
 
   return computeKnownBits(Op, Depth).isNonZero();
diff --git a/llvm/test/CodeGen/RISCV/rvv/known-never-zero.ll b/llvm/test/CodeGen/RISCV/rvv/known-never-zero.ll
index 0aeb41c4799bc..04367028b12d8 100644
--- a/llvm/test/CodeGen/RISCV/rvv/known-never-zero.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/known-never-zero.ll
@@ -12,29 +12,20 @@ define i32 @vscale_known_nonzero() {
 ; CHECK-NEXT:    addi sp, sp, -16
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    sd ra, 8(sp) # 8-byte Folded Spill
-; CHECK-NEXT:    sd s0, 0(sp) # 8-byte Folded Spill
 ; CHECK-NEXT:    .cfi_offset ra, -8
-; CHECK-NEXT:    .cfi_offset s0, -16
 ; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    srli s0, a0, 3
-; CHECK-NEXT:    neg a0, s0
-; CHECK-NEXT:    and a0, s0, a0
+; CHECK-NEXT:    srli a0, a0, 3
+; CHECK-NEXT:    neg a1, a0
+; CHECK-NEXT:    and a0, a0, a1
 ; CHECK-NEXT:    lui a1, 30667
 ; CHECK-NEXT:    addiw a1, a1, 1329
 ; CHECK-NEXT:    call __muldi3
-; CHECK-NEXT:    bnez s0, .LBB0_2
-; CHECK-NEXT:  # %bb.1:
-; CHECK-NEXT:    li a0, 32
-; CHECK-NEXT:    j .LBB0_3
-; CHECK-NEXT:  .LBB0_2:
 ; CHECK-NEXT:    srliw a0, a0, 27
 ; CHECK-NEXT:    lui a1, %hi(.LCPI0_0)
 ; CHECK-NEXT:    addi a1, a1, %lo(.LCPI0_0)
 ; CHECK-NEXT:    add a0, a1, a0
 ; CHECK-NEXT:    lbu a0, 0(a0)
-; CHECK-NEXT:  .LBB0_3:
 ; CHECK-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
-; CHECK-NEXT:    ld s0, 0(sp) # 8-byte Folded Reload
 ; CHECK-NEXT:    addi sp, sp, 16
 ; CHECK-NEXT:    ret
   %x = call i32 @llvm.vscale()