hip: scalarize fixed-vector dot reductions

Author: Mike-Leo-Smith

src/backends/hip/hip_device.cpp

@@ -56,7 +56,7 @@ static constexpr char hip_shader_cache_magic[] = "LCHIPCCH";
 static constexpr auto hip_shader_cache_artifact_version = 2u;
 // Increment whenever the HIP AST/XIR/LLVM lowering contract changes in a way
 // that can alter generated code without changing the kernel AST hash.
-static constexpr auto hip_shader_cache_codegen_revision = 3u;
+static constexpr auto hip_shader_cache_codegen_revision = 4u;
 static constexpr auto hip_shader_cache_max_artifact_size = 1ull << 30u;
 static constexpr auto hip_shader_cache_payload_hash_seed =
     0x4849504341434845ull;

src/backends/hip/llvm_codegen/hip_codegen_llvm_impl_arith.cpp

@@ -98,8 +98,19 @@ llvm::Value *HIPCodegenLLVMImpl::_translate_arithmetic_inst(IB &b, FunctionConte
     };
     auto dot_product_fp = [&](llvm::Value *u, llvm::Value *v) noexcept {
         LUISA_DEBUG_ASSERT(u->getType()->isVectorTy() && v->getType()->isFPOrFPVectorTy() && u->getType() == v->getType());
-        auto zero = llvm::ConstantFP::getNegativeZero(u->getType()->getScalarType());
-        return b.CreateFAddReduce(zero, b.CreateFMul(u, v));
+        // Keep fixed-width vector reduction explicit. The AMDGPU backend can
+        // miscompile llvm.vector.reduce.fadd for a dynamically normalized
+        // float3 in complex divergent kernels, even though the intrinsic is
+        // correct in small isolated kernels. Component-wise LLVM arithmetic
+        // has the same ordered fast-math semantics without that target path.
+        auto vector_type = llvm::cast<llvm::FixedVectorType>(u->getType());
+        auto product = b.CreateFMul(u, v);
+        auto result = b.CreateExtractElement(product, uint64_t{0u});
+        for (auto i = 1u; i < vector_type->getNumElements(); i++) {
+            result = b.CreateFAdd(
+                result, b.CreateExtractElement(product, i));
+        }
+        return result;
     };
     auto inf_nan_mask_and_test = [&](llvm::Type *t) noexcept -> std::pair<llvm::Constant *, llvm::Constant *> {
         LUISA_DEBUG_ASSERT(t->isFPOrFPVectorTy());

src/tests/unit/runtime/test_hip_callable_boundary.cpp

@@ -1,7 +1,15 @@
 #include "ut/ut.hpp"
 
+#include <array>
+#include <chrono>
+#include <cmath>
 #include <cstddef>
 #include <cstdint>
+#include <cstdlib>
+#include <filesystem>
+#include <fstream>
+#include <iterator>
+#include <string>
 #include <utility>
 
 #include <luisa/core/binary_io.h>
@@ -92,6 +100,51 @@ struct CompileResult {
     luisa::vector<uint32_t> values;
 };
 
+[[nodiscard]] luisa::vector<float>
+evaluate_normalized_vectors_through_callable(
+    Device &device,
+    luisa::span<const float3> inputs,
+    uint32_t seed) noexcept {
+    Callable project = [](Float3 direction, UInt value) noexcept {
+        for (auto round = 0u;
+             round < callable_round_count; round++) {
+            auto right_shift = round % 15u + 1u;
+            auto left_shift = round % 13u + 1u;
+            auto multiplier =
+                0x9e3779b1u + round * 0x85ebca6bu;
+            auto additive =
+                0xc2b2ae35u + round * 0x27d4eb2du;
+            value = (value ^ (value >> right_shift)) * multiplier;
+            value ^= value << left_shift;
+            value += additive;
+        }
+        return direction.x + 2.0f * direction.y +
+               4.0f * direction.z +
+               cast<float>(value & 255u);
+    };
+    Kernel1D kernel = [&project](BufferFloat output,
+                                 BufferFloat3 input,
+                                 UInt initial_value) noexcept {
+        const auto index = dispatch_x();
+        const auto direction = normalize(input.read(index));
+        output.write(
+            index,
+            project(direction, initial_value + index));
+    };
+
+    auto shader = device.compile(
+        kernel, ShaderOption{.enable_cache = false});
+    auto input = device.create_buffer<float3>(inputs.size());
+    auto output = device.create_buffer<float>(inputs.size());
+    luisa::vector<float> values(inputs.size());
+    auto stream = device.create_stream();
+    stream << input.copy_from(inputs)
+           << shader(output, input, seed).dispatch(inputs.size())
+           << output.copy_to(values.data())
+           << synchronize();
+    return values;
+}
+
 [[nodiscard]] CompileResult compile_reused_callable(
     Device &device, CapturingBinaryIO &binary_io,
     uint32_t reuse_count, uint32_t seed) noexcept {
@@ -146,6 +199,28 @@ int main(int argc, char *argv[]) {
     DeviceConfig config{.binary_io = &binary_io};
     auto device = context.create_device("hip", &config);
 
+    std::error_code filesystem_error;
+    const auto original_directory =
+        std::filesystem::current_path(filesystem_error);
+    const auto dump_directory =
+        std::filesystem::temp_directory_path(filesystem_error) /
+        ("luisa_hip_callable_boundary_" +
+         std::to_string(
+             std::chrono::steady_clock::now()
+                 .time_since_epoch()
+                 .count()));
+    std::filesystem::create_directories(
+        dump_directory, filesystem_error);
+    std::filesystem::current_path(
+        dump_directory, filesystem_error);
+    expect(!filesystem_error)
+        << "failed to prepare isolated HIP LLVM dump directory";
+#if defined(_WIN32)
+    _putenv_s("LUISA_DUMP_LLVM_IR", "1");
+#else
+    setenv("LUISA_DUMP_LLVM_IR", "1", 1);
+#endif
+
     "HIP preserves shared DSL callable boundaries after optimization"_test =
         [&] {
             constexpr auto seed = 0x12345678u;
@@ -181,4 +256,77 @@ int main(int argc, char *argv[]) {
                 << "shared Luisa Callable was repeatedly expanded into the "
                    "HIP kernel";
         };
+
+    "HIP preserves fixed-vector reductions before callable boundaries"_test =
+        [&] {
+            constexpr auto seed = 0x89abcdefu;
+            constexpr std::array inputs{
+                float3{0.0f, 0.0f, 0.75f},
+                float3{0.25f, -0.5f, 2.0f},
+                float3{-1.0f, 3.0f, 0.125f},
+                float3{4.0f, -2.0f, 1.0f},
+                float3{0.5f, 0.25f, -0.75f},
+                float3{-2.0f, -1.0f, 3.0f},
+                float3{7.0f, 0.125f, -0.25f},
+                float3{0.0f, 0.0f, 1.5f},
+                float3{1.0f, 1.0f, 1.0f},
+                float3{-1.0f, 1.0f, -1.0f},
+                float3{0.125f, 8.0f, 0.5f},
+                float3{3.0f, 2.0f, -4.0f},
+                float3{-0.5f, 0.75f, 0.25f},
+                float3{9.0f, -3.0f, 2.0f},
+                float3{0.0625f, 0.125f, 0.25f},
+                float3{-6.0f, 5.0f, 4.0f}};
+            auto actual = evaluate_normalized_vectors_through_callable(
+                device, inputs, seed);
+            for (auto i = 0u; i < inputs.size(); i++) {
+                const auto &v = inputs[i];
+                const auto squared_length =
+                    v.x * v.x + v.y * v.y + v.z * v.z;
+                const auto inverse_length =
+                    1.0f / std::sqrt(squared_length);
+                const auto direction = v * inverse_length;
+                const auto expected =
+                    direction.x + 2.0f * direction.y +
+                    4.0f * direction.z +
+                    static_cast<float>(
+                        scramble_reference(seed + i) & 255u);
+                expect(std::abs(actual[i] - expected) <= 2.0e-5f)
+                    << "fixed-vector reduction changed across an "
+                       "out-of-line HIP Callable";
+            }
+        };
+
+    "HIP lowers fixed-vector dot products without LLVM reductions"_test =
+        [&] {
+            auto dumped_module_count = 0u;
+            auto retained_vector_reduction = false;
+            for (const auto &entry :
+                 std::filesystem::directory_iterator(dump_directory)) {
+                const auto filename =
+                    entry.path().filename().string();
+                if (!filename.starts_with("hip_kernel_final_") ||
+                    entry.path().extension() != ".ll") {
+                    continue;
+                }
+                ++dumped_module_count;
+                std::ifstream stream{entry.path()};
+                const std::string module{
+                    std::istreambuf_iterator<char>{stream},
+                    std::istreambuf_iterator<char>{}};
+                retained_vector_reduction |=
+                    module.find("llvm.vector.reduce.fadd") !=
+                    std::string::npos;
+            }
+            expect(dumped_module_count == 3u)
+                << "expected one final HIP LLVM module per compiled shader";
+            expect(!retained_vector_reduction)
+                << "fixed-vector dot product retained the target-unstable "
+                   "LLVM reduction intrinsic";
+        };
+
+    std::filesystem::current_path(
+        original_directory, filesystem_error);
+    std::filesystem::remove_all(
+        dump_directory, filesystem_error);
 }