[llvm] [SystemZ] Fix Operand Retrieval for Vector Reduction Intrinsic in `shouldExpandReduction` (PR #88874)

[Dominik Steenken via llvm-commits]

https://github.com/dominik-steenken updated https://github.com/llvm/llvm-project/pull/88874

>From da9a7b34d5a746531f98fa705007c605ed66919c Mon Sep 17 00:00:00 2001
From: Dominik Steenken <dost at de.ibm.com>
Date: Mon, 15 Apr 2024 22:03:53 +0200
Subject: [PATCH] [SystemZ] Fix Operand Retrieval for Vector Reduction
 Intrinsic

In the existing version, SystemZTTIImpl::shouldExpandReduction will
create a `cast` error when handling vector reduction intrinsics that
do not have the vector to reduce as their first operand, such as
`llvm.vector.reduce.fadd` and `llvm.vector.reduce.fmul`.
This commit fixes that problem by moving the cast into the case
statement that handles the specific intrinsic, where the vector
operand position is well-known.
---
 .../SystemZ/SystemZTargetTransformInfo.cpp    | 25 ++++++++-----------
 1 file changed, 11 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/Target/SystemZ/SystemZTargetTransformInfo.cpp b/llvm/lib/Target/SystemZ/SystemZTargetTransformInfo.cpp
index 4c9e78c05dbcac..3cd1e05aa5d18c 100644
--- a/llvm/lib/Target/SystemZ/SystemZTargetTransformInfo.cpp
+++ b/llvm/lib/Target/SystemZ/SystemZTargetTransformInfo.cpp
@@ -18,6 +18,7 @@
 #include "llvm/CodeGen/BasicTTIImpl.h"
 #include "llvm/CodeGen/CostTable.h"
 #include "llvm/CodeGen/TargetLowering.h"
+#include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/Support/Debug.h"
@@ -1323,25 +1324,21 @@ SystemZTTIImpl::getIntrinsicInstrCost(const IntrinsicCostAttributes &ICA,
 }
 
 bool SystemZTTIImpl::shouldExpandReduction(const IntrinsicInst *II) const {
-  // Always expand on Subtargets without vector instructions
+  // Always expand on Subtargets without vector instructions.
   if (!ST->hasVector())
     return true;
 
-  // Always expand for operands that do not fill one vector reg
-  auto *Type = cast<FixedVectorType>(II->getOperand(0)->getType());
-  unsigned NumElts = Type->getNumElements();
-  unsigned ScalarSize = Type->getScalarSizeInBits();
-  unsigned MaxElts = SystemZ::VectorBits / ScalarSize;
-  if (NumElts < MaxElts)
-    return true;
-
-  // Otherwise
+  // Whether or not to expand is a per-intrinsic decision.
   switch (II->getIntrinsicID()) {
-  // Do not expand vector.reduce.add
-  case Intrinsic::vector_reduce_add:
-    // Except for i64, since the performance benefit is dubious there
-    return ScalarSize >= 64;
   default:
     return true;
+  // Do not expand vector.reduce.add...
+  case Intrinsic::vector_reduce_add:
+    auto *VType = cast<FixedVectorType>(II->getOperand(0)->getType());
+    // ...unless the scalar size is i64 or larger,
+    // or the operand vector is not full, since the
+    // performance benefit is dubious in those cases.
+    return VType->getScalarSizeInBits() >= 64 ||
+           VType->getPrimitiveSizeInBits() < SystemZ::VectorBits;
   }
 }