| 137 | } |
| 138 | |
| 139 | void RegisterBitcastOpKernel() { |
| 140 | TF_Status* status = TF_NewStatus(); |
| 141 | { |
| 142 | auto* builder = TF_NewKernelBuilder("Bitcast", tensorflow::DEVICE_CPU, |
| 143 | &BitcastOp_Create, &BitcastOp_Compute, |
| 144 | &BitcastOp_Delete); |
| 145 | TF_RegisterKernelBuilder("BitcastOp", builder, status); |
| 146 | CHECK_EQ(TF_OK, TF_GetCode(status)) |
| 147 | << "Error while registering bitcast kernel"; |
| 148 | } |
| 149 | |
| 150 | #if GOOGLE_CUDA |
| 151 | { |
| 152 | auto* builder = TF_NewKernelBuilder("Bitcast", tensorflow::DEVICE_GPU, |
| 153 | &BitcastOp_Create, &BitcastOp_Compute, |
| 154 | &BitcastOp_Delete); |
| 155 | TF_RegisterKernelBuilder("BitcastOp", builder, status); |
| 156 | CHECK_EQ(TF_OK, TF_GetCode(status)) |
| 157 | << "Error while registering CUDA bitcast kernel"; |
| 158 | } |
| 159 | #endif |
| 160 | |
| 161 | TF_DeleteStatus(status); |
| 162 | } |
| 163 | |
| 164 | // A dummy static variable initialized by a lambda whose side-effect is to |
| 165 | // register the bitcast kernel. |
no test coverage detected