#########################################################
# 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 typing import List, Union, Tuple, Dict, Any
import ml_kernels
from ml_kernels.custom_op_library import (
CustomOpCompilerInfo, CustomOpTensorInterface,
get_custom_op_library_manager)
from ml_kernels.c_function_call_helpers import (
generate_type_list_from_arguments,
load_operator_function, call_operator_function_numpy)
from afe.ir.defines import InputName
from afe.ir.attributes import CustomOpAttrs
[docs]
CustomOperation = ml_kernels.custom_op_library.CustomOperation
[docs]
OperatorFunction = ml_kernels.c_function_call_helpers.OperatorFunction
[docs]
CustomOpDatatype = ml_kernels.custom_op_library.CustomOpDatatype
_DEFAULT_COMPILER = "gcc"
_DEFAULT_AFE_LIB_PATH = "afe_custom_ops"
[docs]
ArgsList = List[Union[np.ndarray, int]]
def _create_args_list(output: np.ndarray,
input_dict: Dict[InputName, np.ndarray]
) -> ArgsList:
"""
Use the output tensor and input tensors to create an arguments list that is used
to feed in to the custom op.
:param output: np.ndarray. Output tensor
:param input_dict: Dict[InputName, np.ndarray]. Input name to its tensor
:return: ArgsList. An arguments list used to feed in to the custom op
"""
# Output is the 1st element in the list
args_list: ArgsList = [output]
int_args: List[int] = [output.ndim, *output.shape]
for _input in input_dict.values():
args_list.append(_input)
input_shape = _input.shape
int_args += [len(input_shape), *input_shape]
args_list += int_args
return args_list
[docs]
def create_custom_op_interfaces(output_shape: Tuple[int, ...],
output_dtype: CustomOpDatatype,
input_shapes: List[Tuple[int, ...]],
input_dtypes: List[CustomOpDatatype],
) -> Tuple[CustomOpTensorInterface, List[CustomOpTensorInterface]]:
"""
Use the output tensor and input tensors shapes and dtypes to create output CustomOpTensorInterface
and input CustomOpTensorInterfaces
:param output_shape: Tuple[int, ...]. Shape of the output tensor
:param output_dtype: CustomOpDatatype. Data type of the output tensor
:param input_shapes: List[Tuple[int, ...]]. List of shape of the input tensor
:param input_dtypes: List[CustomOpDatatype]. List of data type of the input tensor
:return: Tuple[CustomOpTensorInterface, List[CustomOpTensorInterface]].
(output CustomOpTensorInterface, input CustomOpTensorInterfaces)
"""
# Output interface
output_interface = CustomOpTensorInterface(output_shape,
output_dtype)
# Input interfaces
input_interfaces = []
for shape, dtype in zip(input_shapes, input_dtypes):
input_interfaces.append(CustomOpTensorInterface(shape,
dtype))
return output_interface, input_interfaces
[docs]
def create_output(shape: Tuple[int, ...],
dtype: CustomOpDatatype) -> np.ndarray:
"""
Create the output tensor that is used to store the output from custom op.
:param shape: Tuple[int, ...]. Shape of the output tensor
:param dtype: CustomOpDatatype. Data type of the output tensor
:return: np.ndarray. Output tensor
"""
dtype = np.dtype(dtype)
return np.zeros(shape=shape, dtype=dtype)
[docs]
def create_custom_op(custom_op_attrs: Dict[str, Any],
dtype: CustomOpDatatype,
output_interface: CustomOpTensorInterface,
input_interfaces: List[CustomOpTensorInterface],
target_compiler: str,
library_path: str
) -> CustomOperation:
"""
Create a CustomOperation object
:param custom_op_attrs: Dict[str, Any]. Custom op attributes
:param dtype: CustomOpDatatype. Custom op data type
:param output_interface: CustomOpTensorInterface. Output custom op interface
:param input_interfaces: List[CustomOpTensorInterface]. Input custom op interfaces
:param target_compiler: str. Target compiler. Currently support gcc
:param library_path: str. Path to store the codegen and compiled SO file
:return: CustomOperation object
"""
# Prepare input arguments for generating custom op
compiler_flags = custom_op_attrs["compiler_flags"]
code = custom_op_attrs["code"]
func_name = custom_op_attrs["func_name"]
constant_attrs = custom_op_attrs["constant_attrs"]
compiler_info = CustomOpCompilerInfo(compiler=target_compiler,
compiler_flags=compiler_flags,
library_path=library_path)
return CustomOperation(code=code,
func_name=func_name,
datatype=dtype,
constant_attrs=constant_attrs,
output_interface=output_interface,
input_interfaces=input_interfaces,
compiler_info=compiler_info)
[docs]
def create_custom_op_function(attrs: CustomOpAttrs,
input_dict: Dict[InputName, np.ndarray],
output_shape: Tuple[int, ...],
) -> Tuple[ArgsList, CustomOperation]:
"""
Initialize the custom op. Compile the custom op and put it into the
CustomOpLibraryManager. Returned the generated arguments list and
function so it can be used at the execution time.
:param attrs: CustomOpAttrs
:param input_dict: Dict[InputName, np.ndarray]. Input name to its tensor
:param output_shape: Tuple[int, ...]. Output shape
:return: Tuple[ArgsList, CustomOperation]. A tuple contains arguments list
and CustomOperation function
"""
# Convert string attributes to dictionary
custom_op_attrs = attrs.custom_op_attrs
dtype = CustomOpDatatype(custom_op_attrs["datatype"])
# Create an output tensor used to store output of custom op
output = create_output(output_shape, dtype)
# Create an argument list used to feed in to the custom op
args_list = _create_args_list(output, input_dict)
# Create output and input interfaces
input_shapes = []
input_dtypes = []
for _input in input_dict.values():
input_shapes.append(_input.shape)
input_dtypes.append(CustomOpDatatype(_input.dtype.name))
output_interface, input_interfaces = create_custom_op_interfaces(
output_shape, dtype, input_shapes, input_dtypes)
# Create custom op
target_compiler = _DEFAULT_COMPILER
library_path = _DEFAULT_AFE_LIB_PATH
custom_op = create_custom_op(custom_op_attrs, dtype, output_interface, input_interfaces,
target_compiler, library_path)
# Get the CustomOpLibraryManager
custom_op_library_manager = get_custom_op_library_manager()
assert custom_op_library_manager is not None, "Please initialize Custom Op Library Manager"
# Add the custom op into CustomOpLibraryManager, generate the C code and compile it
lib, func = custom_op_library_manager.add_custom_op(custom_op)
# Link the custom op to a Python-callable function
type_list = generate_type_list_from_arguments(args_list)
function = load_operator_function(lib, func, type_list)
return args_list, function
[docs]
def execute_custom_op(attrs: CustomOpAttrs,
input_dict: Dict[InputName, np.ndarray]) -> None:
"""
Execute the custom op
:param attrs: CustomOpAttrs
:param input_dict: Dict[InputName, np.ndarray]. Input name to its tensor
:return: np.ndarray
"""
# Updating the input tensors in existing argument list
# Output is the 1st element in the list so we skip it
for i, _input in enumerate(input_dict.values()):
attrs.args_list[i + 1] = _input
call_operator_function_numpy(attrs.function, attrs.args_list)
return attrs.args_list[0]