#########################################################
# Copyright (C) 2023 SiMa Technologies, Inc.
#
# This material is SiMa proprietary and confidential.
#
# This material may not be copied or distributed without
# the express prior written permission of SiMa.
#
# All rights reserved.
#########################################################
# Code owner: Chris Rodrigues
#########################################################
"""
Compilation of Relay IR to object code for the TVM runtime.
"""
from dataclasses import dataclass
from enum import Enum
import tvm.ir
import tvm.relay
from afe._tvm._relay_transform import convert_layout_for_vector
from afe._tvm._runtime import apply_batch_dimension
# TVM target string specifying how to compile to ARM
# CPU is A65, turn on optimization flags (no Scalable Vector Extension for SiMa A65)
# +v8.2a Support ARM v8.2a instructions (fp16 and advanced SIMD)
# +fullfp16 Enable full half-precision floating point
# +neon Enable Advanced SIMD instructions
_ARM_TARGET = "llvm -device=arm_cpu -mtriple=aarch64-linux-gnu -mcpu=cortex-a65 -mattr=+v8.2a,+fullfp16,+neon"
# C++ compiler for ARM target
_ARM_CXX = "aarch64-linux-gnu-g++"
[docs]
class LibType(Enum):
"""An object file format that can be produced by TVM."""
@dataclass(frozen=True)
[docs]
class CompiledTVMObjectFile:
"""
A reference to an object file that was compiled by TVM.
The file path must be relative to the directory where all the
model compilation's output files are written.
The TVM target string is saved for error checking and diagnostics.
"""
[docs]
def compile_to_arm(model: tvm.ir.module.IRModule, filename: str, lib_type: LibType) -> CompiledTVMObjectFile:
"""
Compile a TVM IRModule to ARM code that can be executed using the TVM runtime.
:param model: Relay IR module to compile
:param filename: Name of file to create
:param lib_type: Object code format to use
:return: A reference to the compiled file
"""
# Use highest known-good optimization level. Some bugs were found with level 3.
with tvm.transform.PassContext(opt_level=2):
model = tvm.relay.transform.InferType()(model)
model = convert_layout_for_vector(model)
lib = tvm.relay.build(model, target=_ARM_TARGET)
if lib_type == LibType.shared_object:
lib.export_library(filename, cc=_ARM_CXX)
else:
lib.export_library(filename)
return CompiledTVMObjectFile(_ARM_TARGET, filename)