add workspace size check

afebb964 · wangshaojie6 · e8d6434d · afebb964
Commit afebb964 authored Jun 10, 2022 by wangshaojie6
Show whitespace changes
Inline Side-by-side

Showing with 6 additions and 0 deletions

example/20_convnd_bwd_weight_xdl/convnd_bwd_weight_xdl_bf16_splitk.cpp ...nvnd_bwd_weight_xdl/convnd_bwd_weight_xdl_bf16_splitk.cpp +6 -0

No files found.
--- a/example/20_convnd_bwd_weight_xdl/convnd_bwd_weight_xdl_bf16_splitk.cpp
+++ b/example/20_convnd_bwd_weight_xdl/convnd_bwd_weight_xdl_bf16_splitk.cpp
@@ -324,6 +324,12 @@ int main(int argc, char* argv[])

    // alloc work space
    size_t bwd_weight_workspace_size = conv->GetWorkSpaceSize(argument.get());
+    if(bwd_weight_workspace_size <= 0)
+    {
+        print_use_msg();
+        exit(1);
+    }
+
    float conv_ave_time              = 0.f;

    DeviceMem wei_work_space_device_buf(bwd_weight_workspace_size);