gpu: jit: reorder: add fallback if allocation fails

bviyer · Dec 2, 2022 · 4ee69b1 · 4ee69b1
1 parent 215a981
commit 4ee69b1
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/gpu/jit/codegen/reorder.hpp b/src/gpu/jit/codegen/reorder.hpp
@@ -179,8 +179,10 @@ bool try_emit_batched_reorder_1d_tile(ngen::HW hw, GeneratorT *host,
 
     const int grf_size = ngen::GRF::bytes(hw);
     op_plan_t plan = grf_size;
-    auto tmp = lex_scope.alloc_reg_buf_data(
-            utils::div_up(int(batch * sizeof(uint32_t)), grf_size));
+    int tmp_regs = utils::div_up(int(batch * sizeof(uint32_t)), grf_size);
+    auto tmp_range = lex_scope.try_alloc_range(tmp_regs);
+    if (tmp_range.isInvalid()) return false;
+    reg_buf_data_t tmp(reg_buf_t(hw, tmp_range));
     using inst_mod_t = ngen::InstructionModifier;
     using reg_data_t = ngen::RegData;
     auto mov = [&](inst_mod_t mod, reg_data_t dst, reg_data_t src) {