Fix Vulkan interleave SPIRV codegen. Fix a bug in Simplify_Shuffle. Fix a bug in Deinterleave.

mcourteaux · mcourteaux · commit cf6312e62709 · 2025-05-24T15:05:23.000+02:00
diff --git a/src/CodeGen_Vulkan_Dev.cpp b/src/CodeGen_Vulkan_Dev.cpp
@@ -2041,31 +2041,21 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::visit(const Shuffle *op) {
         debug(3) << "\n";
 
         if (arg_ids.size() == 1) {
-
             // 1 argument, just do a simple assignment via a cast
             SpvId result_id = cast_type(op->type, op->vectors[0].type(), arg_ids[0]);
             builder.update_id(result_id);
 
         } else if (arg_ids.size() == 2) {
-
-            // 2 arguments, use a composite insert to update even and odd indices
-            uint32_t even_idx = 0;
-            uint32_t odd_idx = 1;
-            SpvFactory::Indices even_indices;
-            SpvFactory::Indices odd_indices;
-            for (int i = 0; i < op_lanes; ++i) {
-                even_indices.push_back(even_idx);
-                odd_indices.push_back(odd_idx);
-                even_idx += 2;
-                odd_idx += 2;
+            // 2 arguments, use vector-shuffle with logical indices indexing into (vec1[0], vec1[1], ..., vec2[0], vec2[1], ...)
+            SpvFactory::Indices logical_indices;
+            for (int i = 0; i < arg_lanes; ++i) {
+                logical_indices.push_back(uint32_t(i));
+                logical_indices.push_back(uint32_t(i + arg_lanes));
             }
 
             SpvId type_id = builder.declare_type(op->type);
-            SpvId value_id = builder.declare_null_constant(op->type);
-            SpvId partial_id = builder.reserve_id(SpvResultId);
             SpvId result_id = builder.reserve_id(SpvResultId);
-            builder.append(SpvFactory::composite_insert(type_id, partial_id, arg_ids[0], value_id, even_indices));
-            builder.append(SpvFactory::composite_insert(type_id, result_id, arg_ids[1], partial_id, odd_indices));
+            builder.append(SpvFactory::vector_shuffle(type_id, result_id, arg_ids[0], arg_ids[1], logical_indices));
             builder.update_id(result_id);
 
         } else {
@@ -2095,7 +2085,7 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::visit(const Shuffle *op) {
     } else if (op->is_extract_element()) {
         int idx = op->indices[0];
         internal_assert(idx >= 0);
-        internal_assert(idx <= op->vectors[0].type().lanes());
+        internal_assert(idx < op->vectors[0].type().lanes());
         if (op->vectors[0].type().is_vector()) {
             SpvFactory::Indices indices = {(uint32_t)idx};
             SpvId type_id = builder.declare_type(op->type);
diff --git a/src/Deinterleave.cpp b/src/Deinterleave.cpp
@@ -298,6 +298,10 @@ class Deinterleaver : public IRGraphMutator {
         } else {
 
             Type t = op->type.with_lanes(new_lanes);
+            internal_assert((op->type.lanes() - starting_lane + lane_stride - 1) / lane_stride == new_lanes)
+                << "Deinterleaving with lane stride " << lane_stride << " and staring lane " << starting_lane
+                << " for var of Type " << op->type << " to " << t << " drops lanes unexpectedly."
+                << " Deinterleaver probably recursed too deep into types of different lane count.";
             if (external_lets.contains(op->name) &&
                 starting_lane == 0 &&
                 lane_stride == 2) {
@@ -392,8 +396,12 @@ class Deinterleaver : public IRGraphMutator {
             int index = indices.front();
             for (const auto &i : op->vectors) {
                 if (index < i.type().lanes()) {
-                    ScopedValue<int> lane(starting_lane, index);
-                    return mutate(i);
+                    if (i.type().lanes() == op->type.lanes()) {
+                        ScopedValue<int> scoped_starting_lane(starting_lane, index);
+                        return mutate(i);
+                    } else {
+                        return Shuffle::make(op->vectors, indices);
+                    }
                 }
                 index -= i.type().lanes();
             }
@@ -405,10 +413,18 @@ class Deinterleaver : public IRGraphMutator {
 };
 
 Expr deinterleave(Expr e, int starting_lane, int lane_stride, int new_lanes, const Scope<> &lets) {
+    debug(3) << "Deinterleave "
+             << "(start:" << starting_lane << ", stide:" << lane_stride << ", new_lanes:" << new_lanes << "): "
+             << e << " of Type: " << e.type() << "\n";
+    Type original_type = e.type();
     e = substitute_in_all_lets(e);
     Deinterleaver d(starting_lane, lane_stride, new_lanes, lets);
     e = d.mutate(e);
     e = common_subexpression_elimination(e);
+    Type final_type = e.type();
+    int expected_lanes = (original_type.lanes() + lane_stride - starting_lane - 1) / lane_stride;
+    internal_assert(original_type.code() == final_type.code()) << "Underlying types not identical after interleaving.";
+    internal_assert(expected_lanes == final_type.lanes()) << "Number of lanes incorrect after interleaving: " << final_type.lanes() << "while expected was " << expected_lanes << ".";
     return simplify(e);
 }
 
@@ -419,12 +435,12 @@ Expr extract_odd_lanes(const Expr &e, const Scope<> &lets) {
 
 Expr extract_even_lanes(const Expr &e, const Scope<> &lets) {
     internal_assert(e.type().lanes() % 2 == 0);
-    return deinterleave(e, 0, 2, (e.type().lanes() + 1) / 2, lets);
+    return deinterleave(e, 0, 2, e.type().lanes() / 2, lets);
 }
 
 Expr extract_mod3_lanes(const Expr &e, int lane, const Scope<> &lets) {
     internal_assert(e.type().lanes() % 3 == 0);
-    return deinterleave(e, lane, 3, (e.type().lanes() + 2) / 3, lets);
+    return deinterleave(e, lane, 3, e.type().lanes() / 3, lets);
 }
 
 }  // namespace
diff --git a/src/Simplify_Let.cpp b/src/Simplify_Let.cpp
@@ -98,7 +98,7 @@ Body Simplify::simplify_let(const LetOrLetStmt *op, ExprInfo *info) {
         Expr new_var = Variable::make(f.new_value.type(), f.new_name);
         Expr replacement = new_var;
 
-        debug(4) << "simplify let " << op->name << " = " << f.value << " in...\n";
+        debug(4) << "simplify let " << op->name << " = (" << f.value.type() << ") " << f.value << " in...\n";
 
         while (true) {
             const Variable *var = f.new_value.template as<Variable>();
@@ -180,6 +180,16 @@ Body Simplify::simplify_let(const LetOrLetStmt *op, ExprInfo *info) {
                 f.new_value = cast->value;
                 new_var = Variable::make(f.new_value.type(), f.new_name);
                 replacement = substitute(f.new_name, Cast::make(cast->type, new_var), replacement);
+            } else if (shuffle && shuffle->is_concat() && is_pure(shuffle)) {
+                // Substitute in all concatenates as they will likely simplify
+                // with other shuffles.
+                // As the structure of this while loop makes it hard to peel off
+                // pure operations from _all_ arguments to the Shuffle, we will
+                // instead subsitute all of the vars that go in the shuffle, and
+                // instead guard against side effects by checking with `is_pure()`.
+                replacement = substitute(f.new_name, shuffle, replacement);
+                f.new_value = Expr();
+                break;
             } else if (shuffle && shuffle->is_slice()) {
                 // Replacing new_value below might free the shuffle
                 // indices vector, so save them now.
diff --git a/src/Simplify_Shuffle.cpp b/src/Simplify_Shuffle.cpp
@@ -289,13 +289,18 @@ Expr Simplify::visit(const Shuffle *op, ExprInfo *info) {
             if (inner_shuffle->is_concat()) {
                 int slice_min = op->indices.front();
                 int slice_max = op->indices.back();
+                if (slice_min > slice_max) {
+                    // Slices can go backward.
+                    std::swap(slice_min, slice_max);
+                }
                 int concat_index = 0;
                 int new_slice_start = -1;
                 vector<Expr> new_concat_vectors;
                 for (const auto &v : inner_shuffle->vectors) {
                     // Check if current concat vector overlaps with slice.
-                    if ((concat_index >= slice_min && concat_index <= slice_max) ||
-                        ((concat_index + v.type().lanes() - 1) >= slice_min && (concat_index + v.type().lanes() - 1) <= slice_max)) {
+                    int overlap_max = std::min(slice_max, concat_index + v.type().lanes() - 1);
+                    int overlap_min = std::max(slice_min, concat_index);
+                    if (overlap_min <= overlap_max) {
                         if (new_slice_start < 0) {
                             new_slice_start = concat_index;
                         }
@@ -305,7 +310,10 @@ Expr Simplify::visit(const Shuffle *op, ExprInfo *info) {
                     concat_index += v.type().lanes();
                 }
                 if (new_concat_vectors.size() < inner_shuffle->vectors.size()) {
-                    return Shuffle::make_slice(Shuffle::make_concat(new_concat_vectors), op->slice_begin() - new_slice_start, op->slice_stride(), op->indices.size());
+                    return Shuffle::make_slice(Shuffle::make_concat(new_concat_vectors),
+                                               op->slice_begin() - new_slice_start,
+                                               op->slice_stride(),
+                                               op->indices.size());
                 }
             }
         }
diff --git a/src/runtime/opencl.cpp b/src/runtime/opencl.cpp
@@ -37,7 +37,7 @@ extern "C" WEAK void *halide_opencl_get_symbol(void *user_context, const char *n
 #ifdef WINDOWS
         "opencl.dll",
 #else
-        "libOpenCL.so",
+        "libOpenCL.so.1",
         "/System/Library/Frameworks/OpenCL.framework/OpenCL",
 #endif
     };
diff --git a/src/runtime/vulkan_internal.h b/src/runtime/vulkan_internal.h
@@ -279,6 +279,8 @@ const char *vk_get_error_name(VkResult error) {
         return "VK_ERROR_FORMAT_NOT_SUPPORTED";
     case VK_ERROR_FRAGMENTED_POOL:
         return "VK_ERROR_FRAGMENTED_POOL";
+    case VK_ERROR_UNKNOWN:
+        return "VK_ERROR_UNKNOWN";
     case VK_ERROR_SURFACE_LOST_KHR:
         return "VK_ERROR_SURFACE_LOST_KHR";
     case VK_ERROR_NATIVE_WINDOW_IN_USE_KHR:
@@ -302,6 +304,8 @@ const char *vk_get_error_name(VkResult error) {
     }
 }
 
+#define vk_report_error(user_context, code, func) (error((user_context)) << "Vulkan: " << (func) << " returned " << vk_get_error_name((code)) << " (code: " << (code) << ") ")
+
 // --------------------------------------------------------------------------
 
 }  // namespace
diff --git a/src/runtime/vulkan_resources.h b/src/runtime/vulkan_resources.h
diff --git a/test/correctness/vector_shuffle.cpp b/test/correctness/vector_shuffle.cpp
