Fix div by zero

86b49567 · Paul · 124ed38d · 86b49567
Commit 86b49567 authored Jun 02, 2022 by Paul
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

src/targets/gpu/compile_gen.cpp src/targets/gpu/compile_gen.cpp +4 -0

No files found.
--- a/src/targets/gpu/compile_gen.cpp
+++ b/src/targets/gpu/compile_gen.cpp
@@ -20,6 +20,10 @@ static std::vector<std::size_t> vector_sizes(const std::vector<shape>& inputs)

 vectorize vectorize::elements(std::size_t axis, const std::vector<shape>& inputs)
 {
+    if (std::all_of(inputs.begin(), inputs.end(), [&](const auto& s) {
+            return s.lens()[axis] == 1;
+        }))
+        return {1, axis};
    auto sizes = vector_sizes(inputs);
    std::vector<std::size_t> max_vec_size;
    std::transform(inputs.begin(),