Use optimized 1x1 conv implementation.

Change-Id: I6be5c3f6fc8b981ab738fb492e8b3318544aa289

tensorflow/lite/micro/kernels/kelvin/conv.cc

diff --git a/tensorflow/lite/micro/kernels/kelvin/conv.cc b/tensorflow/lite/micro/kernels/kelvin/conv.cc
index 0d2f068..03904bc 100644
--- a/tensorflow/lite/micro/kernels/kelvin/conv.cc
+++ b/tensorflow/lite/micro/kernels/kelvin/conv.cc
@@ -16,6 +16,12 @@
 namespace tflite {
 namespace {
 
+constexpr int kFilterHeightIndex = 1;
+constexpr int kFilterWidthIndex = 2;
+constexpr int kFilterInputChannelIndex = 3;
+constexpr int kInputChannelIndex = 3;
+constexpr int kOutputChannelIndex = 3;
+
 void* Init(TfLiteContext* context, const char* buffer, size_t length) {
   TFLITE_DCHECK(context->AllocatePersistentBuffer != nullptr);
   return context->AllocatePersistentBuffer(context, sizeof(OpDataConv));
@@ -85,7 +91,64 @@
           break;
         }
         case kTfLiteInt64: {
-          const auto fn = opt ? kelvin::opt::conv_per_channel_b64 : reference_integer_ops::ConvPerChannel<int64_t>;
+          const auto filter_shape = tflite::micro::GetTensorShape(filter);
+          const auto input_shape = tflite::micro::GetTensorShape(input);
+          const auto output_shape = tflite::micro::GetTensorShape(output);
+          if (opt && filter_shape.Dims(kFilterHeightIndex) == 1 &&
+                     filter_shape.Dims(kFilterWidthIndex) == 1 &&
+                     filter_shape.Dims(kFilterInputChannelIndex) % 32 == 0 &&
+                     filter_shape.Dims(kOutputChannelIndex) % 2 == 0) {
+            kelvin::opt::conv_per_channel_b64_1x1(
+                params_q,
+                data.per_channel_output_multiplier,
+                data.per_channel_output_shift,
+                tflite::micro::GetTensorShape(input),
+                tflite::micro::GetTensorData<int16_t>(input),
+                tflite::micro::GetTensorShape(filter),
+                tflite::micro::GetTensorData<int8_t>(filter),
+                tflite::micro::GetTensorShape(bias),
+                tflite::micro::GetOptionalTensorData<std::int64_t>(bias),
+                tflite::micro::GetTensorShape(output),
+                tflite::micro::GetTensorData<int16_t>(output));
+            break;
+          }
+          // TODO(derekjchow): Check for valid padding
+          if (opt && filter_shape.Dims(kFilterHeightIndex) == 1
+                  && filter_shape.Dims(kFilterInputChannelIndex) % 32 == 0
+                  && filter_shape.Dims(kOutputChannelIndex) % 2 == 0) {
+            if (input_shape.Dims(kInputChannelIndex) ==
+                filter_shape.Dims(kFilterInputChannelIndex)) {
+              kelvin::opt::conv_per_channel_b64_filter1xn_non_group(
+                  params_q,
+                  data.per_channel_output_multiplier,
+                  data.per_channel_output_shift,
+                  tflite::micro::GetTensorShape(input),
+                  tflite::micro::GetTensorData<int16_t>(input),
+                  tflite::micro::GetTensorShape(filter),
+                  tflite::micro::GetTensorData<int8_t>(filter),
+                  tflite::micro::GetTensorShape(bias),
+                  tflite::micro::GetOptionalTensorData<std::int64_t>(bias),
+                  tflite::micro::GetTensorShape(output),
+                  tflite::micro::GetTensorData<int16_t>(output));
+              break;
+            } else if (filter_shape.Dims(kFilterInputChannelIndex) % 32 == 0) {
+              kelvin::opt::conv_per_channel_b64_filter1xn_group(
+                  params_q,
+                  data.per_channel_output_multiplier,
+                  data.per_channel_output_shift,
+                  tflite::micro::GetTensorShape(input),
+                  tflite::micro::GetTensorData<int16_t>(input),
+                  tflite::micro::GetTensorShape(filter),
+                  tflite::micro::GetTensorData<int8_t>(filter),
+                  tflite::micro::GetTensorShape(bias),
+                  tflite::micro::GetOptionalTensorData<std::int64_t>(bias),
+                  tflite::micro::GetTensorShape(output),
+                  tflite::micro::GetTensorData<int16_t>(output));
+              break;
+            }
+          }
+          const auto fn = opt ? kelvin::opt::conv_per_channel_b64
+                              : reference_integer_ops::ConvPerChannel<int64_t>;
           fn(
               params_q,
               data.per_channel_output_multiplier, data.per_channel_output_shift,