[llvm] [DAG] Prevent combination on small negative value for vscale (PR #88108)

[via llvm-commits]

https://github.com/vfdff created https://github.com/llvm/llvm-project/pull/88108

Add a TargetLowering hook for it to use to determine when this is a prefer immediate const for VScale.
On AArch64, rdvl can accept a nagative value, while cntd/cntw/cnth can't.

Set the hook to default true to avoid affecting the other backends.
Fix https://github.com/llvm/llvm-project/issues/84620

>From 1b75ddfbcd4bce9bb3e5e05ff37a1d1ea1bb666b Mon Sep 17 00:00:00 2001
From: zhongyunde 00443407 <zhongyunde at huawei.com>
Date: Tue, 9 Apr 2024 05:25:16 -0400
Subject: [PATCH] [DAG] Prevent combination on small negative value for vscale

Add a TargetLowering hook for it to use to determine when this
is a prefer immediate const for VScale.
On AArch64, rdvl can accept a nagative value, while cntd/cntw/cnth can't.

Set the hook to default true to avoid affecting the other backends.
Fix https://github.com/llvm/llvm-project/issues/84620
---
 llvm/include/llvm/CodeGen/TargetLowering.h       |  2 ++
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp    |  3 ++-
 llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp |  2 ++
 llvm/lib/Target/AArch64/AArch64ISelLowering.cpp  | 10 ++++++++++
 llvm/lib/Target/AArch64/AArch64ISelLowering.h    |  2 ++
 llvm/test/CodeGen/AArch64/sve-vl-arith.ll        |  3 +--
 6 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/TargetLowering.h b/llvm/include/llvm/CodeGen/TargetLowering.h
index a4dc097446186a..c1bec56f60e63b 100644
--- a/llvm/include/llvm/CodeGen/TargetLowering.h
+++ b/llvm/include/llvm/CodeGen/TargetLowering.h
@@ -3847,6 +3847,8 @@ class TargetLowering : public TargetLoweringBase {
   /// legal.  It is frequently not legal in PIC relocation models.
   virtual bool isOffsetFoldingLegal(const GlobalAddressSDNode *GA) const;
 
+  virtual bool isPreferVScaleConst(const APInt Imm) const;
+
   /// On x86, return true if the operand with index OpNo is a CALL or JUMP
   /// instruction, which can use either a memory constraint or an address
   /// constraint. -fasm-blocks "__asm call foo" lowers to
diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 2f46b23a97c62c..9010d7256716fc 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -3979,7 +3979,8 @@ SDValue DAGCombiner::visitSUB(SDNode *N) {
   // canonicalize (sub X, (vscale * C)) to (add X, (vscale * -C))
   if (N1.getOpcode() == ISD::VSCALE && N1.hasOneUse()) {
     const APInt &IntVal = N1.getConstantOperandAPInt(0);
-    return DAG.getNode(ISD::ADD, DL, VT, N0, DAG.getVScale(DL, VT, -IntVal));
+    if (TLI.isPreferVScaleConst(-IntVal))
+      return DAG.getNode(ISD::ADD, DL, VT, N0, DAG.getVScale(DL, VT, -IntVal));
   }
 
   // canonicalize (sub X, step_vector(C)) to (add X, step_vector(-C))
diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index 962f0d98e3be90..7e234d77e7d638 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -502,6 +502,8 @@ TargetLowering::isOffsetFoldingLegal(const GlobalAddressSDNode *GA) const {
   return true;
 }
 
+bool TargetLowering::isPreferVScaleConst(const APInt Imm) const { return true; }
+
 //===----------------------------------------------------------------------===//
 //  Optimization Methods
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index f552f91929201c..c86bd6d54b0232 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -10580,6 +10580,16 @@ bool AArch64TargetLowering::isOffsetFoldingLegal(
   return false;
 }
 
+bool AArch64TargetLowering::isPreferVScaleConst(const APInt Imm) const {
+  if (Subtarget->hasFeature(AArch64::FeatureUseScalarIncVL))
+    return true;
+
+  // Multi of 16 can use instruction rdvl.
+  if (!Imm.isNegatedPowerOf2() && Imm.getLoBits(3).isZero())
+    return false;
+  return true;
+}
+
 bool AArch64TargetLowering::isFPImmLegal(const APFloat &Imm, EVT VT,
                                          bool OptForSize) const {
   bool IsLegal = false;
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.h b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
index 3465f3be887543..2f8e9b5839a0e2 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.h
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
@@ -606,6 +606,8 @@ class AArch64TargetLowering : public TargetLowering {
 
   bool isOffsetFoldingLegal(const GlobalAddressSDNode *GA) const override;
 
+  bool isPreferVScaleConst(const APInt Imm) const override;
+
   bool isFPImmLegal(const APFloat &Imm, EVT VT,
                     bool ForCodeSize) const override;
 
diff --git a/llvm/test/CodeGen/AArch64/sve-vl-arith.ll b/llvm/test/CodeGen/AArch64/sve-vl-arith.ll
index dd4294c8d3bdcc..33869ade84aede 100644
--- a/llvm/test/CodeGen/AArch64/sve-vl-arith.ll
+++ b/llvm/test/CodeGen/AArch64/sve-vl-arith.ll
@@ -204,8 +204,7 @@ define i64 @dech_scalar_i64(i64 %a) {
 ; NO_SCALAR_INC-LABEL: dech_scalar_i64:
 ; NO_SCALAR_INC:       // %bb.0:
 ; NO_SCALAR_INC-NEXT:    cnth x8, all, mul #3
-; NO_SCALAR_INC-NEXT:    neg x8, x8
-; NO_SCALAR_INC-NEXT:    add x0, x0, x8
+; NO_SCALAR_INC-NEXT:    sub x0, x0, x8
 ; NO_SCALAR_INC-NEXT:    ret
 ;
 ; CHECK-LABEL: dech_scalar_i64: