ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-04-19 22:39:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Qianfeng
2023-08-24 01:13:07 +08:00
committed by GitHub
parent 7c71dc7e70
commit 350d64f351
2 changed files with 12 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
client_example/13_batchnorm/batchnorm_bwd_nhwc.cpp
View File

@@ -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});
}

6
client_example/13_batchnorm/batchnorm_fwd_nhwc.cpp
View File

@@ -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});
}