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

[Dominik Steenken via llvm-commits]

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

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 introducing a short loop to find the vector operand instead of assuming that it is the first operand.

>From 5bb6fbeb52884de54aad380caaecf64f782c2fb8 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 introducing a short loop to find
the vector operand instead of assuming that it is the first operand.
---
 .../SystemZ/SystemZTargetTransformInfo.cpp    | 39 ++++++++++++++-----
 1 file changed, 29 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Target/SystemZ/SystemZTargetTransformInfo.cpp b/llvm/lib/Target/SystemZ/SystemZTargetTransformInfo.cpp
index 4c9e78c05dbcac..5da42d7cccd3c1 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,43 @@ 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();
+  // Find the type of the vector operand of the intrinsic
+  // This assumes that each vector reduction intrinsic only
+  // has one vector operand.
+  FixedVectorType *VType = 0x0;
+  for (unsigned I = 0; I < II->getNumOperands(); ++I) {
+    auto *T = II->getOperand(I)->getType();
+    if (T->isVectorTy()) {
+      VType = cast<FixedVectorType>(T);
+      break;
+    }
+  }
+
+  // If we did not find a vector operand, do not continue.
+  if (VType == 0x0)
+    return true;
+
+  // If the vector operand is not a full vector, the reduction
+  // should be expanded.
+  unsigned NumElts = VType->getNumElements();
+  unsigned ScalarSize = VType->getScalarSizeInBits();
   unsigned MaxElts = SystemZ::VectorBits / ScalarSize;
   if (NumElts < MaxElts)
     return true;
 
-  // Otherwise
+  // Handling of full vector operands depends on the
+  // individual intrinsic.
   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:
+    // ...unless the scalar size is i64 or larger, since the
+    // performance benefit is dubious there
+    return ScalarSize >= 64;
   }
 }