Add workspace setting up for batchnorm bwd/fwd client examples (#860)

350d64f3 · Qianfeng · GitHub · 7c71dc7e · 350d64f3 · 350d64f3
Unverified Commit 350d64f3 authored Aug 24, 2023 by Qianfeng Committed by GitHub Aug 24, 2023
Showing with 12 additions and 0 deletions

client_example/13_batchnorm/batchnorm_bwd_nhwc.cpp client_example/13_batchnorm/batchnorm_bwd_nhwc.cpp +6 -0

client_example/13_batchnorm/batchnorm_fwd_nhwc.cpp client_example/13_batchnorm/batchnorm_fwd_nhwc.cpp +6 -0

No files found.
--- a/client_example/13_batchnorm/batchnorm_bwd_nhwc.cpp
+++ b/client_example/13_batchnorm/batchnorm_bwd_nhwc.cpp
@@ -191,6 +191,12 @@ int main(int argc, char* argv[])

        if(op_ptr->IsSupportedArgument(argument_ptr.get()))
        {
+            size_t workspace_sz = op_ptr->GetWorkSpaceSize(argument_ptr.get());
+
+            SimpleDeviceMem workspace(workspace_sz);
+
+            op_ptr->SetWorkSpacePointer(argument_ptr.get(), workspace.GetDeviceBuffer());
+
            invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false});
        }


--- a/client_example/13_batchnorm/batchnorm_fwd_nhwc.cpp
+++ b/client_example/13_batchnorm/batchnorm_fwd_nhwc.cpp
@@ -187,6 +187,12 @@ int main(int argc, char* argv[])

        if(op_ptr->IsSupportedArgument(argument_ptr.get()))
        {
+            size_t workspace_sz = op_ptr->GetWorkSpaceSize(argument_ptr.get());
+
+            SimpleDeviceMem workspace(workspace_sz);
+
+            op_ptr->SetWorkSpacePointer(argument_ptr.get(), workspace.GetDeviceBuffer());
+
            invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false});
        }