From fc8a0ebf15b71a5577d1a34337278a9d9f16334c Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Mon, 21 Jan 2019 16:40:52 +0800 Subject: [PATCH 1/2] [CODEGEN][CUDA] Fix global barrier --- src/codegen/codegen_cuda.cc | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/codegen/codegen_cuda.cc b/src/codegen/codegen_cuda.cc index 0ab56a116eab..327cadefae25 100644 --- a/src/codegen/codegen_cuda.cc +++ b/src/codegen/codegen_cuda.cc @@ -225,8 +225,15 @@ void CodeGenCUDA::PrintStorageSync(const Call* op) { std::string ptr = GetUniqueName("pf"); this->stream << "volatile unsigned* " << ptr << " = &" << vid_global_barrier_state_<< ";\n"; + + this->stream << "#if __CUDA_ARCH__ >= 600\n"; + this->PrintIndent(); + this->stream << "atomicAdd_block(&" << vid_global_barrier_expect_ << ", " << num_blocks << ");\n"; + this->stream << "#else\n"; this->PrintIndent(); - this->stream << vid_global_barrier_expect_ << " += " << num_blocks << ";\n"; + this->stream << "atomicAdd(&" << vid_global_barrier_expect_ << ", " << num_blocks << ");\n"; + this->stream << "#endif\n"; + this->PrintIndent(); this->stream <<"while (" << ptr << "[0] < " << vid_global_barrier_expect_ << ");\n"; this->EndScope(wb); From 04a1a924d6fab9a1788d78b077b5af4958db4612 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Mon, 21 Jan 2019 17:03:09 +0800 Subject: [PATCH 2/2] Fix lint --- src/codegen/codegen_cuda.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/codegen/codegen_cuda.cc b/src/codegen/codegen_cuda.cc index 327cadefae25..d4f83ca4487f 100644 --- a/src/codegen/codegen_cuda.cc +++ b/src/codegen/codegen_cuda.cc @@ -228,7 +228,8 @@ void CodeGenCUDA::PrintStorageSync(const Call* op) { this->stream << "#if __CUDA_ARCH__ >= 600\n"; this->PrintIndent(); - this->stream << "atomicAdd_block(&" << vid_global_barrier_expect_ << ", " << num_blocks << ");\n"; + this->stream << "atomicAdd_block(&" << vid_global_barrier_expect_ << ", " + << num_blocks << ");\n"; this->stream << "#else\n"; this->PrintIndent(); this->stream << "atomicAdd(&" << vid_global_barrier_expect_ << ", " << num_blocks << ");\n";