[llvm] [GISel] Lower scalar G_SELECT in LegalizerHelper (PR #79342)

[Kai Nacke via llvm-commits]

https://github.com/redstar updated https://github.com/llvm/llvm-project/pull/79342

>From c89e1c4ab83f4b2d2b14924dccafd1c6cb5c0621 Mon Sep 17 00:00:00 2001
From: Kai Nacke <kai.peter.nacke at ibm.com>
Date: Wed, 24 Jan 2024 18:29:40 +0000
Subject: [PATCH] [GISel] Lower scalar G_SELECT in LegalizerHelper

The LegalizerHelper only has support to lower G_SELECT with
vector operands. The approach is the same for scalar arguments,
which this PR adds.
---
 .../CodeGen/GlobalISel/LegalizerHelper.cpp    | 26 ++++++++-----
 .../GlobalISel/LegalizerHelperTest.cpp        | 39 +++++++++++++++++++
 2 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index 3b2cf3191092734..089c3b6945a02a7 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -7949,10 +7949,9 @@ LegalizerHelper::LegalizeResult LegalizerHelper::lowerSelect(MachineInstr &MI) {
   // Implement vector G_SELECT in terms of XOR, AND, OR.
   auto [DstReg, DstTy, MaskReg, MaskTy, Op1Reg, Op1Ty, Op2Reg, Op2Ty] =
       MI.getFirst4RegLLTs();
-  if (!DstTy.isVector())
-    return UnableToLegalize;
 
-  bool IsEltPtr = DstTy.getElementType().isPointer();
+  bool IsEltPtr =
+      (DstTy.isVector() ? DstTy.getElementType() : DstTy).isPointer();
   if (IsEltPtr) {
     LLT ScalarPtrTy = LLT::scalar(DstTy.getScalarSizeInBits());
     LLT NewTy = DstTy.changeElementType(ScalarPtrTy);
@@ -7962,7 +7961,7 @@ LegalizerHelper::LegalizeResult LegalizerHelper::lowerSelect(MachineInstr &MI) {
   }
 
   if (MaskTy.isScalar()) {
-    // Turn the scalar condition into a vector condition mask.
+    // Turn the scalar condition into a vector condition mask if needed.
 
     Register MaskElt = MaskReg;
 
@@ -7972,13 +7971,22 @@ LegalizerHelper::LegalizeResult LegalizerHelper::lowerSelect(MachineInstr &MI) {
       MaskElt = MIRBuilder.buildSExtInReg(MaskTy, MaskElt, 1).getReg(0);
 
     // Continue the sign extension (or truncate) to match the data type.
-    MaskElt = MIRBuilder.buildSExtOrTrunc(DstTy.getElementType(),
-                                          MaskElt).getReg(0);
+    MaskElt =
+        MIRBuilder
+            .buildSExtOrTrunc(DstTy.isVector() ? DstTy.getElementType() : DstTy,
+                              MaskElt)
+            .getReg(0);
 
-    // Generate a vector splat idiom.
-    auto ShufSplat = MIRBuilder.buildShuffleSplat(DstTy, MaskElt);
-    MaskReg = ShufSplat.getReg(0);
+    if (DstTy.isVector()) {
+      // Generate a vector splat idiom.
+      auto ShufSplat = MIRBuilder.buildShuffleSplat(DstTy, MaskElt);
+      MaskReg = ShufSplat.getReg(0);
+    } else
+      MaskReg = MaskElt;
     MaskTy = DstTy;
+  } else if (!DstTy.isVector()) {
+    // Cannot handle the case that mask is a vector and dst is a scalar.
+    return UnableToLegalize;
   }
 
   if (MaskTy.getSizeInBits() != DstTy.getSizeInBits()) {
diff --git a/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp b/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp
index d7876b7ce87490c..886780d92e5e473 100644
--- a/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp
+++ b/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp
@@ -3431,6 +3431,45 @@ TEST_F(AArch64GISelMITest, LowerUDIVREM) {
   EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
 }
 
+// Test G_SELECT lowering.
+TEST_F(AArch64GISelMITest, LowerSelect) {
+  setUp();
+  if (!TM)
+    GTEST_SKIP();
+
+  // Declare your legalization info
+  DefineLegalizerInfo(A, { getActionDefinitionsBuilder(G_SELECT).lower(); });
+
+  LLT S1 = LLT::scalar(1);
+  LLT S32 = LLT::scalar(32);
+  auto Tst = B.buildTrunc(S1, Copies[0]);
+  auto SrcA = B.buildTrunc(S32, Copies[1]);
+  auto SrcB = B.buildTrunc(S32, Copies[2]);
+  auto SELECT = B.buildInstr(TargetOpcode::G_SELECT, {S32}, {Tst, SrcA, SrcB});
+
+  AInfo Info(MF->getSubtarget());
+  DummyGISelObserver Observer;
+  LegalizerHelper Helper(*MF, Info, Observer, B);
+  // Perform Legalization
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
+            Helper.lower(*SELECT, 0, S32));
+
+  auto CheckStr = R"(
+  CHECK: [[TST:%[0-9]+]]:_(s1) = G_TRUNC
+  CHECK: [[TRUE:%[0-9]+]]:_(s32) = G_TRUNC
+  CHECK: [[FALSE:%[0-9]+]]:_(s32) = G_TRUNC
+  CHECK: [[MSK:%[0-9]+]]:_(s32) = G_SEXT [[TST]]
+  CHECK: [[M:%[0-9]+]]:_(s32) = G_CONSTANT i32 -1
+  CHECK: [[NEGMSK:%[0-9]+]]:_(s32) = G_XOR [[MSK]]:_, [[M]]:_
+  CHECK: [[TVAL:%[0-9]+]]:_(s32) = G_AND [[TRUE]]:_, [[MSK]]:_
+  CHECK: [[FVAL:%[0-9]+]]:_(s32) = G_AND [[FALSE]]:_, [[NEGMSK]]:_
+  CHECK: [[RES:%[0-9]+]]:_(s32) = G_OR [[TVAL]]:_, [[FVAL]]:_
+  )";
+
+  // Check
+  EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
+}
+
 // Test widening of G_UNMERGE_VALUES
 TEST_F(AArch64GISelMITest, WidenUnmerge) {
   setUp();