[Hold for pipeline test] Support control-flow

src/nnfusion/common/common.hpp

@@ -67,6 +67,7 @@
 #include "nnfusion/core/operators/op_define/maximum.hpp"
 #include "nnfusion/core/operators/op_define/min.hpp"
 #include "nnfusion/core/operators/op_define/minimum.hpp"
+#include "nnfusion/core/operators/op_define/mod.hpp"
 #include "nnfusion/core/operators/op_define/multiply.hpp"
 #include "nnfusion/core/operators/op_define/negative.hpp"
 #include "nnfusion/core/operators/op_define/not.hpp"

src/nnfusion/core/kernels/cuda_gpu/cuda_emitter.hpp

@@ -11,6 +11,7 @@
 #include "nnfusion/engine/async_manager.hpp"
 
 DECLARE_string(fantares_codegen_server);
+DECLARE_string(ftuning_list);
 
 namespace nnfusion
 {
@@ -206,9 +207,14 @@ namespace nnfusion
                     , m_antares_ke_imp(new AntaresKEImp)
                 {
                     GENERIC_OP_LOGGING();
+                    parse_tuning_list();
                     if (!FLAGS_fantares_codegen_server.empty())
                     {
                         // NNFUSION_LOG(INFO) << "Translate for " << ctx->gnode->get_op_type();
+                        if (TuningList.find(ctx->gnode->get_op_type()) == TuningList.end())
+                        {
+                            return;
+                        }
 
                         ir = nnfusion::op::get_translation(ctx->gnode);
 #if 0
@@ -287,6 +293,7 @@ namespace nnfusion
                                                    << ctx->gnode->get_op_type();
                                 log_cache.insert(ctx->gnode->get_op_type());
                             }
+                            return;
                         }
 
                         kernel_info =
@@ -316,13 +323,27 @@ namespace nnfusion
                 std::string ir;
                 bool is_memcpy = false;
 
+                bool parse_tuning_list()
+                {
+                    auto tuninglist_str = FLAGS_ftuning_list;
+                    stringstream ss(tuninglist_str);
+                    while (ss.good())
+                    {
+                        string substr;
+                        getline(ss, substr, ',');
+                        TuningList.insert(substr);
+                    }
+                    NNFUSION_LOG(INFO) << "Kernel Tuning List: " << join(TuningList, ", ");
+                }
+
             protected:
                 // map tensor names and allocate tmp tensor
                 void process_antares_kernel_info();
                 void find_launch_config(const std::string& str, dim3& gridDim, dim3& blockDim);
                 std::vector<AntaresKernelInfo::Pointer> kernel_info;
                 std::unordered_map<std::string, std::string>
                     tensor_name_map; // antares tensor name : kernel tensor name
+                std::unordered_set<std::string> TuningList;
             };
 
             class CacheCudaEmitter : public CudaEmitter

src/nnfusion/core/kernels/cuda_gpu/cuda_kernelops.hpp

@@ -53,6 +53,7 @@ namespace nnfusion
         class GreaterEq;
         class Less;
         class LessEq;
+        class Mod;
         class Not;
         class Relu;
         class ReluBackprop;
@@ -275,6 +276,13 @@ namespace nnfusion
                 static constexpr const char* math_kernel = "x1 != 0 ? fdividef(x0, x1) : 0";
             };
 
+            template <>
+            struct CudaOpMap<nnfusion::op::Mod>
+            {
+                static constexpr const char* op = "fmod";
+                static constexpr const char* math_kernel = nullptr;
+            };
+
             template <>
             struct CudaOpMap<nnfusion::op::Sign>
             {

src/nnfusion/core/kernels/cuda_gpu/kernels/elementwise.cpp

@@ -38,6 +38,7 @@ REGISTER_EW_KERNEL(PowerBackwardExponent)
 REGISTER_EW_KERNEL(Subtract)
 REGISTER_EW_KERNEL(Divide)
 REGISTER_EW_KERNEL(DivNoNan)
+REGISTER_EW_KERNEL(Mod)
 REGISTER_EW_KERNEL(Sign)
 REGISTER_EW_KERNEL(Convert)
 REGISTER_EW_KERNEL(Equal)

src/nnfusion/core/kernels/cuda_gpu/kernels/if.cpp

@@ -0,0 +1,39 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#include "if.hpp"
+#include "../cuda_cudnn.hpp"
+#include "convolution.hpp"
+
+using namespace nnfusion;
+using namespace nnfusion::kernels;
+
+cuda::If::If(shared_ptr<KernelContext> ctx)
+    : KernelEmitter(ctx)
+{
+    std::stringstream tag;
+    tag << "_IfOP";
+    custom_tag = tag.str();
+}
+
+LanguageUnit_p cuda::If::emit_function_body()
+{
+    LanguageUnit_p _lu(new LanguageUnit(get_function_name()));
+    auto& lu = *_lu;
+
+    // function signature:
+    // extern "C" __global__ void kernel(m_context->dtypes[0]* input0, m_context->dtypes[0]* input1, m_context->dtypes[2]* output0)
+    lu << "// TODO\n";
+    return _lu;
+}
+
+LanguageUnit_p cuda::If::emit_dependency()
+{
+    LanguageUnit_p _lu(new LanguageUnit(get_function_name() + "_dep"));
+    _lu->require(header::cuda);
+    return _lu;
+}
+
+REGISTER_KERNEL_EMITTER("If",                                                      // op_name
+                        Device(CUDA_GPU).TypeConstraint(element::f32).Priority(2), // attrs
+                        cuda::If)

src/nnfusion/core/kernels/cuda_gpu/kernels/if.hpp

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#pragma once
+#include "../cuda_emitter.hpp"
+#include "../cuda_langunit.hpp"
+
+namespace nnfusion
+{
+    namespace kernels
+    {
+        namespace cuda
+        {
+            class If : public KernelEmitter
+            {
+            public:
+                If(shared_ptr<KernelContext> ctx);
+
+                LanguageUnit_p emit_function_body() override;
+                LanguageUnit_p emit_dependency() override;
+                // LanguageUnit_p emit_function_signature() override;
+            };
+        } // namespace cuda
+    }     // namespace kernels
+} // namespace nnfusion

src/nnfusion/core/kernels/cuda_gpu/kernels/loop.cpp

@@ -0,0 +1,39 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#include "loop.hpp"
+#include "../cuda_cudnn.hpp"
+#include "convolution.hpp"
+
+using namespace nnfusion;
+using namespace nnfusion::kernels;
+
+cuda::Loop::Loop(shared_ptr<KernelContext> ctx)
+    : KernelEmitter(ctx)
+{
+    std::stringstream tag;
+    tag << "_LoopOP";
+    custom_tag = tag.str();
+}
+
+LanguageUnit_p cuda::Loop::emit_function_body()
+{
+    LanguageUnit_p _lu(new LanguageUnit(get_function_name()));
+    auto& lu = *_lu;
+
+    // function signature:
+    // extern "C" __global__ void kernel(m_context->dtypes[0]* input0, m_context->dtypes[0]* input1, m_context->dtypes[2]* output0)
+    lu << "// TODO\n";
+    return _lu;
+}
+
+LanguageUnit_p cuda::Loop::emit_dependency()
+{
+    LanguageUnit_p _lu(new LanguageUnit(get_function_name() + "_dep"));
+    _lu->require(header::cuda);
+    return _lu;
+}
+
+REGISTER_KERNEL_EMITTER("Loop",                                                    // op_name
+                        Device(CUDA_GPU).TypeConstraint(element::f32).Priority(2), // attrs
+                        cuda::Loop)

src/nnfusion/core/kernels/cuda_gpu/kernels/loop.hpp

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#pragma once
+#include "../cuda_emitter.hpp"
+#include "../cuda_langunit.hpp"
+
+namespace nnfusion
+{
+    namespace kernels
+    {
+        namespace cuda
+        {
+            class Loop : public KernelEmitter
+            {
+            public:
+                Loop(shared_ptr<KernelContext> ctx);
+
+                LanguageUnit_p emit_function_body() override;
+                LanguageUnit_p emit_dependency() override;
+                // LanguageUnit_p emit_function_signature() override;
+            };
+        } // namespace cuda
+    }     // namespace kernels
+} // namespace nnfusion

src/nnfusion/core/operators/CMakeLists.txt

@@ -42,15 +42,18 @@ set(SRC
     op_define/gelu.cpp
     op_define/greater_eq.cpp
     op_define/greater.cpp
+    op_define/if.cpp
     op_define/less_eq.cpp
     op_define/less.cpp
     op_define/log.cpp
+    op_define/loop.cpp
     op_define/lrn.cpp
     op_define/max_pool.cpp
     op_define/max.cpp
     op_define/maximum.cpp
     op_define/min.cpp
     op_define/minimum.cpp
+    op_define/mod.cpp
     op_define/multiply.cpp
     op_define/negative.cpp
     op_define/not_equal.cpp

src/nnfusion/core/operators/generic_op/generic_op_define/Elementwise.cpp

@@ -44,6 +44,7 @@ static const std::unordered_map<std::string, element_op> ElementOpMap = {
      element_op(
          "divnonan",
          "(x0 / x1).when([x1 != const(0).cast(x1.dtype())], const(0).cast(input1[].dtype()))")},
+    {"Mod", element_op("fmod", "")},
     {"Square", element_op("square", "x0 * x0")},
     {"Negative", element_op("negative", "-x0")},
     {"Select", element_op("select", "x2.when([x0 == 0], x1)")},
@@ -174,6 +175,7 @@ REGISTER_ELEM_OP(Subtract)
 REGISTER_ELEM_OP(Multiply)
 REGISTER_ELEM_OP(Divide)
 REGISTER_ELEM_OP(DivNoNan)
+REGISTER_ELEM_OP(Mod)
 REGISTER_ELEM_OP(Square)
 REGISTER_ELEM_OP(Negative)
 REGISTER_ELEM_OP(Select)

src/nnfusion/core/operators/generic_op/generic_op_define/ScatterND.cpp

@@ -52,4 +52,4 @@ REGISTER_OP(ScatterND)
                 {"update_layout", vector_to_string<std::vector<std::string>>(update_layout)},
                 {"output_layout", vector_to_string<std::vector<std::string>>(output_layout)},
             });
-    });
+    });

src/nnfusion/core/operators/op_define/if.cpp

@@ -0,0 +1,58 @@
+//*****************************************************************************
+// Copyright 2017-2020 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#include "if.hpp"
+
+using namespace std;
+using namespace nnfusion::op;
+
+If::If(std::shared_ptr<nnfusion::graph::Graph>& then_branch_graph,
+       std::shared_ptr<nnfusion::graph::Graph>& else_branch_graph)
+    : Op("If")
+    , m_then_branch_graph(then_branch_graph)
+    , m_else_branch_graph(else_branch_graph)
+{
+}
+
+void If::validate_and_infer_types(std::shared_ptr<graph::GNode> gnode)
+{
+    nnfusion::Shape cond_shape = gnode->get_input_shape(0);
+    nnfusion::element::Type cond_et = gnode->get_input_element_type(0);
+    NNFUSION_CHECK(cond_shape.size() == 0)
+        << "The condition tensor of the If operation mush be scalar.";
+    NNFUSION_CHECK(cond_et == nnfusion::element::boolean)
+        << "The condition tensor of the If operation mush be boolean.";
+
+    auto then_branch_outputs = m_then_branch_graph->get_outputs();
+    auto else_branch_outputs = m_else_branch_graph->get_outputs();
+    NNFUSION_CHECK(then_branch_outputs.size() == else_branch_outputs.size())
+        << "The outputs in the then_branch and else_branch must have the same shape and "
+           "same data type.";
+    for (size_t i = 0; i < then_branch_outputs.size(); i++)
+    {
+        NNFUSION_CHECK(then_branch_outputs[i]->get_shape() == else_branch_outputs[i]->get_shape() &&
+                       then_branch_outputs[i]->get_element_type() ==
+                           else_branch_outputs[i]->get_element_type())
+            << "The outputs in the then_branch and else_branch must have the same shape and "
+               "same data type.";
+
+        gnode->set_output_type_and_shape(
+            i, then_branch_outputs[i]->get_element_type(), then_branch_outputs[i]->get_shape());
+    }
+}

src/nnfusion/core/operators/op_define/if.hpp

@@ -0,0 +1,49 @@
+//*****************************************************************************
+// Copyright 2017-2020 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "../op.hpp"
+#include "nnfusion/core/graph/graph.hpp"
+
+namespace nnfusion
+{
+    namespace op
+    {
+        /// \brief If control-flow operation, with same definition as https://github.com/onnx/onnx/blob/master/docs/Changelog.md#If-1.
+        class If : public Op
+        {
+        public:
+            /// \brief Constructs an if operation
+            ///
+            /// \param then_branch_graph The then_branch graph.<br>
+            /// `[f]`
+            /// \param else_branch_graph The else_branch graph.<br>
+            /// `[f]`
+            If(std::shared_ptr<nnfusion::graph::Graph>& then_branch_graph,
+               std::shared_ptr<nnfusion::graph::Graph>& else_branch_graph);
+
+            void validate_and_infer_types(std::shared_ptr<graph::GNode> gnode) override;
+
+        protected:
+            std::shared_ptr<nnfusion::graph::Graph> m_then_branch_graph;
+            std::shared_ptr<nnfusion::graph::Graph> m_else_branch_graph;
+        };
+    }
+}