Source code for afe.backends.mla.afe_to_n2a_compiler.defines

#########################################################
# Copyright (C) 2021 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: Joey Chou
#########################################################
import numpy as np
from ml_dtypes import int4
from typing import Tuple

# n2a_compiler definitions
from mlc.compiler.arch_constants import ARCH_CONSTS
from mlc.compiler.model_graph import PlaceholderName, evaluate_model
from mlc.compiler.model_graph.model_graph import ModelGraph
from mlc.compiler.model_graph.vertex import Vertex
from mlc.compiler.tensor_format import TensorType, DataType
from mlc.compiler.conv2d_global_defines import ConvolutionParameters
from mlc.compiler.model_graph.evaluate import evaluate_model
from mlc.compiler.model_graph.l1_based import (
    L2CachingMode, ModelCompilationProperties, PackParameters,
    get_max_batch_size_and_quadrants_used, set_model_compilation_properties, generate_l1_based_model
)
from mlc.compiler.tessellate_param import (
    TensorDRAMLayout, TensorTessellateParameters, TessellateParameters
)

from mlc.compiler.model_graph.json_utils import export_model_to_json, import_model_from_json

# n2a_compiler operators
import mlc.operator.constructors as compiler_operator
from mlc.operator.placeholder_operators import (
    tuple_placeholder_vertex, TuplePlaceholderOperator
)

# n2a_compiler simulator API
from mlc.sim_runner import SimRunner

# ml_kernels
from ml_kernels.types import bfloat16, is_integer_type, QuantizedWeightDtypes
from ml_kernels.math_helpers import (
    normalize, RoundType, create_random_tensor,
    Activation
)
from ml_kernels.requantization import (
    ArithFoldedRequantization, FractionalZeroRequantization, Narrowing,
    TFLiteRequantization, narrowing_requantization, get_id_requantization
)
from ml_kernels.np_operators import pool_requantization
from ml_kernels.global_constants import (
    compute_bf16_exp_lut, compute_bf16_reciprocal_lut, compute_bf16_rsqrt_lut
)


# AFE definition


CompilerVertex = Vertex



CompilerInputShapes = tuple[int, int, int, int]



CompilerBiasQuantizedDtype = np.int32



CompilerWeightDtype = bfloat16



CompilerBiasDtype = np.float32



CompilerRightShfitDtype = np.uint8



CompilerCustomOpTensorDtype = 'int8'

_DEFAULT_BACKEND_LIB_PATH = "n2a_custom_ops"