From 7e35494f96344cc71b63f69fad4c0356337cf6ab Mon Sep 17 00:00:00 2001
From: Andrew Adams <andrew.b.adams@gmail.com>
Date: Sat, 14 Oct 2023 04:29:33 -0700
Subject: [PATCH] Generate simpler LLVM IR for shuffles that recursively become
 broadcasts (#7902)

* Generate simpler LLVM IR for shuffles that recursively become broadcasts

* Don't re-codegen arg
---
 src/CodeGen_LLVM.cpp | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/src/CodeGen_LLVM.cpp b/src/CodeGen_LLVM.cpp
index b4fd79d88a15..610dd3719a57 100644
--- a/src/CodeGen_LLVM.cpp
+++ b/src/CodeGen_LLVM.cpp
@@ -4062,6 +4062,25 @@ void CodeGen_LLVM::visit(const Shuffle *op) {
         vecs.push_back(codegen(e));
     }
 
+    // Handle simple broadcasts, which can be generated by the recursive calls below.
+    if (op->vectors.size() == 1) {
+        bool all_indices_the_same = true;
+        internal_assert(!op->indices.empty());
+        for (int i : op->indices) {
+            all_indices_the_same &= (i == op->indices[0]);
+        }
+        if (all_indices_the_same) {
+            value = vecs[0];
+            if (value->getType()->isVectorTy()) {
+                value = builder->CreateExtractElement(value, ConstantInt::get(i32_t, op->indices[0]));
+            } else {
+                internal_assert(op->indices[0] == 0);
+            }
+            value = create_broadcast(value, op->indices.size());
+            return;
+        }
+    }
+
     if (op->is_interleave()) {
         value = interleave_vectors(vecs);
     } else if (op->is_concat()) {