[PATCH] D123560: GlobalISel: Implement MoreElements for select of vector conditions

[Matt Arsenault via Phabricator via llvm-commits]

arsenm created this revision.
arsenm added reviewers: aemerson, paquette.
Herald added subscribers: hiraditya, rovka.
Herald added a project: All.
arsenm requested review of this revision.
Herald added a subscriber: wdng.
Herald added a project: LLVM.

https://reviews.llvm.org/D123560

Files:
  llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
  llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp


Index: llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp
===================================================================
--- llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp
+++ llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp
@@ -4092,4 +4092,51 @@
   EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
 }
 
+TEST_F(AArch64GISelMITest, MoreElementsSelect) {
+  setUp();
+  if (!TM)
+    return;
+
+  LLT s1 = LLT::scalar(1);
+  LLT s64 = LLT::scalar(64);
+  LLT v2s1 = LLT::fixed_vector(2, 1);
+  LLT v2s32 = LLT::fixed_vector(2, 32);
+
+  LegalizerInfo LI;
+  DummyGISelObserver Observer;
+  LegalizerHelper Helper(*MF, LI, Observer, B);
+
+  B.setInsertPt(*EntryMBB, EntryMBB->end());
+
+  auto Val0 = B.buildBitcast(v2s32, Copies[0]);
+  auto Val1 = B.buildBitcast(v2s32, Copies[1]);
+
+  // Build select of vectors with scalar condition.
+  auto Zero = B.buildConstant(s64, 0);
+  auto Cond = B.buildICmp(CmpInst::ICMP_EQ, s1, Copies[2], Zero);
+  auto Select = B.buildSelect(v2s32, Cond, Val0, Val1);
+
+  // Splat the condition into a vector select
+  B.setInstr(*Select);
+
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::UnableToLegalize,
+            Helper.moreElementsVector(*Select, 1, LLT::fixed_vector(3, 1)));
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
+            Helper.moreElementsVector(*Select, 1, v2s1));
+
+  auto CheckStr = R"(
+  CHECK: [[BITCAST0:%[0-9]+]]:_(<2 x s32>) = G_BITCAST
+  CHECK: [[BITCAST1:%[0-9]+]]:_(<2 x s32>) = G_BITCAST
+  CHECK: [[ZERO0:%[0-9]+]]:_(s64) = G_CONSTANT i64 0
+  CHECK: [[CMP:%[0-9]+]]:_(s1) = G_ICMP intpred(eq), %{{[0-9]+}}:_(s64), [[ZERO0]]
+  CHECK: [[IMPDEF:%[0-9]+]]:_(<2 x s1>) = G_IMPLICIT_DEF
+  CHECK: [[ZERO1:%[0-9]+]]:_(s64) = G_CONSTANT i64 0
+  CHECK: [[INSERT:%[0-9]+]]:_(<2 x s1>) = G_INSERT_VECTOR_ELT [[IMPDEF]]:_, [[CMP]]:_(s1), [[ZERO1]]
+  CHECK: [[SHUFFLE:%[0-9]+]]:_(<2 x s1>) = G_SHUFFLE_VECTOR [[INSERT]]:_(<2 x s1>), [[IMPDEF]]:_, shufflemask(0, 0)
+  CHECK: [[SELECT:%[0-9]+]]:_(<2 x s32>) = G_SELECT [[SHUFFLE]]:_(<2 x s1>), [[BITCAST0]]:_, [[BITCAST1]]:_
+  )";
+
+  EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
+}
+
 } // namespace
Index: llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
===================================================================
--- llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -4897,10 +4897,26 @@
     moreElementsVectorDst(MI, MoreTy, 0);
     Observer.changedInstr(MI);
     return Legalized;
-  case TargetOpcode::G_SELECT:
-    if (TypeIdx != 0)
-      return UnableToLegalize;
-    if (MRI.getType(MI.getOperand(1).getReg()).isVector())
+  case TargetOpcode::G_SELECT: {
+    Register DstReg = MI.getOperand(0).getReg();
+    Register CondReg = MI.getOperand(1).getReg();
+    LLT DstTy = MRI.getType(DstReg);
+    LLT CondTy = MRI.getType(CondReg);
+    if (TypeIdx == 1) {
+      if (!CondTy.isScalar() ||
+          DstTy.getElementCount() != MoreTy.getElementCount())
+        return UnableToLegalize;
+
+      // This is turning a scalar select into vectors into a vector
+      // select. Broadcast the select condition.
+      auto ShufSplat = MIRBuilder.buildShuffleSplat(MoreTy, CondReg);
+      Observer.changingInstr(MI);
+      MI.getOperand(1).setReg(ShufSplat.getReg(0));
+      Observer.changedInstr(MI);
+      return Legalized;
+    }
+
+    if (CondTy.isVector())
       return UnableToLegalize;
 
     Observer.changingInstr(MI);
@@ -4909,6 +4925,7 @@
     moreElementsVectorDst(MI, MoreTy, 0);
     Observer.changedInstr(MI);
     return Legalized;
+  }
   case TargetOpcode::G_UNMERGE_VALUES:
     return UnableToLegalize;
   case TargetOpcode::G_PHI:


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D123560.422095.patch
Type: text/x-patch
Size: 3713 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20220412/5a1978f1/attachment.bin>