[llvm] [SandboxVec][BottomUpVec] Implement packing of vectors (PR #116447)

[via llvm-commits]

https://github.com/vporpo created https://github.com/llvm/llvm-project/pull/116447

Up until now we could only support packing of scalar elements. This patch fixes this by implementing packing of vector elements, by generating extractelement and insertelement instruction pairs.

>From 3b4a1cea34596e0ea77740b7596e831ecd41d8de Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Tue, 5 Nov 2024 13:17:16 -0800
Subject: [PATCH] [SandboxVec][BottomUpVec] Implement packing of vectors

Up until now we could only support packing of scalar elements.
This patch fixes this by implementing packing of vector elements,
by generating extractelement and insertelement instruction pairs.
---
 .../SandboxVectorizer/Passes/BottomUpVec.cpp  | 26 ++++++++++++++++---
 .../SandboxVectorizer/bottomup_basic.ll       | 23 ++++++++++++++++
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
index 4b82fea2aa2b22..a2ea11be59b8ed 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
@@ -181,12 +181,32 @@ Value *BottomUpVec::createPack(ArrayRef<Value *> ToPack) {
     // An element can be either scalar or vector. We need to generate different
     // IR for each case.
     if (Elm->getType()->isVectorTy()) {
-      llvm_unreachable("Unimplemented");
+      unsigned NumElms =
+          cast<FixedVectorType>(Elm->getType())->getNumElements();
+      for (auto ExtrLane : seq<int>(0, NumElms)) {
+        // We generate extract-insert pairs, for each lane in `Elm`.
+        Constant *ExtrLaneC =
+            ConstantInt::getSigned(Type::getInt32Ty(Ctx), ExtrLane);
+        // This may return a Constant if Elm is a Constant.
+        auto *ExtrI =
+            ExtractElementInst::create(Elm, ExtrLaneC, WhereIt, Ctx, "VPack");
+        if (!isa<Constant>(ExtrI))
+          WhereIt = std::next(cast<Instruction>(ExtrI)->getIterator());
+        Constant *InsertLaneC =
+            ConstantInt::getSigned(Type::getInt32Ty(Ctx), InsertIdx++);
+        // This may also return a Constant if ExtrI is a Constant.
+        auto *InsertI = InsertElementInst::create(
+            LastInsert, ExtrI, InsertLaneC, WhereIt, Ctx, "VPack");
+        if (!isa<Constant>(InsertI)) {
+          LastInsert = InsertI;
+          WhereIt = std::next(cast<Instruction>(LastInsert)->getIterator());
+        }
+      }
     } else {
       Constant *InsertLaneC =
           ConstantInt::getSigned(Type::getInt32Ty(Ctx), InsertIdx++);
-      // This may be folded into a Constant if LastInsert is a Constant. In that
-      // case we only collect the last constant.
+      // This may be folded into a Constant if LastInsert is a Constant. In
+      // that case we only collect the last constant.
       LastInsert = InsertElementInst::create(LastInsert, Elm, InsertLaneC,
                                              WhereIt, Ctx, "Pack");
       if (auto *NewI = dyn_cast<Instruction>(LastInsert))
diff --git a/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll b/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll
index dff27d06d8ed22..7422d287ff3e2a 100644
--- a/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll
+++ b/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll
@@ -187,3 +187,26 @@ define void @cant_vectorize_seeds(ptr %ptr) {
   ret void
 }
 
+define void @pack_vectors(ptr %ptr, ptr %ptr2) {
+; CHECK-LABEL: define void @pack_vectors(
+; CHECK-SAME: ptr [[PTR:%.*]], ptr [[PTR2:%.*]]) {
+; CHECK-NEXT:    [[PTR0:%.*]] = getelementptr <2 x float>, ptr [[PTR]], i32 0
+; CHECK-NEXT:    [[PTR1:%.*]] = getelementptr float, ptr [[PTR]], i32 2
+; CHECK-NEXT:    [[LD0:%.*]] = load <2 x float>, ptr [[PTR0]], align 8
+; CHECK-NEXT:    [[LD1:%.*]] = load float, ptr [[PTR2]], align 4
+; CHECK-NEXT:    [[VPACK:%.*]] = extractelement <2 x float> [[LD0]], i32 0
+; CHECK-NEXT:    [[VPACK1:%.*]] = insertelement <3 x float> poison, float [[VPACK]], i32 0
+; CHECK-NEXT:    [[VPACK2:%.*]] = extractelement <2 x float> [[LD0]], i32 1
+; CHECK-NEXT:    [[VPACK3:%.*]] = insertelement <3 x float> [[VPACK1]], float [[VPACK2]], i32 1
+; CHECK-NEXT:    [[PACK:%.*]] = insertelement <3 x float> [[VPACK3]], float [[LD1]], i32 2
+; CHECK-NEXT:    store <3 x float> [[PACK]], ptr [[PTR0]], align 8
+; CHECK-NEXT:    ret void
+;
+  %ptr0 = getelementptr <2 x float>, ptr %ptr, i32 0
+  %ptr1 = getelementptr float, ptr %ptr, i32 2
+  %ld0 = load <2 x float>, ptr %ptr0
+  %ld1 = load float, ptr %ptr2
+  store <2 x float> %ld0, ptr %ptr0
+  store float %ld1, ptr %ptr1
+  ret void
+}