ecsimsw
TF_source ) kernel_builder 본문
///
REGISTER_KERNEL_BUILDER(Name("SparseApplyAdadelta") \
.Device(DEVICE_CPU) \
.TypeConstraint("T") \
.TypeConstraint("Tindices"), \
SparseApplyAdadeltaOp<T, Tindices>);
///
#define REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, ...) \
constexpr bool should_register_##ctr##__flag = \
SHOULD_REGISTER_OP_KERNEL(#__VA_ARGS__); \
static ::tensorflow::kernel_factory::OpKernelRegistrar \
registrar__body__##ctr##__object( \
should_register_##ctr##__flag \
? ::tensorflow::register_kernel::kernel_builder.Build() \
: nullptr, \
#__VA_ARGS__, \
[](::tensorflow::OpKernelConstruction* context) \
-> ::tensorflow::OpKernel* { \
return new __VA_ARGS__(context); \
});
///
namespace register_kernel {
class Name : public KernelDefBuilder {
public:
// With selective registration, kernels whose implementation class is not used
// by any kernel are disabled with the SHOULD_REGISTER_OP_KERNEL call in
// REGISTER_KERNEL_BUILDER_UNIQ. However, an unused kernel that shares an
// implementation class with a used kernel would get through that mechanism.
//
// This mechanism stops that registration by changing the name of the kernel
// for the unused op to one that is ignored by
// OpKernelRegistrar::InitInternal. Note that this method alone is
// not sufficient - the compiler can't evaluate the entire KernelDefBuilder at
// compilation time, so this method doesn't actually reduce code size.
explicit Name(const char* op)
: KernelDefBuilder(SHOULD_REGISTER_OP(op) ? op : "_no_register") {}
};
///
KernelDefBuilder::KernelDefBuilder(const char* op_name) {
kernel_def_ = new KernelDef;
kernel_def_->set_op(op_name);
}
const KernelDef* KernelDefBuilder::Build() {
KernelDef* r = kernel_def_;
kernel_def_ = nullptr;
return r;
}
KernelDefBuilder& KernelDefBuilder::Device(const char* device_type) {
kernel_def_->set_device_type(device_type);
return *this;
}
KernelDefBuilder& KernelDefBuilder::TypeConstraint(
const char* attr_name, gtl::ArraySlice allowed) {
auto* constraint = kernel_def_->add_constraint();
constraint->set_name(attr_name);
auto* allowed_values = constraint->mutable_allowed_values()->mutable_list();
for (DataType dt : allowed) {
allowed_values->add_type(dt);
}
return *this;
}
KernelDefBuilder& KernelDefBuilder::TypeConstraint(const char* attr_name,
DataType allowed) {
auto* constraint = kernel_def_->add_constraint();
constraint->set_name(attr_name);
constraint->mutable_allowed_values()->mutable_list()->add_type(allowed);
return *this;
}
'Machine Learning > tf_source' 카테고리의 다른 글
TF_source ) matmul_op.cc (0) | 2019.04.30 |
---|---|
TF_source ) 2019-04-25 (0) | 2019.04.25 |
What's the difference between user registers and kernel registers? (0) | 2019.04.20 |
TF_source ) KernelRegistry (0) | 2019.04.20 |
TF_source) Opkernel (0) | 2019.04.12 |