From ad16a866382fa0672977edc04d0cfc5f4c37cae5 Mon Sep 17 00:00:00 2001
From: liuqi <liuqi10@xiaomi.com>
Date: Wed, 29 Nov 2017 18:07:08 +0800
Subject: [PATCH] Add multiple type operation UT for buffer_to_image op.

---
 mace/ops/buffer_to_image.cc      | 2 +-
 mace/ops/buffer_to_image_test.cc | 6 ++++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/mace/ops/buffer_to_image.cc b/mace/ops/buffer_to_image.cc
index fe726d1b..56711794 100644
--- a/mace/ops/buffer_to_image.cc
+++ b/mace/ops/buffer_to_image.cc
@@ -14,6 +14,6 @@ REGISTER_OPENCL_OPERATOR(OpKeyBuilder("BufferToImage")
 REGISTER_OPENCL_OPERATOR(OpKeyBuilder("BufferToImage")
                              .TypeConstraint<half>("T")
                              .Build(),
-                         BufferToImageOp<DeviceType::OPENCL, float>);
+                         BufferToImageOp<DeviceType::OPENCL, half>);
 
 }  //  namespace mace
diff --git a/mace/ops/buffer_to_image_test.cc b/mace/ops/buffer_to_image_test.cc
index ea5fbe21..7bd667ca 100644
--- a/mace/ops/buffer_to_image_test.cc
+++ b/mace/ops/buffer_to_image_test.cc
@@ -15,6 +15,7 @@ void TestBidirectionTransform(const int type, const std::vector<index_t> &input_
       .Input("Input")
       .Output("B2IOutput")
       .AddIntArg("buffer_type", type)
+      .AddIntArg("T", DataTypeToEnum<T>::value)
       .Finalize(net.NewOperatorDef());
 
   // Add input data
@@ -27,6 +28,7 @@ void TestBidirectionTransform(const int type, const std::vector<index_t> &input_
       .Input("B2IOutput")
       .Output("I2BOutput")
       .AddIntArg("buffer_type", type)
+      .AddIntArg("T", DataTypeToEnum<T>::value)
       .Finalize(net.NewOperatorDef());
 
   // Run
@@ -40,6 +42,10 @@ TEST(BufferToImageTest, ArgSmall) {
   TestBidirectionTransform<DeviceType::OPENCL, float>(kernels::ARGUMENT, {1});
 }
 
+TEST(BufferToImageTest, ArgHalfSmall) {
+  TestBidirectionTransform<DeviceType::OPENCL, half>(kernels::ARGUMENT, {1});
+}
+
 TEST(BufferToImageTest, ArgMedia) {
   TestBidirectionTransform<DeviceType::OPENCL, float>(kernels::ARGUMENT, {11});
 }
-- 
GitLab