compiler/loco/src/Service/ShapeInference.test.cpp

   1 /*
   2  * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *    http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include "loco/Service/ShapeInference.h"
  18 #include "GraphTestcase.h"
  19
  20 #include <vector>
  21
  22 #include <gtest/gtest.h>
  23
  24 // This test validates whether framework works as expected.
  25 TEST(ShapeInferenceTest, framework)
  26 {
  27   // Mock-up Shape Inference Rule
  28   struct SampleShapeInferenceRule final : public loco::ShapeInferenceRule
  29   {
  30   public:
  31     SampleShapeInferenceRule(std::vector<const loco::Node *> *nodes) : _nodes{nodes}
  32     {
  33       // DO NOTHING
  34     }
  35
  36   public:
  37     // Accept all the dialects
  38     bool recognize(const loco::Dialect *) const final { return true; }
  39
  40     bool infer(const loco::Node *node, loco::NodeShape &shape) const final
  41     {
  42       // Record the order of inference
  43       _nodes->emplace_back(node);
  44
  45       if (_nodes->size() != 1)
  46       {
  47         return false;
  48       }
  49
  50       // Set the first node as Tensor<1>
  51       loco::TensorShape tensor_shape;
  52
  53       tensor_shape.rank(1);
  54       tensor_shape.dim(0) = 4;
  55
  56       shape.set(tensor_shape);
  57
  58       return true;
  59     }
  60
  61   private:
  62     std::vector<const loco::Node *> *_nodes;
  63   };
  64
  65   GraphTestcase<GraphCode::Identity> testcase;
  66
  67   std::vector<const loco::Node *> nodes;
  68
  69   SampleShapeInferenceRule rule{&nodes};
  70
  71   loco::apply(&rule).to(testcase.graph());
  72
  73   // Framework SHOULD visit all the nodes
  74   ASSERT_EQ(2, nodes.size());
  75   // Framework SHOULD visit "pull" before "push"
  76   ASSERT_EQ(testcase.pull_node, nodes.at(0));
  77   ASSERT_EQ(testcase.push_node, nodes.at(1));
  78
  79   // Framework SHOULD make an annotation if "rule" returns TRUE
  80   ASSERT_TRUE(loco::shape_known(testcase.pull_node));
  81   ASSERT_EQ(loco::Domain::Tensor, loco::shape_get(testcase.pull_node).domain());
  82   ASSERT_EQ(1, loco::shape_get(testcase.pull_node).as<loco::TensorShape>().rank());
  83   ASSERT_EQ(4, loco::shape_get(testcase.pull_node).as<loco::TensorShape>().dim(0));
  84
  85   // Framework SHOULD NOT make any annotation if "rule" returns FALSE
  86   ASSERT_FALSE(loco::shape_known(testcase.push_node));
  87 }