Return nullptr when block index is invalid (#1649)

13332998 · Po Yen Chen · GitHub · bec6fbc6 · 13332998
Unverified Commit 13332998 authored Nov 11, 2024 by Po Yen Chen Committed by GitHub Nov 11, 2024
Show whitespace changes
Inline Side-by-side

Showing with 9 additions and 1 deletion

include/ck_tile/ops/fmha/block/page_block_navigator.hpp include/ck_tile/ops/fmha/block/page_block_navigator.hpp +9 -1

No files found.
--- a/include/ck_tile/ops/fmha/block/page_block_navigator.hpp
+++ b/include/ck_tile/ops/fmha/block/page_block_navigator.hpp
@@ -230,7 +230,15 @@ struct PageBlockNavigator
    CK_TILE_HOST_DEVICE
    DataType* get_block_ptr(index_t block_index) const
    {
-        return physical_blocks + physical_block_indices[block_index] * block_stride + fixed_offset;
+        if(block_index < num_blocks)
+        {
+            return physical_blocks + physical_block_indices[block_index] * block_stride +
+                   fixed_offset;
+        }
+        else
+        {
+            return nullptr;
+        }
    }

    CK_TILE_HOST_DEVICE int32_t get_block_index(const WindowOrigin& global_window_origin) const