compiler.hpp 3.27 KB
Newer Older
1
//*****************************************************************************
2
// Copyright 2017-2020 Intel Corporation
3 4 5 6 7 8 9 10 11 12 13 14 15
//
// 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.
//*****************************************************************************
16

17 18
// NOTE: This file follows nGraph format style.
// Follows nGraph naming convention for public APIs only, else MLIR naming convention.
19

20 21
#pragma once

22 23 24
#include "contrib/mlir/runtime/cpu/memory_manager.hpp"
#include "ngraph/check.hpp"
#include "ngraph/descriptor/tensor.hpp"
25 26
#include "ngraph/node.hpp"

27 28 29 30
#include <mlir/IR/Builders.h>
#include <mlir/IR/Module.h>
#include <mlir/IR/Types.h>

31 32 33
#include <typeindex>
#include <unordered_map>
#include <vector>
Nagy Mostafa's avatar
Nagy Mostafa committed
34

35 36
namespace ngraph
{
37 38 39 40 41 42 43 44
    namespace descriptor
    {
        class Tensor;
    }
    namespace element
    {
        class Type;
    }
45 46 47 48
    namespace op
    {
        class CompiledKernel;
    }
49 50
    namespace runtime
    {
51
        namespace ngmlir
52
        {
53 54 55 56 57
            /// MLIR Compiler. Given an nGraph sub-graph, represented as CompiledKernel node, it
            /// translates the graph down to nGraph dialect and applies core optimizations.
            ///
            /// The compiler owns the MLIR module until compilation is done. After that,
            /// the module can be grabbed and plugged into MLIR backends.
58 59
            class MLIRCompiler
            {
Nagy Mostafa's avatar
Nagy Mostafa committed
60
            public:
61 62
                /// Initializes MLIR environment. It must be called only once.
                static void init();
Nagy Mostafa's avatar
Nagy Mostafa committed
63

64
            public:
65 66
                MLIRCompiler(const ngraph::op::CompiledKernel* compiled_kernel,
                             mlir::MLIRContext& context)
67
                    : m_compiledKernel(compiled_kernel)
68
                    , m_context(context)
nmostafa's avatar
nmostafa committed
69
                {
70 71
                    NGRAPH_CHECK(initialized,
                                 "Cannot instantiate a compiler without initializing MLIR");
nmostafa's avatar
nmostafa committed
72
                }
73

74 75
                /// Compiles a subgraph with MLIR
                void compile();
nmostafa's avatar
nmostafa committed
76

77
                mlir::OwningModuleRef& get_module() { return m_module; }
78
            private:
79
                // Converts an nGraph sub-graph to MLIR nGraph dialect.
80
                void buildNgDialectModule();
81
                // Applies any nGraph dialect optimizations
82
                void optimizeNgDialect();
83

84
            private:
85
                // Sub-graph to be compiled and executed with MLIR.
86
                const ngraph::op::CompiledKernel* m_compiledKernel;
87 88 89

                // MLIR context that holds all the MLIR information related to the sub-graph
                // compilation.
90
                mlir::MLIRContext& m_context;
91
                mlir::OwningModuleRef m_module;
92

93 94
                // Global initialization for MLIR compiler
                static bool initialized;
95 96 97 98
            };
        }
    }
}