Commit 35c977cd authored by wangshaojie6's avatar wangshaojie6
Browse files

Merge branch 'add_get_work_space' into bwd_weight_bf16_splitk

parents b902f28d 86449091
......@@ -299,7 +299,7 @@ int main(int argc, char* argv[])
// alloc work space
size_t bwd_weight_workspace_size = conv->GetWorkSpaceSize(argument.get());
float ave_time = 0.f;
if(bwd_weight_workspace_size > 0)
if(std::is_same<InDataType, ck::bhalf_t>::value && split_k > 1)
{
DeviceMem wei_work_space_device_buf(bwd_weight_workspace_size);
wei_work_space_device_buf.SetZero();
......
......@@ -1269,14 +1269,6 @@ struct DeviceConvndBwdWeightXdl_C_Shuffle_Input_N_Hi_Wi_C_Weight_K_Y_X_C_Output_
arg.filter_spatial_lengths_[1] * arg.filter_spatial_lengths_[2] *
sizeof(float);
}
else
{
WorkSpaceSize = arg.Conv_K_ * 0;
}
}
else
{
WorkSpaceSize = arg.Conv_K_ * 0;
}
return WorkSpaceSize;
}
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment