From 6772190f8ba6589c37b6729e1b5ab9c6bbc4d044 Mon Sep 17 00:00:00 2001
From: liuqi <liuqi10@xiaomi.com>
Date: Wed, 6 Sep 2017 18:18:49 +0800
Subject: [PATCH] Optimize a computation of batch_norm kernel and ops' BUILD
 file.

---
 mace/kernels/batch_norm.h            |  2 +-
 mace/kernels/neon/batch_norm_neon.cc |  2 +-
 mace/ops/BUILD                       | 18 +++++++++---------
 3 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/mace/kernels/batch_norm.h b/mace/kernels/batch_norm.h
index d81fdae9..136163f9 100644
--- a/mace/kernels/batch_norm.h
+++ b/mace/kernels/batch_norm.h
@@ -47,7 +47,7 @@ struct BatchNormFunctor : public BatchNormFunctorBase<D, T> {
       new_offset = offset[c] - mean[c] * new_scale;
 
       for (TIndex i = 0; i < n; ++i) {
-        TIndex pos = i * channel * sample_size + c * sample_size;
+        TIndex pos = (i * channel + c) * sample_size;
         const T* input_sample_ptr = input + pos;
         T* output_sample_ptr = output + pos;
         for (TIndex j = 0; j < sample_size; ++j) {
diff --git a/mace/kernels/neon/batch_norm_neon.cc b/mace/kernels/neon/batch_norm_neon.cc
index d307173f..7121c23d 100644
--- a/mace/kernels/neon/batch_norm_neon.cc
+++ b/mace/kernels/neon/batch_norm_neon.cc
@@ -41,7 +41,7 @@ struct BatchNormFunctor<DeviceType::NEON> : public BatchNormFunctorBase<DeviceTy
       float32x4_t new_scale_f = vdupq_n_f32(new_scale);
       float32x4_t new_offset_f = vdupq_n_f32(new_offset);
       for (TIndex i = 0; i < n; ++i) {
-        TIndex pos = i * channel * sample_size + c * sample_size;
+        TIndex pos = (i * channel + c) * sample_size;
         const float* input_sample_ptr = input + pos;
         float* output_sample_ptr = output + pos;
 
diff --git a/mace/ops/BUILD b/mace/ops/BUILD
index 0394fa1d..a408f7d9 100644
--- a/mace/ops/BUILD
+++ b/mace/ops/BUILD
@@ -12,7 +12,7 @@ load("//mace:mace.bzl", "if_android")
 cc_library(
     name = "test",
     testonly = 1,
-    srcs = [
+    hdrs = [
         "ops_test_util.h",
     ],
     deps = [
@@ -23,14 +23,14 @@ cc_library(
 
 cc_library(
     name = "ops",
-    srcs = [
-        "batch_norm.cc",
-        "relu.cc",
-    ],
-    hdrs = glob([
-        "relu.h",
-        "batch_norm.h",
-    ]),
+    srcs = glob(
+        ["*.cc"],
+        exclude = ["*_test.cc"],
+    ),
+    hdrs = glob(
+        ["*.h"],
+        exclude = ["ops_test_util.h"],
+    ),
     copts = ["-std=c++11"],
     deps = [
         "//mace/core",
-- 
GitLab