Adding Support for tensorHalfBatchNorm

llvm/projects/hpvm-tensor-rt/tensor_runtime/include/half_precision_api.h

@@ -299,6 +299,94 @@ void* tensorHalfConvolution(void* input_ptr, void* filter_ptr,
 
 
 
+void* tensorHalfBatchNorm(void* input_ptr, void* gamma_ptr, void* beta_ptr,
+           		  void* mean_ptr, void* variance_ptr, double epsilon){
+
+  INFO("*** TensorHalfBatchNorm \n");
+  profileEvent("tensorHalfBatchNorm");
+
+  Tensor* input = (Tensor*) input_ptr;
+  Tensor* gamma = (Tensor*) gamma_ptr;
+  Tensor* beta = (Tensor*) beta_ptr;
+  Tensor* mean = (Tensor*) mean_ptr;
+  Tensor* variance = (Tensor*) variance_ptr;
+  
+  float alpha_val = 1.0f, beta_val = 0.0f;
+  hostToDeviceCopy(input);
+  hostToDeviceCopy(gamma);
+  hostToDeviceCopy(beta);
+  hostToDeviceCopy(mean);
+  hostToDeviceCopy(variance);
+
+  size_t* input_dims = input->dims.dim_sizes;
+
+  profileEvent("F2H_start");
+  
+  Tensor* input_half = (Tensor*) create4DTensor(CUDNN_DATA_HALF, CUDNN_TENSOR_NCHW,
+						input_dims[0], input_dims[1],
+						input_dims[2], input_dims[3]);
+
+  Tensor* gamma_half = (Tensor*) create4DTensor(CUDNN_DATA_HALF, CUDNN_TENSOR_NCHW,
+						1, input_dims[1],
+						1, 1);
+
+  Tensor* beta_half = (Tensor*) create4DTensor(CUDNN_DATA_HALF, CUDNN_TENSOR_NCHW,
+						1, input_dims[1],
+						1, 1);
+
+  Tensor* mean_half = (Tensor*) create4DTensor(CUDNN_DATA_HALF, CUDNN_TENSOR_NCHW,
+					       1, input_dims[1],
+					       1, 1);
+
+  Tensor* variance_half = (Tensor*) create4DTensor(CUDNN_DATA_HALF, CUDNN_TENSOR_NCHW,
+						   1, input_dims[1],
+						   1, 1);
+
+
+  changeTensorPlacement(input_half, DEVICE);
+  changeTensorPlacement(gamma_half, DEVICE);
+  changeTensorPlacement(beta_half, DEVICE);
+  changeTensorPlacement(mean_half, DEVICE);
+  changeTensorPlacement(variance_half, DEVICE);
+
+
+  f2h((float*) input->gpu_data, input_half->num_elems, (half*) input_half->gpu_data);
+  f2h((float*) gamma->gpu_data, gamma_half->num_elems, (half*) gamma_half->gpu_data);
+  f2h((float*) beta->gpu_data,  beta_half->num_elems, (half*) beta_half->gpu_data);
+  f2h((float*) mean->gpu_data,  mean_half->num_elems, (half*) mean_half->gpu_data);
+  f2h((float*) variance->gpu_data, variance_half->num_elems, (half*) variance_half->gpu_data);
+
+
+  profileEvent("F2H_end");
+
+
+  checkCUDNN(cudnnBatchNormalizationForwardInference(cudnnHandle, CUDNN_BATCHNORM_SPATIAL,
+						     &alpha_val, &beta_val,
+						     input_half->tensor_desc, input_half->gpu_data,
+						     input_half->tensor_desc, input_half->gpu_data,
+						     gamma->tensor_desc, gamma->gpu_data,
+						     beta->gpu_data, mean->gpu_data,
+						     variance->gpu_data, epsilon));
+
+						     //gamma_half->tensor_desc, gamma_half->gpu_data,
+						     // beta_half->gpu_data, mean_half->gpu_data,
+						     //variance_half->gpu_data, epsilon));
+
+
+  profileEvent("H2F_start");
+  // NOTE: Transforming half precision output to single precision
+  h2f((half*) input_half->gpu_data, input->num_elems, (float*) input->gpu_data);
+
+  profileEvent("H2F_end");
+
+  
+  profileEvent("tensorHalfBatchNorm_end", true);
+
+
+  return input;
+}
+
+
 
 
 void* tensorHalfPooling(void* input_ptr,

llvm/projects/hpvm-tensor-rt/tensor_runtime/include/tensor_runtime.h

@@ -104,6 +104,9 @@ extern "C"{
   void* tensorBatchNorm(void* input_ptr, void* gamma_ptr, void* beta_ptr,
 			void* mean_ptr, void* variance_ptr, double epsilon);
 
+  void* tensorHalfBatchNorm(void* input_ptr, void* gamma_ptr, void* beta_ptr,
+			void* mean_ptr, void* variance_ptr, double epsilon);
+
   
   /* Error injection API - used for accuracy tuning */
   void* tensorAddError(void* x_ptr, int error_scale);