[Kernel] matx runtime compatible interface

.github/workflows/test_py_kernel.yml

@@ -17,7 +17,7 @@ jobs:
         with:
           python-version: '3.11'
       - name: install requirements
-        run: pip3 install msgpack && pip3 install numpy && pip3 install sympy
+        run: pip3 install msgpack && pip3 install numpy && pip3 install sympy && pip3 install jinja2
       - name: Add directory to PATH
         run: bash ./ci/download_mlir_binary.sh && echo "$GITHUB_WORKSPACE/mlir_build/pre_build/bin" >> $GITHUB_PATH
       - name: Test Python Script

include/matxscript/runtime/mlir/convert_memref.h

@@ -0,0 +1,55 @@
+// 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.
+
+#pragma once
+
+#include "cstdlib"
+#include "matxscript/runtime/container/ndarray.h"
+#include "matxscript/runtime/data_type.h"
+#include "matxscript/runtime/dlpack.h"
+#include "memory"
+#include "string"
+#include "vector"
+
+namespace matxscript {
+namespace runtime {
+namespace mlir {
+
+std::shared_ptr<void> alloc_memref_descriptor_ptr(size_t ndim,
+                                                  DLManagedTensor* dl_managed_tensor = nullptr);
+std::shared_ptr<void> convert_from_raw_ptr(void* raw_memref, bool ok_to_free_data = false);
+std::shared_ptr<void> convert_from_dl_managed_tensor(DLManagedTensor* dl_managed_tensor);
+DLManagedTensor* convert_to_dl_managed_tensor(std::shared_ptr<void>& memref_ptr,
+                                              size_t ndim,
+                                              DLDataType dtype);
+
+NDArray convert_to_ndarray(std::shared_ptr<void>& memref_ptr, size_t ndim, DLDataType dtype);
+NDArray convert_to_ndarray(void* memref_ptr,
+                           size_t ndim,
+                           DLDataType dtype,
+                           bool ok_to_free_data = false);
+std::shared_ptr<void> convert_from_ndarray(NDArray& nd);
+
+DLDataType cvt_str_to_dl_dtype(const std::string& str);
+
+bool is_overlapping(void* target, std::initializer_list<void*> others);
+
+}  // namespace mlir
+}  // namespace runtime
+}  // namespace matxscript

include/matxscript/runtime/mlir/func_loader.h

@@ -0,0 +1,39 @@
+// 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.
+
+#pragma once
+
+#include <dlfcn.h>
+#include <matxscript/runtime/runtime_value.h>
+#include <unordered_map>
+
+namespace matxscript {
+namespace runtime {
+namespace mlir {
+
+using lib_path = std::string;
+using func_name = std::string;
+using lib_ptr_func_map_pair = std::pair<void*, std::unordered_map<func_name, void*>>;
+static std::unordered_map<lib_path, lib_ptr_func_map_pair> lib_func_map;
+
+void* load_func(const std::string& func_name, const std::string& share_lib_path);
+
+}  // namespace mlir
+}  // namespace runtime
+}  // namespace matxscript

include/matxscript/runtime/mlir/memref_cpp_interface.h

@@ -0,0 +1,147 @@
+// 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.
+
+#pragma once
+#include <iostream>
+#include "cstdlib"
+
+namespace matxscript {
+namespace runtime {
+namespace mlir {
+
+using memref_size_t = intptr_t;
+using data_t = void;
+using data_array = data_t*;
+
+template <typename T, size_t N>
+struct MemRefDescriptor {
+  T* allocated;
+  T* aligned;
+  memref_size_t offset;
+  memref_size_t sizes[N];
+  memref_size_t strides[N];
+};
+
+typedef struct {
+  data_t* allocated;
+  data_t* aligned;
+  memref_size_t offset;
+  memref_size_t rest[];
+} MemrefCInterface;
+
+class MemrefCPPInterface {
+ public:
+  // MemRefDescriptor layout:
+  // |allocated_ptr|aligned_ptr|offset|sizes|strides|
+  // | data_array  | data_array|memref_size_t*(1+2*ndim)|
+  explicit MemrefCPPInterface(void* memref_des_ptr, size_t ndim)
+      : _memref_des_ptr(memref_des_ptr), _ndim(ndim) {
+    // view_as_data_arrays layout:
+    // |allocated_ptr|aligned_ptr| rest ...
+    // | data_array  | data_array| data_array ...
+    auto ci = static_cast<MemrefCInterface*>(memref_des_ptr);
+    allocated_ptr = &(ci->allocated);
+    aligned_ptr = &(ci->aligned);
+
+    // rest layout:
+    // |offset|sizes|strides|
+    // |memref_size_t|memref_size_t*ndim|memref_size_t*ndim|
+    offset_ptr = &(ci->offset);
+    sizes_array = &(ci->rest[0]);
+    strides_array = &(ci->rest[ndim]);
+  }
+
+  void* get_memref_des_ptr() const;
+  size_t get_ndim() const;
+
+  data_t* get_allocated() const;
+  data_t* get_aligned() const;
+  memref_size_t get_offset() const;
+  memref_size_t* get_sizes() const;
+  memref_size_t* get_strides() const;
+
+  void set_allocated_ptr(data_t* new_allocated_ptr);
+  void set_aligned_ptr(data_t* new_aligned_ptr);
+  void set_offset(memref_size_t new_offset);
+  template <typename T>
+  void set_sizes(T* shape) {
+    for (int i = 0; i < _ndim; i++) {
+      sizes_array[i] = static_cast<memref_size_t>(shape[i]);
+    }
+  }
+
+  template <typename T>
+  void set_strides(T* strides) {
+    for (int i = 0; i < _ndim; i++) {
+      strides_array[i] = static_cast<memref_size_t>(strides[i]);
+    }
+  }
+
+  template <typename T>
+  void compute_strides(T* shape) {
+    uint64_t last = 1;
+    strides_array[_ndim - 1] = static_cast<memref_size_t>(last);
+
+    for (int64_t i = static_cast<int64_t>(_ndim) - 2; i >= 0; --i) {
+      last = last * shape[i + 1];
+      strides_array[i] = static_cast<memref_size_t>(last);
+    }
+  }
+
+  template <typename T>
+  void display() {
+    std::cout << "get_memref_des_ptr() = " << get_memref_des_ptr() << std::endl;
+    std::cout << "get_ndim()           = " << get_ndim() << std::endl;
+    std::cout << "allocated_ptr        = " << allocated_ptr << std::endl;
+    std::cout << "get_allocated()      = " << get_allocated() << std::endl;
+    std::cout << "aligned_ptr          = " << aligned_ptr << std::endl;
+    std::cout << "get_aligned()        = " << get_aligned() << std::endl;
+    std::cout << "offset_ptr           = " << offset_ptr << std::endl;
+    std::cout << "get_offset()         = " << get_offset() << std::endl;
+    std::cout << "get_sizes()          = " << get_sizes() << std::endl;
+    std::cout << "get_strides()        = " << get_strides() << std::endl;
+    size_t total = 1;
+    for (int i = 0; i < get_ndim(); i++) {
+      std::cout << "memrefCInterface.size_ptr[" << i << "]" << (get_sizes())[i] << std::endl;
+      total *= (get_sizes())[i];
+    }
+    for (int i = 0; i < get_ndim(); i++) {
+      std::cout << "memrefCInterface.stride_ptr[" << i << "]" << (get_strides())[i] << std::endl;
+    }
+    auto* data_ptr = static_cast<T*>(get_aligned());
+    for (int i = 0; i < total; i++) {
+      std::cout << "memrefCInterface.get_aligned()[" << i << "]" << data_ptr[i] << std::endl;
+    }
+
+    std::cout << std::endl;
+  }
+
+ private:
+  void* _memref_des_ptr;
+  size_t _ndim;
+  data_t** allocated_ptr;
+  data_t** aligned_ptr;
+  memref_size_t* offset_ptr;
+  memref_size_t* sizes_array;
+  memref_size_t* strides_array;
+};
+
+}  // namespace mlir
+}  // namespace runtime
+}  // namespace matxscript

python/matx/kernel/_template/template_func.py

@@ -19,19 +19,19 @@
 
 from typing import Dict, Any
 
-from matx.kernel.kernel_parser import KernelParser
+from matx.kernel.kernel_parser import KernelTemplateParser
 
 
 class TemplateFunc:
 
     def __init__(self, func):
-        self.fun_dict: Dict[Any, KernelParser] = {}
+        self.fun_dict: Dict[Any, KernelTemplateParser] = {}
         self.func = func
 
-    def get_function(self, args_type_list) -> KernelParser:
+    def get_function(self, args_type_list) -> KernelTemplateParser:
         args_types = tuple(args_type_list)
         if args_types not in self.fun_dict:
-            p = KernelParser(self.func, args_types)
+            p = KernelTemplateParser(self.func, args_types)
             p.parse()
             self.fun_dict[args_types] = p
         return self.fun_dict[args_types]