Skip to content

Commit

Permalink
[AOT] Add AOTLowerMain pass to lower a Relay main into TIR
Browse files Browse the repository at this point in the history
This is a pass refactored out of the AOTExecutorCodegen. Instead of
combining all of the functionality of the AOTExecutorCodegen into a
single monolithic pass, this pass only handles the lowering of the
Relay main function into TIR. Tests for the pass are included.
  • Loading branch information
mbaret committed Sep 6, 2022
1 parent b3edb6e commit 11a4573
Show file tree
Hide file tree
Showing 9 changed files with 1,508 additions and 13 deletions.
1 change: 1 addition & 0 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -305,6 +305,7 @@ tvm_file_glob(GLOB_RECURSE RELAY_PASS_SRCS
tvm_file_glob(GLOB RELAY_BACKEND_SRCS
src/relay/backend/*.cc
src/relay/backend/vm/*.cc
src/relay/backend/aot/*.cc
)
tvm_file_glob(GLOB_RECURSE RELAY_IR_SRCS
src/relay/ir/*.cc
Expand Down
21 changes: 21 additions & 0 deletions python/tvm/relay/backend/_aot.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you 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.
"""The AOT FFI namespace.
"""
import tvm._ffi

tvm._ffi._init_api("relay.backend.aot", __name__)
43 changes: 43 additions & 0 deletions python/tvm/relay/backend/aot.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you 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.
# pylint: disable=invalid-name
"""AOT passes"""
from tvm.ir.transform import Pass
from .utils import CallType

from . import _aot


def AOTLowerMain(mod_name: str, config: object, call_type: CallType) -> Pass:
"""Lower a Relay main function into an AOT TIR main function.
Parameters
----------
mod_name: str
The name of the module.
config : CompilationConfig
The compilation configuration.
call_type : CallType
The calling convention to use.
Returns
-------
Pass
The AOTLowerMain pass.
"""
return _aot.AOTLowerMain(mod_name, config, call_type.value)
7 changes: 7 additions & 0 deletions python/tvm/relay/backend/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,13 @@
# specific language governing permissions and limitations
# under the License.
"""Utility backend functions."""
from enum import Enum


class CallType(Enum):
Packed = 0
CPacked = 1
Unpacked = 2


def _is_valid_modname(mod_name):
Expand Down
Loading

0 comments on commit 11a4573

Please sign in to comment.