Skip to content

Commit

Permalink
[OpenCL] Fix OpenCL get_valid_counts errors due to intrinsic atomic_a…
Browse files Browse the repository at this point in the history
…dd (apache#5857)

* [OpenCL] Fix atomic add used by get_valid_counts

* Rename l -> load, add flag to enable atomics

* Opencl doesn't do data rearrangement
  • Loading branch information
Trevor Morris committed Jun 30, 2020
1 parent e24c9ba commit 4fecf40
Show file tree
Hide file tree
Showing 4 changed files with 43 additions and 2 deletions.
31 changes: 31 additions & 0 deletions src/target/source/codegen_opencl.cc
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,11 @@ std::string CodeGenOpenCL::Finish() {
"#endif\n\n";
}

// Enable atomic_add used by get_valid_counts. Only needed for OpenCL < 1.1.
if (enable_atomics_) {
decl_stream << "#pragma OPENCL EXTENSION cl_khr_global_int32_base_atomics : enable\n"
"#pragma OPENCL EXTENSION cl_khr_global_int32_extended_atomics : enable\n\n";
}
return CodeGenC::Finish();
}

Expand Down Expand Up @@ -224,6 +229,32 @@ std::string CodeGenOpenCL::CastFromTo(std::string value, DataType from, DataType
return os.str();
}

void CodeGenOpenCL::VisitExpr_(const CallNode* op, std::ostream& os) {
if (op->op.same_as(builtin::address_of())) {
// Overload tvm_address_of to add storage scope (e.g. __global).
const LoadNode* load = op->args[0].as<LoadNode>();
CHECK(op->args.size() == 1 && load);
os << "((";
auto it = alloc_storage_scope_.find(load->buffer_var.get());
if (it != alloc_storage_scope_.end()) {
PrintStorageScope(it->second, os);
}
this->PrintType(load->dtype.element_of(), os);
os << " *)" << this->GetVarID(load->buffer_var.get()) << " + ";
this->PrintExpr(load->index, os);
os << ')';
} else if (op->op.same_as(builtin_call_extern_)) {
auto func = Downcast<StringImm>(op->args[0]);
// Enable atomics extension if used.
if (func->value == "atomic_add") {
enable_atomics_ = true;
}
CodeGenC::VisitExpr_(op, os);
} else {
CodeGenC::VisitExpr_(op, os);
}
}

void CodeGenOpenCL::VisitExpr_(const BroadcastNode* op, std::ostream& os) { // NOLINT(*)
std::string v = PrintExpr(op->value);
os << "((";
Expand Down
3 changes: 3 additions & 0 deletions src/target/source/codegen_opencl.h
Original file line number Diff line number Diff line change
Expand Up @@ -54,13 +54,16 @@ class CodeGenOpenCL final : public CodeGenC {
std::string CastFromTo(std::string value, DataType from, DataType target); // NOLINT(*)

// overload visitor
void VisitExpr_(const CallNode* op, std::ostream& os) final; // NOLINT(*)
void VisitExpr_(const BroadcastNode* op, std::ostream& os) final; // NOLINT(*)
void VisitExpr_(const FloatImmNode* op, std::ostream& os) final; // NOLINT(*)

private:
// whether enable fp16 and fp64 extension
bool enable_fp16_{false};
bool enable_fp64_{false};
// Whether to enable atomics extension.
bool enable_atomics_{false};
};

} // namespace codegen
Expand Down
4 changes: 2 additions & 2 deletions tests/python/relay/test_op_level5.py
Original file line number Diff line number Diff line change
Expand Up @@ -270,8 +270,8 @@ def verify_get_valid_counts(dshape, score_threshold, id_index, score_index):
intrp = relay.create_executor("debug", ctx=ctx, target=target)
out = intrp.evaluate(func)(np_data)
tvm.testing.assert_allclose(out[0].asnumpy(), np_out1, rtol=1e-3, atol=1e-04)
# get_valid_count for cuda doesn't do data rearrangement
if target == 'cuda':
# get_valid_count for cuda, opencl doesn't do data rearrangement
if target in ['cuda', 'opencl']:
return
tvm.testing.assert_allclose(out[1].asnumpy(), np_out2, rtol=1e-3, atol=1e-04)
tvm.testing.assert_allclose(out[2].asnumpy(), np_out3, rtol=1e-3, atol=1e-04)
Expand Down
7 changes: 7 additions & 0 deletions topi/python/topi/cuda/nms.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,10 +34,17 @@ def cuda_atomic_add_rule(op):
return tvm.tir.call_pure_extern("int32", "atomicAdd", op.args[0], op.args[1])
raise RuntimeError("only support int32, float32 and float64")

def opencl_atomic_add_rule(op):
if op.dtype == "int32":
return tvm.tir.call_pure_extern("int32", "atomic_add", op.args[0], op.args[1])
raise RuntimeError("only support int32")

tvm.target.intrin.register_intrin_rule(
"cuda", "atomic_add", cuda_atomic_add_rule, override=True)

tvm.target.intrin.register_intrin_rule(
"opencl", "atomic_add", opencl_atomic_add_rule, override=True)

tvm.ir.register_op_attr("tir.atomic_add", "TCallEffectKind", tvm.tir.CallEffectKind.Opaque)

def atomic_add(x, y):
Expand Down

0 comments on commit 4fecf40

Please sign in to comment.