Kernel use kernelop wrapper

include/matxscript/runtime/type_helper_macros.h

@@ -77,7 +77,7 @@ struct TypeAsHelper {
       return v.template As<TO_TYPE>();
     }
     bool state = v.template Is<TO_TYPE>();
-    if (std::is_same<double, TO_TYPE>::value) {
+    if (std::is_same<double, TO_TYPE>::value || std::is_floating_point<TO_TYPE>::value) {
       state |= v.template Is<int64_t>();
     }
     if (!state) {

python/matx/kernel/codegen/cpp_template/__init__.py

@@ -0,0 +1,21 @@
+#  Copyright 2023 ByteDance Ltd. and/or its affiliates.
+#
+#  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.
+
+
+from .matx_api import render_matx_api_code

python/matx/kernel/codegen/cpp_template/base.py

@@ -0,0 +1,164 @@
+#  Copyright 2023 ByteDance Ltd. and/or its affiliates.
+#
+#  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.
+from dataclasses import dataclass
+from typing import List, TYPE_CHECKING
+import os
+import time
+import jinja2
+import numpy as np
+
+import matx.kernel.typing.utils as typing_utils
+from matx.kernel.typing import STR_TO_PYTYPE
+from matx.kernel.parser.utils import FuncReturnKind
+
+if TYPE_CHECKING:
+    from matx.kernel.kernel_parser import KernelParser
+
+TEMPLATE_DIR = os.path.dirname(os.path.abspath(__file__))
+JINJA2_ENV = jinja2.Environment(loader=jinja2.FileSystemLoader(TEMPLATE_DIR))
+
+
+@dataclass
+class MatxInterfaceCodegenMetaData:
+    """Class for keeping track of data necessary for matx interface func codegen."""
+    file_name: str
+    line_no: int
+
+    unique_id: int
+    python_func_name: str
+    func_return_kind: FuncReturnKind
+
+    arg_names: List[str]
+    arg_types: List[str]
+    arg_len: int
+
+    mlir_func_name: str
+    mlir_arg_types: List[str]
+
+    matx_func_name: str
+    matx_arg_types: List[str]
+    matx_arg_names: List[str]
+    matx_rt_type: str
+
+    c_api_func_name: str
+    c_api_arg_types: List[str]
+
+    return_ndim: int
+    return_dtype: str
+
+    def __init__(
+            self,
+            file_name: str,
+            line_no: int,
+            unique_id: int,
+            py_func_name: str,
+            func_return_kind: FuncReturnKind,
+            arg_names: List[str],
+            arg_types: List[str],
+            rt_type: str,
+            return_ndim: int,
+            return_dtype: str) -> None:
+        self.file_name = file_name
+        self.line_no = line_no
+
+        self.unique_id = unique_id
+        self.python_func_name = py_func_name
+        self.func_return_kind = func_return_kind
+
+        self.arg_names = arg_names
+        self.arg_types = arg_types
+        self.arg_len = len(arg_names)
+
+        self.mlir_func_name = f"_mlir_ciface_{py_func_name}"
+        self.mlir_arg_types = [t if t != "NDArray" else "void *" for t in arg_types]
+
+        self.matx_func_name = f"_{unique_id}_{py_func_name}_"
+        self.matx_arg_types = [*arg_types, 'void *']
+        self.matx_arg_names = [*arg_names, "handle_2_71828182846"]
+        self.matx_rt_type = rt_type
+
+        self.c_api_func_name = f"{self.matx_func_name}__c_api"
+
+        self.return_ndim = return_ndim
+        self.return_dtype = return_dtype
+
+
+PYTYPE_TO_CPP_TYPE_STR = {
+    bool: "bool",
+    int: "int32_t",
+    float: "float",
+    np.bool_: "bool",
+    np.int8: "int8_t",
+    np.int16: "int16_t",
+    np.int32: "int32_t",
+    np.int64: "int64_t",
+    np.intc: "int32_t",
+    np.uint8: "uint8_t",
+    np.uint16: "uint16_t",
+    np.uint32: "uint32_t",
+    np.uint64: "uint64_t",
+    np.uintc: "uint32_t",
+    # todo support float16
+    # np.float16 has no corresponding python builtin ctypes
+    np.float16: "__fp16",
+    np.float32: "float",
+    np.float64: "double",
+    np.longlong: "int64_t",
+    np.ulonglong: "uint64_t"
+}
+
+
+def cvt_to_cpp_type_str(t):
+    if typing_utils.is_scalar_type(t):
+        return PYTYPE_TO_CPP_TYPE_STR[t.dtype]
+    elif typing_utils.is_ndarray_type(t):
+        return "NDArray"
+    elif typing_utils.is_symbol(t):
+        return "int64_t"
+    else:
+        raise SyntaxError(f"Unsupported type {t}")
+
+
+def make_meta_data(parser: 'KernelParser') -> MatxInterfaceCodegenMetaData:
+    file_name: str = parser.file_name
+    line_no: int = parser.line_no
+
+    _nanoseconds = int(time.time() * 1e9)
+    unique_id: int = int(_nanoseconds / 100) + 0x01b21dd213814000
+    python_func_name: str = parser.func_name
+    func_return_kind: FuncReturnKind = parser.graph.func_return_kind
+
+    arg_names: List[str] = [k for k in parser.args.keys()]
+    arg_types: List[str] = [cvt_to_cpp_type_str(t) for t in parser.arg_types]
+
+    if func_return_kind.is_void():
+        return_type: str = "void"
+    elif func_return_kind.is_scalar():
+        return_type: str = PYTYPE_TO_CPP_TYPE_STR[STR_TO_PYTYPE[parser.graph.return_dtype_str]]
+    elif func_return_kind.is_dynamic_tensor():
+        return_type: str = "void"
+    else:
+        raise SyntaxError(f"Unsupported return type {func_return_kind}")
+
+    return_ndim: int = len(parser.graph.return_shape)
+    return_dtype: str = parser.graph.return_dtype_str
+    return MatxInterfaceCodegenMetaData(
+        file_name, line_no, unique_id, python_func_name, func_return_kind,
+        arg_names, arg_types, return_type, return_ndim, return_dtype
+    )

python/matx/kernel/codegen/cpp_template/function_meta_data.py

@@ -45,13 +45,14 @@ class CInterfaceCodegenData:
     return_ndim: int
     return_dtype: str
     input_types: List[str]
+    input_args: List[str]
     lib_path: str
     func_return_kind: 'FuncReturnKind'
     free_return: bool
     debug: bool
 
     def __init__(self, unique_id: int, func_name: str, return_type: str, return_ndim: int,
-                 return_dtype: str, input_types: List[str], lib_path: str,
+                 return_dtype: str, input_types: List[str], input_args: List[str], lib_path: str,
                  func_return_kind: 'FuncReturnKind', debug: bool = False):
         self.env = jinja2.Environment(loader=jinja2.FileSystemLoader(TEMPLATE_DIR))
         self.template = self.env.get_template('cpp_header.txt')
@@ -61,6 +62,7 @@ def __init__(self, unique_id: int, func_name: str, return_type: str, return_ndim
         self.return_ndim = return_ndim
         self.return_dtype = return_dtype
         self.input_types = input_types
+        self.input_args = input_args
         self.lib_path = lib_path
         self.func_return_kind = func_return_kind
         self.debug = DEBUG or debug
@@ -72,6 +74,7 @@ def code(self):
                                       return_ndim=self.return_ndim,
                                       return_dtype=self.return_dtype,
                                       input_types=self.input_types,
+                                      input_args=self.input_args,
                                       lib_path=self.lib_path,
                                       func_return_kind=self.func_return_kind,
                                       debug=self.debug)
@@ -114,13 +117,14 @@ def cvt_to_cpp_type_str(t):
         raise SyntaxError(f"Unsupported type {t}")
 
 
-def from_kernel_parser(parser: 'KernelParser', lib_path: str) -> CInterfaceCodegenData:
+def get_codegen_data(parser: 'KernelParser', lib_path: str) -> CInterfaceCodegenData:
     nanoseconds = int(time.time() * 1e9)
     unique_id: int = int(nanoseconds / 100) + 0x01b21dd213814000
     func_name: str = parser.func_name
     return_ndim: int = len(parser.graph.return_shape)
     return_dtype: str = parser.graph.return_dtype_str
     input_types: List[str] = [cvt_to_cpp_type_str(t) for t in parser.arg_types]
+    input_args: List[str] = [k for k in parser.args.keys()]
     func_return_kind: 'FuncReturnKind' = parser.graph.func_return_kind
     if func_return_kind.is_void():
         return_type: str = "void"
@@ -131,4 +135,4 @@ def from_kernel_parser(parser: 'KernelParser', lib_path: str) -> CInterfaceCodeg
     else:
         raise SyntaxError(f"Unsupported return type {func_return_kind}")
     return CInterfaceCodegenData(unique_id, func_name, return_type, return_ndim,
-                                 return_dtype, input_types, lib_path, func_return_kind)
+                                 return_dtype, input_types, input_args, lib_path, func_return_kind)

python/matx/kernel/codegen/cpp_template/matx_api.py

@@ -0,0 +1,63 @@
+#  Copyright 2023 ByteDance Ltd. and/or its affiliates.
+#
+#  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.
+
+
+from .base import make_meta_data, MatxInterfaceCodegenMetaData, JINJA2_ENV
+from .matx_api_func import MatxAPIFuncCodegen
+from .matx_c_api_func import MatxCAPIFuncCodegen
+from typing import Tuple, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from matx.kernel.kernel_parser import KernelParser
+
+
+class MatxApiCodegen:
+    """Class for keeping track of data necessary for c++ interface codegen."""
+
+    def __init__(self, meta_data: MatxInterfaceCodegenMetaData) -> None:
+        self.meta_data = meta_data
+        self.matx_api_func_code_gen = MatxAPIFuncCodegen(meta_data)
+        self.matx_c_api_func_code_gen = MatxCAPIFuncCodegen(meta_data)
+
+    def func_name(self):
+        return self.meta_data.python_func_name
+
+    def code(self):
+        code_template = JINJA2_ENV.get_template('matx_api.txt')
+        mlir_func_signature = self.matx_api_func_code_gen.gen_mlir_func_signature()
+        matx_api_declaration = self.matx_api_func_code_gen.func_declaration()
+        matx_c_api_declaration = self.matx_c_api_func_code_gen.func_declaration()
+        matx_api_definition = self.matx_api_func_code_gen.func_definition()
+        matx_c_api_definition = self.matx_c_api_func_code_gen.func_definition()
+        return code_template.render(
+            mlir_func_signature=mlir_func_signature,
+            matx_api_declaration=matx_api_declaration,
+            matx_c_api_declaration=matx_c_api_declaration,
+            matx_api_definition=matx_api_definition,
+            matx_c_api_definition=matx_c_api_definition,
+            c_interface_func_name=self.meta_data.c_api_func_name,
+            py_func_name=self.meta_data.python_func_name
+        )
+
+
+def render_matx_api_code(parser: 'KernelParser') -> Tuple[str,
+                                                          MatxInterfaceCodegenMetaData]:
+    meta_data = make_meta_data(parser)
+    codegen_class = MatxApiCodegen(meta_data)
+    return codegen_class.code(), meta_data

python/matx/kernel/codegen/cpp_template/matx_api.txt

@@ -0,0 +1,76 @@
+// Copyright 2023 ByteDance Ltd. and/or its affiliates.
+//
+// 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.
+#define _GLIBCXX_USE_CXX11_ABI 0
+
+#include "matxscript/runtime/mlir/convert_memref.h"
+#include "matxscript/runtime/mlir/func_loader.h"
+#include "matxscript/runtime/codegen_all_includes.h"
+#include <math.h>
+
+#if defined(__clang__) || (defined(__GNUC__) && __GNUC__ >= 10 )
+#if __has_builtin(__fp16)
+static_assert(sizeof (__fp16) == 2, "__fp16(a.k.a float16) size must be 2 bytes, your machine/gcc may not support __fp16 type");
+#else
+#include <float.h>
+#if __has_builtin(_Float16)
+static_assert(sizeof (_Float16) == 2, "_Float16(a.k.a float16) size must be 2 bytes, your machine/gcc may not support _Float16 type");
+using __fp16 = _Float16;
+#endif
+#endif
+#else
+using __fp16 = uint16_t;
+#endif
+
+static_assert(sizeof (float) == 4, "float(a.k.a float32) size must be 4 bytes, your machine/gcc may not support float type");
+static_assert(sizeof (double) == 8, "double(a.k.a float64) size must be 8 bytes, your machine/gcc may not support double type");
+using namespace ::matxscript::runtime;
+using namespace matxscript::runtime::mlir;
+extern "C" void* __matxscript_module_ctx = NULL;
+
+extern "C" MATX_DLL MATXScriptFuncRegistry __matxscript_func_registry__;
+extern "C" {{mlir_func_signature}};
+
+namespace {
+{{matx_api_declaration}}
+
+{{matx_c_api_declaration}}
+
+{{matx_api_definition}}
+
+{{matx_c_api_definition}}
+
+}  // namespace
+
+extern "C" {
+
+MATX_DLL MATXScriptBackendPackedCFunc __matxscript_func_array__[] = {
+    (MATXScriptBackendPackedCFunc){{c_interface_func_name}},
+};
+MATX_DLL MATXScriptFuncRegistry __matxscript_func_registry__ = {
+    "1\000{{py_func_name}}\000",    __matxscript_func_array__,
+};
+
+} // extern C
+
+extern "C" {
+
+MATX_DLL const char* __matxscript_closures_names__ = "1\000{{py_func_name}}\000";
+
+} // extern C
+