#ifdef USE_ACL

#include <vector>

#include "caffe/layers/acl_batch_norm_layer.hpp"

namespace caffe {

template <typename Dtype>

void ACLBatchNormLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,

const vector<Blob<Dtype>*>& top) {

BatchNormLayer<Dtype>::LayerSetUp(bottom, top);

this->force_bypass_acl_path_= bypass_acl_class_layer & FLAGS_ENABLE_ACL_BN;

}

template <typename Dtype>

void ACLBatchNormLayer<Dtype>::SetupACLLayer(const vector<Blob<Dtype>*>& bottom,

const vector<Blob<Dtype>*>& top){

if (!this->init_layer_) return;

this->init_layer_=false;

// Initialize ACL.

if (Caffe::arm_gpu_mode()) {

new_gpulayer();

}else{

new_cpulayer();

}

this->force_bypass_acl_path_=false;

TensorShape in_shape ((unsigned int)bottom[0]->width(), (unsigned int)bottom[0]->height(),(unsigned int)bottom[0]->channels(),(unsigned int)bottom[0]->num());

TensorShape out_shape((unsigned int)top[0]->width(), (unsigned int)top[0]->height(),(unsigned int)top[0]->channels(),(unsigned int)top[0]->num());

TensorShape mean_shape((unsigned int)this->channels_);

TensorShape var_shape=mean_shape;

TensorShape beta_shape=mean_shape;

TensorShape gamma_shape=mean_shape;

Dtype beta_val[beta_shape.total_size()];

Dtype gamma_val[gamma_shape.total_size()];

for (int i=0;i<beta_shape.total_size();++i) {

beta_val[i]=0.0;

}

for (int i=0;i<gamma_shape.total_size();++i) {

gamma_val[i]=1.0;

}

if (Caffe::arm_gpu_mode()) {

Dtype *top_data = top[0]->mutable_gpu_data();

const Dtype* bottom_data = bottom[0]->gpu_data();

// use the stored mean/variance estimates.

const Dtype scale_factor = this->blobs_[2]->cpu_data()[0] == 0 ?

0 : 1 / this->blobs_[2]->cpu_data()[0];

caffe_cpu_scale(this->variance_.count(), scale_factor,

this->blobs_[0]->gpu_data(), this->mean_.mutable_gpu_data());

caffe_cpu_scale(this->variance_.count(), scale_factor,

this->blobs_[1]->gpu_data(), this->variance_.mutable_gpu_data());

new_tensor(this->gpu().input,in_shape,(void*)bottom_data);

new_tensor(this->gpu().output,out_shape,(void*)top_data);

new_tensor(this->gpu().mean,mean_shape);

new_tensor(this->gpu().var,var_shape);

new_tensor(this->gpu().beta,beta_shape);

new_tensor(this->gpu().gamma,gamma_shape);

tensor_mem(this->gpu().mean,(void*)this->mean_.mutable_gpu_data());

tensor_mem(this->gpu().var,(void*)this->variance_.mutable_gpu_data());

tensor_mem(this->gpu().beta,(void*)beta_val);

tensor_mem(this->gpu().gamma,(void*)gamma_val);

this->gpu().mean->commit();

this->gpu().var->commit();

this->gpu().beta->commit();

this->gpu().gamma->commit();

#ifdef USE_PROFILING

logtime_util log_time(ACL_CONFIG_INFO);

#endif //USE_PROFILING

this->gpu().layer->configure(this->gpu().input,this->gpu().output,this->gpu().mean,this->gpu().var,this->gpu().beta,this->gpu().gamma,this->eps_);

}else{

Dtype *top_data = top[0]->mutable_cpu_data();

const Dtype* bottom_data = bottom[0]->cpu_data();

// use the stored mean/variance estimates.

const Dtype scale_factor = this->blobs_[2]->cpu_data()[0] == 0 ?

0 : 1 / this->blobs_[2]->cpu_data()[0];

caffe_cpu_scale(this->variance_.count(), scale_factor,

this->blobs_[0]->cpu_data(), this->mean_.mutable_cpu_data());

caffe_cpu_scale(this->variance_.count(), scale_factor,

this->blobs_[1]->cpu_data(), this->variance_.mutable_cpu_data());

new_tensor(this->cpu().input,in_shape,(void*)bottom_data);

new_tensor(this->cpu().output,out_shape,(void*)top_data);

new_tensor(this->cpu().mean,mean_shape);

new_tensor(this->cpu().var,var_shape);

new_tensor(this->cpu().beta,beta_shape);

new_tensor(this->cpu().gamma,gamma_shape);

tensor_mem(this->cpu().mean,(void*)this->mean_.mutable_cpu_data());

tensor_mem(this->cpu().var,(void*)this->variance_.mutable_cpu_data());

tensor_mem(this->cpu().beta,(void*)beta_val);

tensor_mem(this->cpu().gamma,(void*)gamma_val);

this->cpu().mean->commit();

this->cpu().var->commit();

this->cpu().beta->commit();

this->cpu().gamma->commit();

#ifdef USE_PROFILING

logtime_util log_time(ACL_CONFIG_INFO);

#endif //USE_PROFILING

this->cpu().layer->configure(this->cpu().input,this->cpu().output,this->cpu().mean,this->cpu().var,this->cpu().beta,this->cpu().gamma,this->eps_);

}

}

template <typename Dtype>

void ACLBatchNormLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,

const vector<Blob<Dtype>*>& top) {

BatchNormLayer<Dtype>::Reshape(bottom, top);

}

template <typename Dtype>

void ACLBatchNormLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,

const vector<Blob<Dtype>*>& top) {

if(Caffe::arm_gpu_mode()){

Forward_gpu(bottom, top);

return;

}

#ifdef USE_PROFILING

logtime_util log_time(ACL_BN_INFO);

#endif //USE_PROFILING

if (this->force_bypass_acl_path_||!this->use_global_stats_) {

BatchNormLayer<Dtype>::Forward_cpu(bottom,top);

return;

}

const Dtype* bottom_data = bottom[0]->cpu_data();

Dtype* top_data = top[0]->mutable_cpu_data();

SetupACLLayer(bottom,top);

tensor_mem(this->cpu().input,(void*)(bottom_data));

cpu_run();

tensor_mem((void*)(top_data),this->cpu().output);

}

template <typename Dtype>

void ACLBatchNormLayer<Dtype>::Forward_gpu(const vector<Blob<Dtype>*>& bottom,

const vector<Blob<Dtype>*>& top) {

if (this->force_bypass_acl_path_||!this->use_global_stats_) {

BatchNormLayer<Dtype>::Forward_cpu(bottom,top);

return;

}

#ifdef USE_PROFILING

logtime_util log_time(ACL_BN_INFO);

#endif //USE_PROFILING

const Dtype* bottom_data = bottom[0]->gpu_data();

Dtype* top_data = top[0]->mutable_gpu_data();

SetupACLLayer(bottom,top);

tensor_mem(this->gpu().input,(void*)(bottom_data));

gpu_run();

tensor_mem((void*)(top_data),this->gpu().output);

}

template <typename Dtype>

ACLBatchNormLayer<Dtype>::~ACLBatchNormLayer() {

}

INSTANTIATE_CLASS(ACLBatchNormLayer);

} // namespace caffe

#endif // USE_ACL