Support string type inputs from Python application code. (#2629)

* Support string type inputs from Python application code.

---------
Signed-off-by: Yasushi Negishi <negishi@jp.ibm.com>

include/onnx-mlir/Runtime/OnnxDataTypeMetaData.inc

@@ -16,7 +16,7 @@ OM_TYPE_METADATA_DEF(ONNX_TYPE_UINT16,      4,  sizeof(uint16_t), "uint16_t")
 OM_TYPE_METADATA_DEF(ONNX_TYPE_INT16,       5,  sizeof(int16_t),  "int16_t")
 OM_TYPE_METADATA_DEF(ONNX_TYPE_INT32,       6,  sizeof(int32_t),  "int32_t")
 OM_TYPE_METADATA_DEF(ONNX_TYPE_INT64,       7,  sizeof(int64_t),  "int64_t")
-OM_TYPE_METADATA_DEF(ONNX_TYPE_STRING,      8,  0,                "const char *")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_STRING,      8,  sizeof(const char *), "const char *")
 OM_TYPE_METADATA_DEF(ONNX_TYPE_BOOL,        9,  sizeof(bool),     "_Bool")
 OM_TYPE_METADATA_DEF(ONNX_TYPE_FLOAT16,     10, 2,                "_Float16")
 OM_TYPE_METADATA_DEF(ONNX_TYPE_DOUBLE,      11, sizeof(double),   "double")

src/Runtime/CMakeLists.txt

@@ -3,6 +3,7 @@
 
 add_subdirectory(jni)
 add_subdirectory(omp)
+add_subdirectory(python)
 
 # TODO: should add for each accelerator its subdirectory that implements InitAccel##name
 # and ShutdownAccel##name.
@@ -77,112 +78,3 @@ set_target_properties(OMExecutionSession
   PROPERTIES
   POSITION_INDEPENDENT_CODE TRUE
   )
-
-# TODO: Remove pybind11::python_link_helper after cmake version bumped to 3.16+
-add_onnx_mlir_library(OMPyExecutionSessionBase
-  PyExecutionSessionBase.cpp
-
-  EXCLUDE_FROM_OM_LIBS
-
-  LINK_LIBS PUBLIC
-  OMExecutionSession
-  OMMlirUtilities
-  pybind11::embed
-  pybind11::python_link_helper
-  )
-if(MSVC)
-  target_link_libraries(OMPyExecutionSessionBase
-    PRIVATE pybind11::windows_extras
-  )
-endif()
-set_target_properties(OMPyExecutionSessionBase
-  PROPERTIES
-  POSITION_INDEPENDENT_CODE TRUE
-  )
-target_compile_definitions(OMPyExecutionSessionBase
-  PRIVATE
-  $<TARGET_PROPERTY:onnx,COMPILE_DEFINITIONS>
-  )
-target_include_directories(OMPyExecutionSessionBase
-  PRIVATE
-  $<TARGET_PROPERTY:onnx,INCLUDE_DIRECTORIES>
-  )
-
-# When running on ubi8 image, shared lib backend tests fail with
-# the following error:
-#
-#    [libprotobuf ERROR google/protobuf/descriptor_database.cc:641] File already exists in database: onnx/onnx-ml.proto
-#    [libprotobuf FATAL google/protobuf/descriptor.cc:1371] CHECK failed: GeneratedDatabase()->Add(encoded_file_descriptor, size):
-#    terminate called after throwing an instance of 'google::protobuf::FatalException'
-#      what():  CHECK failed: GeneratedDatabase()->Add(encoded_file_descriptor, size):
-#    Aborted (core dumped)
-#
-# This is because test.py loads (among others) the following
-# two .so shared libs:
-#
-#   - onnx_cpp2py_export.cpython-39-s390x-linux-gnu.so
-#     (import onnx)
-#   - PyRuntime.cpython-39-s390x-linux-gnu.so
-#     (from PyRuntime import OMExecutionSession)
-#
-# Both libs share the same libprotobuf.so when loaded by test.py.
-# However, they were both built with the same onnx-ml.pb.cc generated
-# from onnx-ml.proto and the protobuf runtime requires all compiled-in
-# .proto files have unique names. Hence the error.
-#
-# PyRuntime doesn't really need onnx beyond the onnx::TensorProto::*
-# types so we remove onnx from its target_link_libraries. But that
-# also removes some of the compile definitions and include directories
-# which we add back through target_compile_definitions and
-# target_include_directories.
-pybind11_add_module(PyRuntime PyExecutionSession.cpp)
-add_dependencies(PyRuntime onnx_proto)
-target_compile_options(PyRuntime
-  PRIVATE
-  $<$<OR:$<CXX_COMPILER_ID:Clang>,$<CXX_COMPILER_ID:AppleClang>,$<CXX_COMPILER_ID:GNU>>:-frtti -fexceptions>
-  $<$<CXX_COMPILER_ID:MSVC>:/EHsc /GR>
-  )
-target_compile_definitions(PyRuntime
-  PRIVATE
-  $<TARGET_PROPERTY:onnx,COMPILE_DEFINITIONS>
-  )
-target_include_directories(PyRuntime
-  PRIVATE
-  $<TARGET_PROPERTY:onnx,INCLUDE_DIRECTORIES>
-  )
-target_link_libraries(PyRuntime
-  PRIVATE
-  OMPyExecutionSessionBase
-  )
-llvm_update_compile_flags(PyRuntime)
-
-install(TARGETS PyRuntime
-  DESTINATION lib
-  )
-
-pybind11_add_module(PyCompileAndRuntime PyOMCompileExecutionSession.cpp)
-add_dependencies(PyCompileAndRuntime onnx_proto)
-target_compile_options(PyCompileAndRuntime
-  PRIVATE
-  $<$<OR:$<CXX_COMPILER_ID:Clang>,$<CXX_COMPILER_ID:AppleClang>,$<CXX_COMPILER_ID:GNU>>:-frtti -fexceptions>
-  $<$<CXX_COMPILER_ID:MSVC>:/EHsc /GR>
-  )
-target_compile_definitions(PyCompileAndRuntime
-  PRIVATE
-  $<TARGET_PROPERTY:onnx,COMPILE_DEFINITIONS>
-  )
-target_include_directories(PyCompileAndRuntime
-  PRIVATE
-  ${ONNX_MLIR_SRC_ROOT}/include
-  $<TARGET_PROPERTY:onnx,INCLUDE_DIRECTORIES>
-  )
-target_link_libraries(PyCompileAndRuntime
-  PRIVATE
-  OMCompiler
-  OMPyExecutionSessionBase
-  )
-llvm_update_compile_flags(PyCompileAndRuntime)
-
-install(TARGETS PyCompileAndRuntime
-  DESTINATION lib
-  )

src/Runtime/ExecutionSession.cpp

@@ -4,7 +4,7 @@
 
 //===------- ExecutionSession.cpp - ExecutionSession Implementation -------===//
 //
-// Copyright 2019-2023 The IBM Research Authors.
+// Copyright 2019-2024 The IBM Research Authors.
 //
 // =============================================================================
 //

src/Runtime/python/CMakeLists.txt

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+
+file(GENERATE
+  OUTPUT ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}/PyRuntime.py
+  INPUT ${CMAKE_CURRENT_SOURCE_DIR}/PyRuntime.py
+  )
+
+file(GENERATE
+  OUTPUT ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}/PyCompileAndRuntime.py
+  INPUT ${CMAKE_CURRENT_SOURCE_DIR}/PyCompileAndRuntime.py
+  )
+
+# TODO: Remove pybind11::python_link_helper after cmake version bumped to 3.16+
+add_onnx_mlir_library(OMPyExecutionSessionBase
+  PyExecutionSessionBase.cpp
+
+  EXCLUDE_FROM_OM_LIBS
+
+  LINK_LIBS PUBLIC
+  OMExecutionSession
+  OMMlirUtilities
+  pybind11::embed
+  pybind11::python_link_helper
+  )
+if(MSVC)
+  target_link_libraries(OMPyExecutionSessionBase
+    PRIVATE pybind11::windows_extras
+  )
+endif()
+set_target_properties(OMPyExecutionSessionBase
+  PROPERTIES
+  POSITION_INDEPENDENT_CODE TRUE
+  )
+target_compile_definitions(OMPyExecutionSessionBase
+  PRIVATE
+  $<TARGET_PROPERTY:onnx,COMPILE_DEFINITIONS>
+  )
+target_include_directories(OMPyExecutionSessionBase
+  PRIVATE
+  $<TARGET_PROPERTY:onnx,INCLUDE_DIRECTORIES>
+  )
+
+# When running on ubi8 image, shared lib backend tests fail with
+# the following error:
+#
+#    [libprotobuf ERROR google/protobuf/descriptor_database.cc:641] File already exists in database: onnx/onnx-ml.proto
+#    [libprotobuf FATAL google/protobuf/descriptor.cc:1371] CHECK failed: GeneratedDatabase()->Add(encoded_file_descriptor, size):
+#    terminate called after throwing an instance of 'google::protobuf::FatalException'
+#      what():  CHECK failed: GeneratedDatabase()->Add(encoded_file_descriptor, size):
+#    Aborted (core dumped)
+#
+# This is because test.py loads (among others) the following
+# two .so shared libs:
+#
+#   - onnx_cpp2py_export.cpython-39-s390x-linux-gnu.so
+#     (import onnx)
+#   - PyRuntimeC.cpython-39-s390x-linux-gnu.so
+#     (from PyRuntimeC import OMExecutionSession)
+#
+# Both libs share the same libprotobuf.so when loaded by test.py.
+# However, they were both built with the same onnx-ml.pb.cc generated
+# from onnx-ml.proto and the protobuf runtime requires all compiled-in
+# .proto files have unique names. Hence the error.
+#
+# PyRuntimeC doesn't really need onnx beyond the onnx::TensorProto::*
+# types so we remove onnx from its target_link_libraries. But that
+# also removes some of the compile definitions and include directories
+# which we add back through target_compile_definitions and
+# target_include_directories.
+pybind11_add_module(PyRuntimeC PyExecutionSession.cpp)
+add_dependencies(PyRuntimeC onnx_proto)
+target_compile_options(PyRuntimeC
+  PRIVATE
+  $<$<OR:$<CXX_COMPILER_ID:Clang>,$<CXX_COMPILER_ID:AppleClang>,$<CXX_COMPILER_ID:GNU>>:-frtti -fexceptions>
+  $<$<CXX_COMPILER_ID:MSVC>:/EHsc /GR>
+  )
+target_compile_definitions(PyRuntimeC
+  PRIVATE
+  $<TARGET_PROPERTY:onnx,COMPILE_DEFINITIONS>
+  )
+target_include_directories(PyRuntimeC
+  PRIVATE
+  $<TARGET_PROPERTY:onnx,INCLUDE_DIRECTORIES>
+  )
+target_link_libraries(PyRuntimeC
+  PRIVATE
+  OMPyExecutionSessionBase
+  )
+llvm_update_compile_flags(PyRuntimeC)
+
+install(TARGETS PyRuntimeC
+  DESTINATION lib
+  )
+
+pybind11_add_module(PyCompileAndRuntimeC PyOMCompileExecutionSession.cpp)
+add_dependencies(PyCompileAndRuntimeC onnx_proto)
+target_compile_options(PyCompileAndRuntimeC
+  PRIVATE
+  $<$<OR:$<CXX_COMPILER_ID:Clang>,$<CXX_COMPILER_ID:AppleClang>,$<CXX_COMPILER_ID:GNU>>:-frtti -fexceptions>
+  $<$<CXX_COMPILER_ID:MSVC>:/EHsc /GR>
+  )
+target_compile_definitions(PyCompileAndRuntimeC
+  PRIVATE
+  $<TARGET_PROPERTY:onnx,COMPILE_DEFINITIONS>
+  )
+target_include_directories(PyCompileAndRuntimeC
+  PRIVATE
+  ${ONNX_MLIR_SRC_ROOT}/include
+  $<TARGET_PROPERTY:onnx,INCLUDE_DIRECTORIES>
+  )
+target_link_libraries(PyCompileAndRuntimeC
+  PRIVATE
+  OMCompiler
+  OMPyExecutionSessionBase
+  )
+llvm_update_compile_flags(PyCompileAndRuntimeC)
+
+install(TARGETS PyCompileAndRuntimeC
+  DESTINATION lib
+  )

src/Runtime/python/PyCompileAndRuntime.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python3
+
+############# PyCompileAndRuntime.py #######################################
+#
+# Copyright 2021-2024 The IBM Research Authors.
+#
+################################################################################
+# commom class `PyOMRuntime` called by python scripts
+################################################################################
+import numpy as np
+
+try:
+    from PyCompileAndRuntimeC import (
+        OMCompileExecutionSession as OMCompileExecutionSession_,
+    )
+except ImportError:
+    raise ImportError(
+        "Looks like you did not build the PyCompileAndRuntimeC target, build it by running `make PyCompileAndRuntimeC`."
+        "You may need to set ONNX_MLIR_HOME to `onnx-mlir/build/Debug` since `make PyCompileAndRuntimeC` outputs to `build/Debug` by default"
+    )
+
+
+class OMCompileExecutionSession(OMCompileExecutionSession_):
+    def run(self, inputs):
+        # Prepare arguments to call sess.run
+        pyrun_inputs = []
+        pyrun_shapes = []
+        pyrun_strides = []
+        for inp in inputs:
+            pyrun_inputs.append(inp.ravel())
+            pyrun_shapes.append(np.array(inp.shape, dtype=np.int64))
+            pyrun_strides.append(np.array(inp.strides, dtype=np.int64))
+        return super(OMCompileExecutionSession, self).run(
+            pyrun_inputs, pyrun_shapes, pyrun_strides
+        )

src/Runtime/PyExecutionSession.hpp → src/Runtime/python/PyExecutionSession.hpp

@@ -26,15 +26,16 @@ class PyExecutionSession : public onnx_mlir::PyExecutionSessionBase {
 };
 } // namespace onnx_mlir
 
-PYBIND11_MODULE(PyRuntime, m) {
+PYBIND11_MODULE(PyRuntimeC, m) {
   py::class_<onnx_mlir::PyExecutionSession>(m, "OMExecutionSession")
       .def(py::init<const std::string &, const std::string &, const bool>(),
           py::arg("shared_lib_path"), py::arg("tag") = "",
           py::arg("use_default_entry_point") = 1)
       .def("entry_points", &onnx_mlir::PyExecutionSession::pyQueryEntryPoints)
       .def("set_entry_point", &onnx_mlir::PyExecutionSession::pySetEntryPoint,
           py::arg("name"))
-      .def("run", &onnx_mlir::PyExecutionSession::pyRun, py::arg("input"))
+      .def("run", &onnx_mlir::PyExecutionSession::pyRun, py::arg("input"),
+          py::arg("shape"), py::arg("stride"))
       .def("input_signature", &onnx_mlir::PyExecutionSession::pyInputSignature)
       .def("output_signature",
           &onnx_mlir::PyExecutionSession::pyOutputSignature);