From 026ac752ba70f85402911305a972965f5b262900 Mon Sep 17 00:00:00 2001 From: KeDengMS Date: Tue, 25 Jun 2019 19:31:32 -0700 Subject: [PATCH 1/2] Implementation of TVM codegen library --- .gitmodules | 2 +- cgmanifest.json | 4 +- cmake/CMakeLists.txt | 8 +- cmake/external/tvm | 2 +- cmake/onnxruntime_codegen.cmake | 17 +- cmake/onnxruntime_unittests.cmake | 6 + onnxruntime/core/codegen/common/common.cc | 258 ++++++++++++++++++ onnxruntime/core/codegen/common/common.h | 151 ++++++++++ onnxruntime/core/codegen/common/creator.h | 76 ++++++ onnxruntime/core/codegen/common/dispatcher.h | 74 +++++ onnxruntime/core/codegen/common/dump_array.h | 62 +++++ onnxruntime/core/codegen/common/handle.h | 22 ++ onnxruntime/core/codegen/common/op_macro.h | 98 +++++++ onnxruntime/core/codegen/common/profile.h | 37 +++ onnxruntime/core/codegen/common/registry.h | 70 +++++ onnxruntime/core/codegen/common/settings.cc | 74 +++++ onnxruntime/core/codegen/common/settings.h | 38 +++ onnxruntime/core/codegen/common/target_info.h | 33 +++ onnxruntime/core/codegen/common/utils.cc | 54 ++++ onnxruntime/core/codegen/common/utils.h | 20 ++ onnxruntime/core/codegen/mti/common.h | 16 ++ .../core/codegen/mti/debug/tvm_print.cc | 83 ++++++ .../core/codegen/mti/debug/tvm_print.h | 19 ++ .../core/codegen/mti/math/binary_ops.cc | 70 +++++ .../core/codegen/mti/math/binary_ops.h | 42 +++ onnxruntime/core/codegen/mti/math/gemm.cc | 28 ++ onnxruntime/core/codegen/mti/math/gemm.h | 16 ++ .../core/codegen/mti/math/logsoftmax.cc | 18 ++ .../core/codegen/mti/math/logsoftmax.h | 11 + .../core/codegen/mti/math/matmul_ops.cc | 138 ++++++++++ .../core/codegen/mti/math/matmul_ops.h | 16 ++ .../core/codegen/mti/math/reduce_ops.cc | 90 ++++++ .../core/codegen/mti/math/reduce_ops.h | 72 +++++ onnxruntime/core/codegen/mti/math/softmax.cc | 18 ++ onnxruntime/core/codegen/mti/math/softmax.h | 11 + .../core/codegen/mti/math/unary_ops.cc | 142 ++++++++++ onnxruntime/core/codegen/mti/math/unary_ops.h | 36 +++ onnxruntime/core/codegen/mti/mti_tvm_utils.cc | 162 +++++++++++ onnxruntime/core/codegen/mti/mti_tvm_utils.h | 64 +++++ onnxruntime/core/codegen/mti/nn/conv_ops.cc | 199 ++++++++++++++ onnxruntime/core/codegen/mti/nn/conv_ops.h | 39 +++ onnxruntime/core/codegen/mti/nn/lstm.cc | 140 ++++++++++ onnxruntime/core/codegen/mti/nn/lstm.h | 35 +++ onnxruntime/core/codegen/mti/nn/pool_ops.cc | 58 ++++ onnxruntime/core/codegen/mti/nn/pool_ops.h | 32 +++ .../core/codegen/mti/tensor/cast_ops.cc | 37 +++ .../core/codegen/mti/tensor/cast_ops.h | 15 + .../core/codegen/mti/tensor/concat_ops.cc | 83 ++++++ .../core/codegen/mti/tensor/concat_ops.h | 15 + onnxruntime/core/codegen/mti/tensor/crop.cc | 58 ++++ onnxruntime/core/codegen/mti/tensor/crop.h | 17 ++ onnxruntime/core/codegen/mti/tensor/gather.cc | 53 ++++ onnxruntime/core/codegen/mti/tensor/gather.h | 17 ++ .../core/codegen/mti/tensor/pad_ops.cc | 121 ++++++++ onnxruntime/core/codegen/mti/tensor/pad_ops.h | 34 +++ .../core/codegen/mti/tensor/reshape_ops.cc | 48 ++++ .../core/codegen/mti/tensor/reshape_ops.h | 16 ++ onnxruntime/core/codegen/mti/tensor/slice.cc | 48 ++++ onnxruntime/core/codegen/mti/tensor/slice.h | 17 ++ onnxruntime/core/codegen/mti/tensor/split.cc | 72 +++++ onnxruntime/core/codegen/mti/tensor/split.h | 25 ++ onnxruntime/core/codegen/mti/tensor/tile.cc | 40 +++ onnxruntime/core/codegen/mti/tensor/tile.h | 16 ++ .../core/codegen/mti/tensor/transpose.cc | 16 ++ .../core/codegen/mti/tensor/transpose.h | 16 ++ onnxruntime/core/codegen/mti/tensor/where.cc | 36 +++ onnxruntime/core/codegen/mti/tensor/where.h | 17 ++ .../core/codegen/target/codegen_context.cc | 27 ++ .../core/codegen/target/codegen_context.h | 44 +++ .../target/generic/op_ir_creator/all_ops.h | 47 ++++ .../generic/op_ir_creator/math/binary_ops.cc | 46 ++++ .../target/generic/op_ir_creator/math/clip.cc | 31 +++ .../target/generic/op_ir_creator/math/gemm.cc | 39 +++ .../generic/op_ir_creator/math/logsoftmax.cc | 32 +++ .../generic/op_ir_creator/math/matmul.cc | 23 ++ .../math/quantize/matmul_integer.cc | 37 +++ .../generic/op_ir_creator/math/reduce_ops.cc | 111 ++++++++ .../generic/op_ir_creator/math/softmax.cc | 32 +++ .../generic/op_ir_creator/math/unary_ops.cc | 136 +++++++++ .../op_ir_creator/math/variadic_ops.cc | 36 +++ .../target/generic/op_ir_creator/nn/conv.cc | 131 +++++++++ .../target/generic/op_ir_creator/nn/lstm.cc | 64 +++++ .../generic/op_ir_creator/nn/pool_ops.cc | 93 +++++++ .../generic/op_ir_creator/tensor/cast.cc | 40 +++ .../generic/op_ir_creator/tensor/concat.cc | 30 ++ .../generic/op_ir_creator/tensor/crop.cc | 45 +++ .../generic/op_ir_creator/tensor/gather.cc | 30 ++ .../generic/op_ir_creator/tensor/pad.cc | 49 ++++ .../op_ir_creator/tensor/reshape_ops.cc | 96 +++++++ .../generic/op_ir_creator/tensor/slice.cc | 107 ++++++++ .../generic/op_ir_creator/tensor/split.cc | 65 +++++ .../generic/op_ir_creator/tensor/transpose.cc | 46 ++++ .../generic/op_ir_creator/tensor/where.cc | 28 ++ .../target/generic/scheduler/all_schedules.h | 20 ++ .../generic/scheduler/ort_type_schedule.cc | 22 ++ .../generic/scheduler/schedule_utils.cc | 164 +++++++++++ .../target/generic/scheduler/schedule_utils.h | 60 ++++ .../generic/scheduler/tvm_rule_schedule.cc | 41 +++ .../generic/weight_layout/transpose_2d.cc | 65 +++++ .../generic/weight_layout/transpose_2d.h | 33 +++ .../weight_layout/vertical_stripes_2d.cc | 77 ++++++ .../weight_layout/vertical_stripes_2d.h | 40 +++ .../core/codegen/target/ort_tvm_utils.cc | 186 +++++++++++++ .../core/codegen/target/ort_tvm_utils.h | 31 +++ .../core/codegen/target/tvm_ir_builder.cc | 125 +++++++++ .../core/codegen/target/tvm_ir_builder.h | 64 +++++ .../core/codegen/target/tvm_op_creator.cc | 37 +++ .../core/codegen/target/tvm_op_creator.h | 84 ++++++ .../codegen/target/tvm_schedule_builder.cc | 104 +++++++ .../codegen/target/tvm_schedule_builder.h | 46 ++++ .../core/codegen/target/tvm_scheduler.cc | 79 ++++++ .../core/codegen/target/tvm_scheduler.h | 128 +++++++++ .../core/codegen/target/weight_layout.cc | 92 +++++++ .../core/codegen/target/weight_layout.h | 68 +++++ onnxruntime/core/codegen/tvm/tvm_compiler.cc | 93 ------- onnxruntime/core/codegen/tvm/tvm_compiler.h | 36 --- onnxruntime/core/codegen/tvm/tvm_kernel.h | 126 --------- onnxruntime/core/codegen/tvm/tvm_utils.cc | 32 --- onnxruntime/core/codegen/tvm/tvm_utils.h | 17 -- onnxruntime/test/testdata/fuse_add_1.pb | Bin 407 -> 0 bytes onnxruntime/test/testdata/fuse_mul_1.pb | Bin 0 -> 169 bytes onnxruntime/test/tvm/tvm_basic_test.cc | 98 +++---- .../test/tvm/tvm_demo/demo_compiler.cc | 226 +++++++++++++++ onnxruntime/test/tvm/tvm_demo/demo_compiler.h | 31 +++ 124 files changed, 6794 insertions(+), 374 deletions(-) create mode 100644 onnxruntime/core/codegen/common/common.cc create mode 100644 onnxruntime/core/codegen/common/common.h create mode 100644 onnxruntime/core/codegen/common/creator.h create mode 100644 onnxruntime/core/codegen/common/dispatcher.h create mode 100644 onnxruntime/core/codegen/common/dump_array.h create mode 100644 onnxruntime/core/codegen/common/handle.h create mode 100644 onnxruntime/core/codegen/common/op_macro.h create mode 100644 onnxruntime/core/codegen/common/profile.h create mode 100644 onnxruntime/core/codegen/common/registry.h create mode 100644 onnxruntime/core/codegen/common/settings.cc create mode 100644 onnxruntime/core/codegen/common/settings.h create mode 100644 onnxruntime/core/codegen/common/target_info.h create mode 100644 onnxruntime/core/codegen/common/utils.cc create mode 100644 onnxruntime/core/codegen/common/utils.h create mode 100644 onnxruntime/core/codegen/mti/common.h create mode 100644 onnxruntime/core/codegen/mti/debug/tvm_print.cc create mode 100644 onnxruntime/core/codegen/mti/debug/tvm_print.h create mode 100644 onnxruntime/core/codegen/mti/math/binary_ops.cc create mode 100644 onnxruntime/core/codegen/mti/math/binary_ops.h create mode 100644 onnxruntime/core/codegen/mti/math/gemm.cc create mode 100644 onnxruntime/core/codegen/mti/math/gemm.h create mode 100644 onnxruntime/core/codegen/mti/math/logsoftmax.cc create mode 100644 onnxruntime/core/codegen/mti/math/logsoftmax.h create mode 100644 onnxruntime/core/codegen/mti/math/matmul_ops.cc create mode 100644 onnxruntime/core/codegen/mti/math/matmul_ops.h create mode 100644 onnxruntime/core/codegen/mti/math/reduce_ops.cc create mode 100644 onnxruntime/core/codegen/mti/math/reduce_ops.h create mode 100644 onnxruntime/core/codegen/mti/math/softmax.cc create mode 100644 onnxruntime/core/codegen/mti/math/softmax.h create mode 100644 onnxruntime/core/codegen/mti/math/unary_ops.cc create mode 100644 onnxruntime/core/codegen/mti/math/unary_ops.h create mode 100644 onnxruntime/core/codegen/mti/mti_tvm_utils.cc create mode 100644 onnxruntime/core/codegen/mti/mti_tvm_utils.h create mode 100644 onnxruntime/core/codegen/mti/nn/conv_ops.cc create mode 100644 onnxruntime/core/codegen/mti/nn/conv_ops.h create mode 100644 onnxruntime/core/codegen/mti/nn/lstm.cc create mode 100644 onnxruntime/core/codegen/mti/nn/lstm.h create mode 100644 onnxruntime/core/codegen/mti/nn/pool_ops.cc create mode 100644 onnxruntime/core/codegen/mti/nn/pool_ops.h create mode 100644 onnxruntime/core/codegen/mti/tensor/cast_ops.cc create mode 100644 onnxruntime/core/codegen/mti/tensor/cast_ops.h create mode 100644 onnxruntime/core/codegen/mti/tensor/concat_ops.cc create mode 100644 onnxruntime/core/codegen/mti/tensor/concat_ops.h create mode 100644 onnxruntime/core/codegen/mti/tensor/crop.cc create mode 100644 onnxruntime/core/codegen/mti/tensor/crop.h create mode 100644 onnxruntime/core/codegen/mti/tensor/gather.cc create mode 100644 onnxruntime/core/codegen/mti/tensor/gather.h create mode 100644 onnxruntime/core/codegen/mti/tensor/pad_ops.cc create mode 100644 onnxruntime/core/codegen/mti/tensor/pad_ops.h create mode 100644 onnxruntime/core/codegen/mti/tensor/reshape_ops.cc create mode 100644 onnxruntime/core/codegen/mti/tensor/reshape_ops.h create mode 100644 onnxruntime/core/codegen/mti/tensor/slice.cc create mode 100644 onnxruntime/core/codegen/mti/tensor/slice.h create mode 100644 onnxruntime/core/codegen/mti/tensor/split.cc create mode 100644 onnxruntime/core/codegen/mti/tensor/split.h create mode 100644 onnxruntime/core/codegen/mti/tensor/tile.cc create mode 100644 onnxruntime/core/codegen/mti/tensor/tile.h create mode 100644 onnxruntime/core/codegen/mti/tensor/transpose.cc create mode 100644 onnxruntime/core/codegen/mti/tensor/transpose.h create mode 100644 onnxruntime/core/codegen/mti/tensor/where.cc create mode 100644 onnxruntime/core/codegen/mti/tensor/where.h create mode 100644 onnxruntime/core/codegen/target/codegen_context.cc create mode 100644 onnxruntime/core/codegen/target/codegen_context.h create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/all_ops.h create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/math/binary_ops.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/math/clip.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/math/gemm.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/math/logsoftmax.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/math/matmul.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/math/quantize/matmul_integer.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/math/reduce_ops.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/math/softmax.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/math/unary_ops.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/math/variadic_ops.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/nn/conv.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/nn/lstm.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/nn/pool_ops.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/cast.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/concat.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/crop.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/gather.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/pad.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/reshape_ops.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/slice.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/split.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/transpose.cc create mode 100644 onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/where.cc create mode 100644 onnxruntime/core/codegen/target/generic/scheduler/all_schedules.h create mode 100644 onnxruntime/core/codegen/target/generic/scheduler/ort_type_schedule.cc create mode 100644 onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.cc create mode 100644 onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.h create mode 100644 onnxruntime/core/codegen/target/generic/scheduler/tvm_rule_schedule.cc create mode 100644 onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.cc create mode 100644 onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.h create mode 100644 onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.cc create mode 100644 onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.h create mode 100644 onnxruntime/core/codegen/target/ort_tvm_utils.cc create mode 100644 onnxruntime/core/codegen/target/ort_tvm_utils.h create mode 100644 onnxruntime/core/codegen/target/tvm_ir_builder.cc create mode 100644 onnxruntime/core/codegen/target/tvm_ir_builder.h create mode 100644 onnxruntime/core/codegen/target/tvm_op_creator.cc create mode 100644 onnxruntime/core/codegen/target/tvm_op_creator.h create mode 100644 onnxruntime/core/codegen/target/tvm_schedule_builder.cc create mode 100644 onnxruntime/core/codegen/target/tvm_schedule_builder.h create mode 100644 onnxruntime/core/codegen/target/tvm_scheduler.cc create mode 100644 onnxruntime/core/codegen/target/tvm_scheduler.h create mode 100644 onnxruntime/core/codegen/target/weight_layout.cc create mode 100644 onnxruntime/core/codegen/target/weight_layout.h delete mode 100644 onnxruntime/core/codegen/tvm/tvm_compiler.cc delete mode 100644 onnxruntime/core/codegen/tvm/tvm_compiler.h delete mode 100644 onnxruntime/core/codegen/tvm/tvm_kernel.h delete mode 100644 onnxruntime/core/codegen/tvm/tvm_utils.cc delete mode 100644 onnxruntime/core/codegen/tvm/tvm_utils.h delete mode 100644 onnxruntime/test/testdata/fuse_add_1.pb create mode 100644 onnxruntime/test/testdata/fuse_mul_1.pb create mode 100644 onnxruntime/test/tvm/tvm_demo/demo_compiler.cc create mode 100644 onnxruntime/test/tvm/tvm_demo/demo_compiler.h diff --git a/.gitmodules b/.gitmodules index bf0b0903613cf..100ba7d91b7cf 100644 --- a/.gitmodules +++ b/.gitmodules @@ -9,7 +9,7 @@ url = https://github.com/onnx/onnx [submodule "cmake/external/tvm"] path = cmake/external/tvm - url = https://github.com/dmlc/tvm.git + url = https://github.com/microsoft/onnxruntime-tvm.git [submodule "cmake/external/date"] path = cmake/external/date url = https://github.com/HowardHinnant/date.git diff --git a/cgmanifest.json b/cgmanifest.json index 05a503e03c993..9b3d4b9d8c475 100644 --- a/cgmanifest.json +++ b/cgmanifest.json @@ -103,8 +103,8 @@ "component":{ "type":"git", "git":{ - "commitHash":"c2b36154778503a509a70a3b5309b201969eccab", - "repositoryUrl":"https://github.com/dmlc/tvm.git" + "commitHash":"fd4801612817f96e890058656834deb925fc064a", + "repositoryUrl":"https://github.com/microsoft/onnxruntime-tvm.git" } } }, diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 9762200a58557..d00f78e2d47da 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -66,7 +66,7 @@ option(onnxruntime_USE_LLVM "Build tvm with LLVM" OFF) option(onnxruntime_USE_OPENMP "Build with OpenMP support" OFF) option(onnxruntime_BUILD_SHARED_LIB "Build a shared library" OFF) option(onnxruntime_ENABLE_MICROSOFT_INTERNAL "Use this option to enable/disable microsoft internal only code" OFF) -option(onnxruntime_USE_NUPHAR "Build with Nupha" OFF) +option(onnxruntime_USE_NUPHAR "Build with Nuphar" OFF) option(onnxruntime_USE_BRAINSLICE "Build with BrainSlice" OFF) option(onnxruntime_USE_TENSORRT "Build with TensorRT support" OFF) option(onnxruntime_ENABLE_LTO "Enable link time optimization" ON) @@ -426,10 +426,8 @@ else() string(APPEND CMAKE_C_FLAGS " -Wall -Wextra -ffunction-sections -fdata-sections") if(onnxruntime_DEV_MODE) - if(NOT onnxruntime_USE_TVM) - string(APPEND CMAKE_CXX_FLAGS " -Werror") - string(APPEND CMAKE_C_FLAGS " -Werror") - endif() + string(APPEND CMAKE_CXX_FLAGS " -Werror") + string(APPEND CMAKE_C_FLAGS " -Werror") endif() check_cxx_compiler_flag(-Wunused-but-set-variable HAS_UNUSED_BUT_SET_VARIABLE) check_cxx_compiler_flag(-Wunused-parameter HAS_UNUSED_PARAMETER) diff --git a/cmake/external/tvm b/cmake/external/tvm index c2b3615477850..fd4801612817f 160000 --- a/cmake/external/tvm +++ b/cmake/external/tvm @@ -1 +1 @@ -Subproject commit c2b36154778503a509a70a3b5309b201969eccab +Subproject commit fd4801612817f96e890058656834deb925fc064a diff --git a/cmake/onnxruntime_codegen.cmake b/cmake/onnxruntime_codegen.cmake index d63d367f2bb9a..56ce2f416db97 100644 --- a/cmake/onnxruntime_codegen.cmake +++ b/cmake/onnxruntime_codegen.cmake @@ -1,15 +1,24 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. +file(GLOB_RECURSE onnxruntime_codegen_common_srcs + "${ONNXRUNTIME_ROOT}/core/codegen/common/*.h" + "${ONNXRUNTIME_ROOT}/core/codegen/common/*.cc" +) + file(GLOB_RECURSE onnxruntime_codegen_tvm_srcs CONFIGURE_DEPENDS - "${ONNXRUNTIME_ROOT}/core/codegen/tvm/*.h" - "${ONNXRUNTIME_ROOT}/core/codegen/tvm/*.cc" + "${ONNXRUNTIME_ROOT}/core/codegen/mti/*.h" + "${ONNXRUNTIME_ROOT}/core/codegen/mti/*.cc" + "${ONNXRUNTIME_ROOT}/core/codegen/target/*.h" + "${ONNXRUNTIME_ROOT}/core/codegen/target/*.cc" ) +source_group(TREE ${ONNXRUNTIME_ROOT}/core FILES ${onnxruntime_codegen_common_srcs} ${onnxruntime_codegen_tvm_srcs}) + #onnxruntime_codegen_tvm depends on onnxruntime framework -add_library(onnxruntime_codegen_tvm ${onnxruntime_codegen_tvm_srcs}) +add_library(onnxruntime_codegen_tvm ${onnxruntime_codegen_common_srcs} ${onnxruntime_codegen_tvm_srcs}) set_target_properties(onnxruntime_codegen_tvm PROPERTIES FOLDER "ONNXRuntime") -target_include_directories(onnxruntime_codegen_tvm PRIVATE ${ONNXRUNTIME_ROOT} ${TVM_INCLUDES}) +target_include_directories(onnxruntime_codegen_tvm PRIVATE ${ONNXRUNTIME_ROOT} ${TVM_INCLUDES} ${MKLML_INCLUDE_DIR} ${eigen_INCLUDE_DIRS}) onnxruntime_add_include_to_target(onnxruntime_codegen_tvm onnxruntime_common onnxruntime_framework gsl onnx onnx_proto protobuf::libprotobuf) target_compile_options(onnxruntime_codegen_tvm PRIVATE ${DISABLED_WARNINGS_FOR_TVM}) # need onnx to build to create headers that this project includes diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index ce3c2c905ba27..d8276c16582d7 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -408,6 +408,12 @@ if(WIN32) $ ) endif() + if (onnxruntime_USE_TVM) + add_custom_command( + TARGET ${test_data_target} POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy $ $ + ) + endif() endif() add_library(onnx_test_data_proto ${TEST_SRC_DIR}/proto/tml.proto) diff --git a/onnxruntime/core/codegen/common/common.cc b/onnxruntime/core/codegen/common/common.cc new file mode 100644 index 0000000000000..757c1677dd2e5 --- /dev/null +++ b/onnxruntime/core/codegen/common/common.cc @@ -0,0 +1,258 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/common/common.h" + +#include "core/framework/tensorprotoutils.h" +#include "core/graph/graph.h" +#include "core/graph/schema_registry.h" +#include +#include + +namespace onnxruntime { + +NodeKey GetKey(const onnxruntime::Node* node) { + ORT_ENFORCE(nullptr != node); + ORT_ENFORCE(node->OutputDefs().size() > 0); + return node->OutputDefs()[0]->Name(); +} + +NodeKey GetKey(const onnxruntime::Node& node) { + ORT_ENFORCE(node.OutputDefs().size() > 0); + return node.OutputDefs()[0]->Name(); +} + +NodeKey GetKey(const onnxruntime::NodeArg* def) { + // NodeArg's name is unique. + ORT_ENFORCE(nullptr != def); + return def->Name(); +} + +bool IsRecurrentNode(const onnxruntime::Node& node) { + auto op_type = node.OpType(); + return (op_type == "LSTM" || op_type == "RNN" || op_type == "GRU" || + op_type == "Scan" || op_type == "Loop"); +} + +bool IsAliasNode(const onnxruntime::Node& node) { + auto op_type = node.OpType(); + return (op_type == "Flatten" || op_type == "Identity" || op_type == "Reshape" || + op_type == "Squeeze" || op_type == "Unsqueeze"); +} + +std::string NormalizeCppName(const std::string& name) { + std::string normalized_name = name; + for (char c : {'.', ' ', '+', '-', '*', '/', '\\', '='}) + std::replace(normalized_name.begin(), normalized_name.end(), c, '_'); + return normalized_name; +} + +std::string NormalizeNodeArgName(const NodeArg* def) { + return NormalizeCppName(def->Name()); +} + +bool IsFusedNode(const Node& node) { + if (node.NodeType() == Node::Type::Fused) { + return true; + } + return false; +} + +// A unified API to get Subgraph +const Graph* GetSubgraph(const Node& node) { + if (node.NodeType() == Node::Type::Fused) { + return &(node.GetFunctionBody()->Body()); + } else if (node.OpType() == "Scan") { + return node.GetGraphAttribute("body"); + } + // return nullptr implying no subgraph + return nullptr; +} + +bool HasLoop(const Node& node) { + auto op_type = node.OpType(); + if (op_type == "LSTM" || + op_type == "GRU" || + op_type == "RNN" || + op_type == "Scan") { + return true; + } + return false; +} + +// Return the corresponding input node for the NodeArg of the given node +const onnxruntime::Node* GetInputNode(const Node& node, const NodeArg* def) { + const auto& input_name = def->Name(); + const onnxruntime::Node* input_node = nullptr; + // search input node set to see if input_name is in their outputs (weights are not from node) + for (auto iter = node.InputNodesBegin(); iter != node.InputNodesEnd(); ++iter) { + const onnxruntime::Node& p = *iter; + bool found = false; + p.ForEachWithIndex( + p.OutputDefs(), + [&found, &input_name](const onnxruntime::NodeArg& out_def, size_t) { + if (input_name == out_def.Name()) { + found = true; + } + return Status::OK(); + }); + if (found) + input_node = &p; + } + return input_node; +} + +// create capacity from subgraph +std::unique_ptr ToCapacity(const onnxruntime::GraphViewer& graph, + std::unique_ptr& subgraph) { + auto meta_def = std::make_unique<::onnxruntime::IndexedSubGraph::MetaDef>(); + static int fuse_count = 0; + meta_def->name = "Fuse" + std::to_string(fuse_count++); + meta_def->domain = "Fuse"; + + std::set node_indices(subgraph->nodes.begin(), subgraph->nodes.end()); + + const auto& start_node_index = subgraph->nodes.front(); + const auto& start_node = *graph.GetNode(start_node_index); + const auto& end_node_index = subgraph->nodes.back(); + const auto& end_node = *graph.GetNode(end_node_index); + meta_def->name += start_node.OpType() + std::to_string(start_node_index); + meta_def->name += "_With" + std::to_string(subgraph->nodes.size()) + "Nodes_"; + meta_def->name += end_node.OpType() + std::to_string(end_node_index); + + for (const auto& node_index : subgraph->nodes) { + const auto& node = *graph.GetNode(node_index); + // handle current graph's inputs + node.ForEachWithIndex( + node.InputDefs(), + [&meta_def, &node, &node_indices](const onnxruntime::NodeArg& def, size_t) { + const onnxruntime::Node* input_node = GetInputNode(node, &def); + bool input_from_subgraph = (input_node && node_indices.count(input_node->Index())); + if (!input_from_subgraph) { + // input is from weights or outside of graph + meta_def->inputs.push_back(def.Name()); + } + return Status::OK(); + }); + + // Handle outouts + // two cases are considerd as outputs + // 1. Output NodeArg is not used by any Node + // 2. Output NodeArg is used by at least one Node out of this subgraph. + // Note a NodeArg can be used by Nodes in and out of the subgraph at the same time. + + auto InsertOutputToSubgraph = [&meta_def](const NodeArg* def) { + if (std::find(meta_def->outputs.begin(), meta_def->outputs.end(), def->Name()) == + meta_def->outputs.end()) { + meta_def->outputs.push_back(def->Name()); + } + }; + + std::unordered_set input_names_from_the_output_node; + + for (auto o_iter = node.OutputEdgesBegin(); o_iter != node.OutputEdgesEnd(); ++o_iter) { + const auto& p = *o_iter; + const Node& out_node = p.GetNode(); + + // preprocess for the case 1 + out_node.ForEachWithIndex( + out_node.InputDefs(), + [&input_names_from_the_output_node](const onnxruntime::NodeArg& in_def, size_t) { + input_names_from_the_output_node.insert(in_def.Name()); + return Status::OK(); + }); + + // handle the case 2 + if (node_indices.count(out_node.Index()) == 0) { + const NodeArg* def = node.OutputDefs()[p.GetSrcArgIndex()]; + InsertOutputToSubgraph(def); + } + } + + // handle case 1 + node.ForEachWithIndex( + node.OutputDefs(), + [&](const onnxruntime::NodeArg& def, size_t) { + if (input_names_from_the_output_node.count(def.Name()) == 0) { + InsertOutputToSubgraph(&def); + } + return Status::OK(); + }); + } + + // Handle subgraph's initializers + const auto& all_initializers = graph.GetAllInitializedTensors(); + for (const auto& node_index : subgraph->nodes) { + const auto& node = *graph.GetNode(node_index); + // check whether it is an immediate nested subgraph + auto immediate_nested_subgraph = GetSubgraph(node); + // If so, copy the immediate nested subgraph's initializers to meta_def->inputs. + // Note we don't need recursion here, since Ort did recursion for us by handling subgraph early than the current graph. + // Therefore, the all inner nested subgraph's initializers should be already in the immediate nested subgraph's inputs. + if (nullptr != immediate_nested_subgraph) { + for (auto& n : immediate_nested_subgraph->Nodes()) { + n.ForEachWithIndex( + n.InputDefs(), + [&meta_def, &all_initializers](const onnxruntime::NodeArg& def, size_t) { + auto iter = all_initializers.find(def.Name()); + if (iter != all_initializers.end()) { + meta_def->inputs.push_back(def.Name()); + } + return Status::OK(); + }); + } + } + } + + meta_def->since_version = 1; + meta_def->status = ONNX_NAMESPACE::EXPERIMENTAL; + std::unique_ptr finished_subgraph(subgraph.release()); + finished_subgraph->SetMetaDef(meta_def); + return std::make_unique(std::move(finished_subgraph)); +} + +int64_t ShapeRank(const NodeArg* def) { + ORT_ENFORCE_DEBUG(nullptr != def); + return gsl::narrow_cast(def->Shape()->dim_size()); +} + +bool ShapeHasValue(const NodeArg* def, int i) { + ORT_ENFORCE_DEBUG(nullptr != def); + ORT_ENFORCE_DEBUG(i >= 0); + ORT_ENFORCE_DEBUG(i < def->Shape()->dim_size()); + return def->Shape()->dim(i).has_dim_value(); +} + +bool ShapeHasSymbol(const NodeArg* def, int i) { + ORT_ENFORCE_DEBUG(nullptr != def); + ORT_ENFORCE_DEBUG(i >= 0); + ORT_ENFORCE_DEBUG(i < def->Shape()->dim_size()); + return def->Shape()->dim(i).has_dim_param(); +} + +int64_t ShapeValue(const NodeArg* def, int i) { + ORT_ENFORCE_DEBUG(ShapeHasValue(def, i)); + return def->Shape()->dim(i).dim_value(); +} + +const std::string& ShapeSymbol(const NodeArg* def, int i) { + ORT_ENFORCE_DEBUG(ShapeHasSymbol(def, i)); + return def->Shape()->dim(i).dim_param(); +} + +ONNX_NAMESPACE::TensorProto_DataType TensorProtoDataType(const NodeArg* def) { + ORT_ENFORCE_DEBUG(nullptr != def); + return static_cast(def->TypeAsProto()->tensor_type().elem_type()); +} + +// Convert GraphNodes to internal NodePtrs without check lifetime. +// Please use it only locally when GraphNodes still exist +std::vector ConvertGraphNodesToNodePtrs(const GraphNodes& graph_nodes) { + std::vector nodes; + for (auto& node : graph_nodes) { + nodes.push_back(&node); + } + return nodes; +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/common/common.h b/onnxruntime/core/codegen/common/common.h new file mode 100644 index 0000000000000..11ad05325a381 --- /dev/null +++ b/onnxruntime/core/codegen/common/common.h @@ -0,0 +1,151 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/framework/compute_capability.h" +#include "core/framework/tensor.h" +#include "core/graph/graph_nodes.h" +#include "core/graph/graph_viewer.h" + +#ifndef NDEBUG +#define ORT_ENFORCE_DEBUG(...) ORT_ENFORCE(__VA_ARGS__) +#else +#define ORT_ENFORCE_DEBUG(...) +#endif // !NDEBUG + +// DYN_PROMOTE is a simplified llvm::dyn_cast, which does not need RTTI +// DYN_PROMOTE is faster than dynamic_cast and also has smaller binary size +// Please use DYN_PROMOTE in a critical path. +#define DYN_PROMOTE(BASE) \ + template \ + inline const ToType* Promote(const BASE* base) { \ + if (ToType::IsType(base)) \ + return static_cast(base); \ + return nullptr; \ + } \ + \ + template \ + inline ToType* Promote(BASE* base) { \ + if (ToType::IsType(base)) \ + return static_cast(base); \ + return nullptr; \ + } \ + \ + template \ + inline ToType* Promote(const std::unique_ptr& base) { \ + if (ToType::IsType(base.get())) \ + return static_cast(base); \ + return nullptr; \ + } \ + \ + template \ + inline ToType* Promote(const std::shared_ptr& base) { \ + if (ToType::IsType(base.get())) \ + return static_cast(base); \ + return nullptr; \ + } + +// DYN_PROMOTE_BASE is a macro inserted in the base class to support DYN_PROMOTE +// TYPE_ID is required for DYN_PROMOTE and TYPE_ID is a enum class +// TYPE_ID_VAR is a corresponding variable name for in the base class +#define DYN_PROMOTE_BASE(BASE, TYPE_ID, TYPE_ID_VAR) \ + inline const TYPE_ID TypeID() const { \ + return TYPE_ID_VAR; \ + } \ + \ + static inline bool IsType(const BASE*) { \ + return true; \ + } + +// DYN_PROMOTE_DERIVED is a macro inserted in a derived class to support DYN_PROMOTE +// TYPE_ID is required for DYN_PROMOTE and TYPE_ID is a enum class +// TYPE_ID_VALUE is corresponding TYPE_ID::value of a derived class. +#define DYN_PROMOTE_DERIVED(BASE, TYPE_ID, TYPE_ID_VALUE) \ + static inline bool IsType(const BASE* base) { \ + ORT_ENFORCE_DEBUG(nullptr != base); \ + return base->TypeID() == TYPE_ID::TYPE_ID_VALUE; \ + } + +// DYNAMIC_PROMOTE is a dynamic_cast needing RTTI +// DYNAMIC_PROMOTE is usually slower than than DYN_PROMOTE. +// Please use DYNAMIC_PROMOTE in a non-critical path. +#define DYNAMIC_PROMOTE(BASE) \ + template \ + inline const X* Promote(const BASE* base) { \ + auto derived = dynamic_cast(base); \ + ORT_ENFORCE(nullptr != derived); \ + return derived; \ + } \ + \ + template \ + inline X* Promote(BASE* base) { \ + auto derived = dynamic_cast(base); \ + ORT_ENFORCE(nullptr != derived); \ + return derived; \ + } \ + \ + template \ + inline X* Promote(const std::unique_ptr& base) { \ + auto derived = dynamic_cast(base.get()); \ + ORT_ENFORCE(nullptr != derived); \ + return derived; \ + } \ + \ + template \ + inline X* Promote(const std::shared_ptr& base) { \ + auto derived = dynamic_cast(base.get()); \ + ORT_ENFORCE(nullptr != derived); \ + return derived; \ + } + +namespace onnxruntime { + +// Nodekey is used as a key for maps +using NodeKey = std::string; + +NodeKey GetKey(const onnxruntime::Node* node); +NodeKey GetKey(const onnxruntime::Node& node); +NodeKey GetKey(const onnxruntime::NodeArg* def); + +bool IsRecurrentNode(const onnxruntime::Node& node); + +bool IsAliasNode(const onnxruntime::Node& node); + +// Helper function that creates ComputeCapability for subgraphs +std::unique_ptr ToCapacity(const onnxruntime::GraphViewer& graph, + std::unique_ptr& subgraph); + +bool IsFusedNode(const Node& node); + +bool HasLoop(const Node& node); + +const Graph* GetSubgraph(const Node& node); + +std::string NormalizeCppName(const std::string& name); + +std::string NormalizeNodeArgName(const NodeArg* def); + +// Return the corresponding input node for the NodeArg of the given node +const onnxruntime::Node* GetInputNode(const Node& node, const NodeArg* def); + +int64_t ShapeRank(const NodeArg* def); + +bool ShapeHasValue(const NodeArg* def, int i); + +bool ShapeHasSymbol(const NodeArg* def, int i); + +int64_t ShapeValue(const NodeArg* def, int i); + +const std::string& ShapeSymbol(const NodeArg* def, int i); + +ONNX_NAMESPACE::TensorProto_DataType TensorProtoDataType(const NodeArg* def); + +// Convert GraphNodes to internal NodePtrs without check lifetime. +// Please use it only locally when GraphNodes still exist +std::vector ConvertGraphNodesToNodePtrs(const GraphNodes& graph_nodes); + +enum : int { + Dimension_Unknown = -1, +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/common/creator.h b/onnxruntime/core/codegen/common/creator.h new file mode 100644 index 0000000000000..d15e86b5a481f --- /dev/null +++ b/onnxruntime/core/codegen/common/creator.h @@ -0,0 +1,76 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/codegen/common/dispatcher.h" + +// TODO rename this file to creator_base +namespace onnxruntime { +namespace codegen { + +// It is a base class for TVM Op IR builder, weight layout builder, TVM scheduler +// CreatorBase is a template class of compiler pass +// for 1) TVM IR builder +// 2) Weight layout transformer +// 3) TVM Scheduler, etc. +// CreatorBase is similor to OpXXCreate in llvm IR builder + +template +class CreatorBase { + public: + CreatorBase(const std::string& name) + : name_(name) {} + + ~CreatorBase() = default; + + virtual RETURN_TYPE Evaluate(INPUT_TYPE, + NODE_TYPE, + CONTEXT_TYPE, + OUTPUT_TYPE) = 0; + + const std::string& Name() const { + return name_; + } + + protected: + std::string name_; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(CreatorBase); +}; + +// macro to stringize +#define STRINGIZE_NX(OP) #OP +#define STRINGIZE(OP) STRINGIZE_NX(OP) + +// macro returns class name +#define CREATOR_CLASS(OP, POSTFIX) \ + OP##POSTFIX + +// macro returns class name as string +#define CREATOR_STRING(OP, POSTFIX) \ + STRINGIZE(CREATOR_CLASS(OP, POSTFIX)) + +// macro returns class constructor name +#define CREATOR_CLASS_FUNC(OP, POSTFIX) \ + OP##POSTFIX() + +// macro declares a creator class inheriting the template class CreatorBase +// with corresponding template parameters +#define DECLARE_CREATOR_CLASS(OP, POSTFIX, INPUT, NODE, CONTEXT, OUTPUT, RETURN) \ + class CREATOR_CLASS(OP, POSTFIX) : public onnxruntime::codegen::CreatorBase { \ + public: \ + CREATOR_CLASS_FUNC(OP, POSTFIX) : CreatorBase(CREATOR_STRING(OP, POSTFIX)) {} \ + RETURN Evaluate(INPUT, \ + NODE, \ + CONTEXT, \ + OUTPUT) override; \ + \ + private: \ + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(CREATOR_CLASS(OP, POSTFIX)); \ + }; + +} // namespace codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/common/dispatcher.h b/onnxruntime/core/codegen/common/dispatcher.h new file mode 100644 index 0000000000000..b4313cecad3a8 --- /dev/null +++ b/onnxruntime/core/codegen/common/dispatcher.h @@ -0,0 +1,74 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/common/common.h" +#include +#include +#include + +namespace onnxruntime { +namespace codegen { + +// DispatcherBase is a customized unordered_map +// that provides all codegen-related functionality +// including 1) dispatching a pass +// 2) dump corresponding name +// DispatcherBase may or may not keep ownership, +// depending on the template parameter, CONTENT_TYPE. + +template +class DispatcherBase { + public: + DispatcherBase(const std::string& name) + : name_(name) {} + + const std::string& Name() const { + return name_; + } + + bool Contains(const std::string& name) const { + return contents_.count(name) > 0; + } + + void ForEach(std::function + func) { + for (auto& p : contents_) { + func(p.first, p.second); + } + } + + bool Register(const std::string& name, + CONTENT_TYPE op) { + if (!Contains(name)) { + contents_.emplace(name, op); + return true; + } + return false; + } + + CONTENT_TYPE Get(const std::string& key) const { + auto iter = contents_.find(key); + if (iter != contents_.end()) { + return iter->second; + } + return nullptr; + } + + const std::unordered_map GetContents() const { + return contents_; + } + + std::unordered_map GetMutableContents() { + return contents_; + } + + protected: + std::string name_; + std::unordered_map contents_; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(DispatcherBase); +}; + +} // namespace codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/common/dump_array.h b/onnxruntime/core/codegen/common/dump_array.h new file mode 100644 index 0000000000000..8e51cd36d0087 --- /dev/null +++ b/onnxruntime/core/codegen/common/dump_array.h @@ -0,0 +1,62 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include +#include +#include +#include + +namespace onnxruntime { + +template +void DumpArrayRecursive(const T1* data, int64_t& data_offset, const std::vector& shape, int idx) { + int dim = static_cast(shape.size()); + if (dim == 0) { + std::cout << "[]\n"; + return; + } + + assert(idx < dim); + int sz = shape[idx]; + + std::cout << "["; + if (idx < dim - 1) { + for (auto i = 0; i < sz; ++i) { + DumpArrayRecursive(data, data_offset, shape, idx + 1); + if (i < sz - 1) { + std::cout << ","; + // print multiple newlines after ',' when necessary + for (int j = idx + 1; j < dim; j++) + std::cout << "\n"; + // print leading spaces before "[" when necessary + for (int j = 0; j < idx + 1; ++j) + std::cout << " "; + } + } + } else { + for (auto i = 0; i < sz; ++i) { + if (std::is_same::value || std::is_same::value) + std::cout << std::setw(3) << static_cast(*(data + data_offset)); + else + std::cout << std::setw(12) << std::setprecision(8) << *(data + data_offset); + data_offset++; + if (i < sz - 1) + std::cout << ","; + } + } + std::cout << "]"; +} + +// A helper function to dump multidimensional arrays in a way similar to numpy +template +void DumpArray(const std::string& tag, const T1* data, const std::vector& shape) { + std::cout << tag << "\n"; + int64_t data_offset = 0; + DumpArrayRecursive(data, data_offset, shape, 0); + assert(data_offset == TotalSize(shape)); + std::cout << std::endl; +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/common/handle.h b/onnxruntime/core/codegen/common/handle.h new file mode 100644 index 0000000000000..7caad27dcbe01 --- /dev/null +++ b/onnxruntime/core/codegen/common/handle.h @@ -0,0 +1,22 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/codegen/common/target_info.h" +#include +#include + +namespace onnxruntime { +namespace codegen { + +using DomainVersionLookupFunc = std::function; + +struct CodeGenHandle { + CodeGenTarget* codegen_target; + DomainVersionLookupFunc domain_version_lookup_func = + // by default, always uses the latest opset implemented + [](const std::string&) { return INT_MAX; }; +}; + +} // namespace codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/common/op_macro.h b/onnxruntime/core/codegen/common/op_macro.h new file mode 100644 index 0000000000000..91a0e803e521a --- /dev/null +++ b/onnxruntime/core/codegen/common/op_macro.h @@ -0,0 +1,98 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +namespace onnxruntime { + +#define LIST_BINARY_OPS() \ + BINARY_OP(Add) \ + BINARY_OP(Div) \ + BINARY_OP(Mul) \ + BINARY_OP(PRelu) \ + BINARY_OP(Sub) + +#define LIST_BINARY_CMP_OPS() \ + BINARY_CMP_OP(Equal) \ + BINARY_CMP_OP(Greater) \ + BINARY_CMP_OP(Less) + +#define LIST_POOL_OPS() \ + POOL_OP(MaxPool) \ + POOL_OP(AveragePool) \ + POOL_OP(GlobalMaxPool) \ + POOL_OP(GlobalAveragePool) + +#define LIST_REDUCE_OPS() \ + REDUCE_INDEXED_OP(ArgMax) \ + REDUCE_INDEXED_OP(ArgMin) \ + REDUCE_OP(ReduceL1) \ + REDUCE_OP(ReduceL2) \ + REDUCE_OP(ReduceLogSum) \ + REDUCE_OP(ReduceLogSumExp) \ + REDUCE_OP(ReduceMax) \ + REDUCE_OP(ReduceMean) \ + REDUCE_OP(ReduceMin) \ + REDUCE_OP(ReduceProd) \ + REDUCE_OP(ReduceSum) \ + REDUCE_OP(ReduceSumSquare) + +#define LIST_UNARY_OPS() \ + UNARY_OP(Abs) \ + UNARY_OP(Affine) \ + UNARY_OP(Ceil) \ + UNARY_OP(Elu) \ + UNARY_OP(Exp) \ + UNARY_OP(Floor) \ + UNARY_OP(HardSigmoid) \ + UNARY_OP(LeakyRelu) \ + UNARY_OP(Log) \ + UNARY_OP(Neg) \ + UNARY_OP(ParametricSoftplus) \ + UNARY_OP(Reciprocal) \ + UNARY_OP(Relu) \ + UNARY_OP(ScaledTanh) \ + UNARY_OP(Selu) \ + UNARY_OP(Sigmoid) \ + UNARY_OP(Softplus) \ + UNARY_OP(Softsign) \ + UNARY_OP(Sqrt) \ + UNARY_OP(Tanh) \ + UNARY_OP(ThresholdedRelu) + +#define LIST_VARIADIC_OPS() \ + VARIADIC_OP(Max) \ + VARIADIC_OP(Min) \ + VARIADIC_OP(Sum) + +#define LIST_ALL_GENERIC_OPS() \ + LIST_BINARY_OPS() \ + LIST_BINARY_CMP_OPS() \ + LIST_REDUCE_OPS() \ + LIST_POOL_OPS() \ + LIST_UNARY_OPS() \ + LIST_VARIADIC_OPS() \ + ADD_OP_ITEM(Cast) \ + ADD_OP_ITEM(Clip) \ + ADD_OP_ITEM(Concat) \ + ADD_OP_ITEM(Conv) \ + ADD_OP_ITEM(Crop) \ + ADD_OP_ITEM(Dropout) \ + ADD_OP_ITEM(Flatten) \ + ADD_OP_ITEM(Gather) \ + ADD_OP_ITEM(Gemm) \ + ADD_OP_ITEM(Identity) \ + ADD_OP_ITEM(LogSoftmax) \ + ADD_OP_ITEM(LSTM) \ + ADD_OP_ITEM(MatMul) \ + ADD_OP_ITEM(MatMulInteger) \ + ADD_OP_ITEM(Pad) \ + ADD_OP_ITEM(Reshape) \ + ADD_OP_ITEM(Slice) \ + ADD_OP_ITEM(Softmax) \ + ADD_OP_ITEM(Split) \ + ADD_OP_ITEM(Squeeze) \ + ADD_OP_ITEM(Transpose) \ + ADD_OP_ITEM(Unsqueeze) \ + ADD_OP_ITEM(Where) + +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/common/profile.h b/onnxruntime/core/codegen/common/profile.h new file mode 100644 index 0000000000000..642ae83db723b --- /dev/null +++ b/onnxruntime/core/codegen/common/profile.h @@ -0,0 +1,37 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +// uncomment this line or use -DCODEGEN_ENABLE_PROFILER in compiler options to enable profiler events in codegen +//#define CODEGEN_ENABLE_PROFILER + +#ifdef CODEGEN_ENABLE_PROFILER +#include "core/common/profiler.h" + +namespace onnxruntime { + +class ProfilerEvent { + public: + ProfilerEvent(const std::string& name) : name_(name) { + ts_ = profiling::Profiler::Instance().StartTime(); + } + + ~ProfilerEvent() { + profiling::Profiler::Instance().EndTimeAndRecordEvent(profiling::EventCategory::NODE_EVENT, name_, ts_); + } + + private: + TimePoint ts_; + const std::string name_; +}; + +} // namespace onnxruntime + +#define CODEGEN_PROFILER_EVENT(name) onnxruntime::ProfilerEvent name##_profiler_event(#name) + +#else + +#define CODEGEN_PROFILER_EVENT(name) + +#endif diff --git a/onnxruntime/core/codegen/common/registry.h b/onnxruntime/core/codegen/common/registry.h new file mode 100644 index 0000000000000..1ec06d4d8d96c --- /dev/null +++ b/onnxruntime/core/codegen/common/registry.h @@ -0,0 +1,70 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/common/common.h" +#include +#include +#include + +namespace onnxruntime { +namespace codegen { + +// RegistryBase is a customized unordered_map +// that keep ownership of passes, +// including 1) IR builder passes +// 2) Weight layout transformer passes +// 3) Scheduler passses, etc. + +template +class RegistryBase { + public: + RegistryBase() = default; + + bool Contains(const std::string& name) const { + return contents_.count(name) > 0; + } + + CONTENT_TYPE* Get(const std::string& name) const { + if (contents_.find(name) != contents_.end()) + return contents_.at(name).get(); + return nullptr; + } + + CONTENT_TYPE* RegisterOrGet( + const std::string& name, + std::unique_ptr&& ptr) { + if (!Contains(name)) + contents_.emplace(name, std::move(ptr)); + return Get(name); + } + + CONTENT_TYPE* RegisterOrGet( + std::unique_ptr&& ptr) { + return RegisterOrGet(ptr->Name(), std::move(ptr)); + } + + bool Register( + const std::string& name, + std::unique_ptr&& ptr) { + if (!Contains(name)) { + contents_.emplace(name, std::move(ptr)); + return true; + } + return false; + } + + bool Register( + std::unique_ptr&& ptr) { + return Register(ptr->Name(), std::move(ptr)); + } + + protected: + std::unordered_map> contents_; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(RegistryBase); +}; + +// Put common Registry Management utilities if these is any + +} // namespace codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/common/settings.cc b/onnxruntime/core/codegen/common/settings.cc new file mode 100644 index 0000000000000..c046f2892088d --- /dev/null +++ b/onnxruntime/core/codegen/common/settings.cc @@ -0,0 +1,74 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/common/settings.h" + +#include "core/common/logging/logging.h" +#include +#include + +namespace onnxruntime { +namespace codegen { + +CodeGenSettings& CodeGenSettings::Instance() { + static CodeGenSettings settings; + return settings; +} + +CodeGenSettings::CodeGenSettings() {} + +void CodeGenSettings::InsertOptions(const std::map& options) { + for (const auto& option : options) { + const auto& key = option.first; + const auto& value = option.second; + + auto iter = options_.find(key); + // found existing ones + if (iter != options_.end()) { + if (iter->second != value) { + LOGS_DEFAULT(CODEGEN_SETTINGS_LOG_LEVEL) << "CodeGenSettings: option" + << key << " is overridded from: " + << iter->second << " to: " << value; + iter->second = value; + } + } else { + options_.insert(std::make_pair(key, value)); + } + } +} + +void CodeGenSettings::DumpOptions() const { + std::ostringstream stream; + stream << "CodeGenSettings: dump all options" << std::endl; + for (const auto& option : options_) { + stream << " " << option.first << " = " << option.second << std::endl; + } + LOGS_DEFAULT(CODEGEN_SETTINGS_LOG_LEVEL) << stream.str(); +} + +std::string CodeGenSettings::GetOptionValue(const std::string& key) const { + const auto& iter = options_.find(key); + if (iter == options_.end()) { + LOGS_DEFAULT(CODEGEN_SETTINGS_LOG_LEVEL) << "CodeGenSettings::GetOptionValue: unrecognized option" << key; + return ""; + } + return iter->second; +} + +bool CodeGenSettings::HasOption(const std::string& key) const { + return options_.count(key) > 0; +} + +bool CodeGenSettings::OptionMatches(const std::string& key, const std::string& value) const { + if (!HasOption(key)) + return false; + +#ifdef _WIN32 + return 0 == _stricmp(options_.at(key).c_str(), value.c_str()); +#else + return 0 == strcasecmp(options_.at(key).c_str(), value.c_str()); +#endif +} + +} // namespace codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/common/settings.h b/onnxruntime/core/codegen/common/settings.h new file mode 100644 index 0000000000000..95a2282ccb1ff --- /dev/null +++ b/onnxruntime/core/codegen/common/settings.h @@ -0,0 +1,38 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include + +namespace onnxruntime { +namespace codegen { + +// use log level warning as default to make sure logs are outputted +#define CODEGEN_SETTINGS_LOG_LEVEL WARNING + +// This stores codegen settings to control dumps, execution preference, etc. +// CodeGenSettings could come from command line options or environment variables +// Or could come from a static variables in source code +class CodeGenSettings { + public: + // generic built-in options + constexpr static const char* kDumpAllOptions = "dump_all_options"; + constexpr static const char* kCodeGenDumpModule = "codegen_dump_module"; // dump tvm module + constexpr static const char* kCodeGenDumpLower = "codegen_dump_lower"; // dump lowered func + constexpr static const char* kCodeGenDumpSchedule = "codegen_dump_schedule"; // dump scheduler + + void InsertOptions(const std::map& options); + void DumpOptions() const; + std::string GetOptionValue(const std::string& key) const; + bool HasOption(const std::string& key) const; + bool OptionMatches(const std::string& key, const std::string& value) const; + static CodeGenSettings& Instance(); + + private: + CodeGenSettings(); + + std::map options_; +}; + +} // namespace codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/common/target_info.h b/onnxruntime/core/codegen/common/target_info.h new file mode 100644 index 0000000000000..da063545f0a1e --- /dev/null +++ b/onnxruntime/core/codegen/common/target_info.h @@ -0,0 +1,33 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { + +// CodeGenTarget holds meta info for backend code generation +// and will be lowered to a target of corresponding backend +// code generation, e.g. TVM's Target. +class CodeGenTarget { + public: + CodeGenTarget() {} + CodeGenTarget(const std::string& target_name) + : target_name_(target_name) {} + + virtual int NaturalVectorWidth(int /*bits*/) const { + return 1; + } + + const std::string& GetTargetName() const { + return target_name_; + } + + virtual ~CodeGenTarget() = default; + + private: + std::string target_name_{"unknown"}; // default name is unknown +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/common/utils.cc b/onnxruntime/core/codegen/common/utils.cc new file mode 100644 index 0000000000000..45c2436a18a82 --- /dev/null +++ b/onnxruntime/core/codegen/common/utils.cc @@ -0,0 +1,54 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/common/utils.h" + +#include +#include + +namespace onnxruntime { + +std::unique_ptr GetEnv(const char* var) { + char* val = nullptr; +#if _MSC_VER + size_t len; + + if (_dupenv_s(&val, &len, var)) { + // Something went wrong, just return nullptr. + return nullptr; + } +#else + val = getenv(var); +#endif // _MSC_VER + + if (val == nullptr) { + return nullptr; + } + + // On windows, we will have to explicitly free val. Instead of returning val + // to its caller and make distinguish between windows and linux, we return + // a unique_ptr, and it will be destroyed automatically after the caller + // completes. + size_t len_val = strlen(val) + 1; + auto p = std::make_unique(len_val); + // use explicit loop to get ride of VC's warning on unsafe copy + for (size_t i = 0; i < len_val; ++i) { + p[i] = val[i]; + } + return p; +} + +bool IsEnvVarDefined(const char* var) { + auto val = GetEnv(var); + return val != nullptr; +} + +int64_t TotalSize(const std::vector& shape) { + int64_t total = 1; + for (auto s : shape) { + total *= s; + } + return total; +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/common/utils.h b/onnxruntime/core/codegen/common/utils.h new file mode 100644 index 0000000000000..40f300888d680 --- /dev/null +++ b/onnxruntime/core/codegen/common/utils.h @@ -0,0 +1,20 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include +#include + +namespace onnxruntime { + +// Holding utility functions that are not tied to TVM and ORT + +std::unique_ptr GetEnv(const char* var); + +// Check if an environment variable is set +bool IsEnvVarDefined(const char* var); + +int64_t TotalSize(const std::vector& shape); + +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/common.h b/onnxruntime/core/codegen/mti/common.h new file mode 100644 index 0000000000000..be3896bed48fb --- /dev/null +++ b/onnxruntime/core/codegen/mti/common.h @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#define MTI_ASSERT(condition) \ + if (!(condition)) { \ + std::string error_msg = "Not satsified: " #condition \ + ": line " + std::to_string(__LINE__) + \ + " in file " + std::string(__FILE__) + "\n"; \ + throw std::runtime_error(error_msg); \ + } + diff --git a/onnxruntime/core/codegen/mti/debug/tvm_print.cc b/onnxruntime/core/codegen/mti/debug/tvm_print.cc new file mode 100644 index 0000000000000..0491636032b47 --- /dev/null +++ b/onnxruntime/core/codegen/mti/debug/tvm_print.cc @@ -0,0 +1,83 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/debug/tvm_print.h" + +#include "core/codegen/common/utils.h" +#include "core/codegen/common/dump_array.h" +#include "core/codegen/mti/common.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +TVM_REGISTER_GLOBAL("tvm.contrib.onnxruntime.print") + .set_body([](tvm::TVMArgs args, tvm::TVMRetValue* /*ret*/) { + DLTensor* X = args[0]; + DLTensor* Y = args[1]; + + DLDataType dtype = X->dtype; + std::vector shape; + int64_t total_size = 1; + for (int i = 0; i < X->ndim; ++i) { + shape.push_back(X->shape[i]); + total_size *= X->shape[i]; + } + + // pass X to Y + memcpy(static_cast(Y->data) + Y->byte_offset, + static_cast(X->data) + X->byte_offset, + total_size * dtype.bits / 8); + + if (tvm::runtime::TypeMatch(dtype, kDLFloat, 32)) { + float* data = reinterpret_cast(static_cast(X->data) + X->byte_offset); + DumpArray("float tensor:", data, shape); + } else if (tvm::runtime::TypeMatch(dtype, kDLInt, 8)) { + int8_t* data = reinterpret_cast(static_cast(X->data) + X->byte_offset); + DumpArray("int8 tensor:", data, shape); + } else if (tvm::runtime::TypeMatch(dtype, kDLInt, 16)) { + int16_t* data = reinterpret_cast(static_cast(X->data) + X->byte_offset); + DumpArray("int16 tensor:", data, shape); + } else if (tvm::runtime::TypeMatch(dtype, kDLInt, 32)) { + int32_t* data = reinterpret_cast(static_cast(X->data) + X->byte_offset); + DumpArray("int32 tensor:", data, shape); + } else if (tvm::runtime::TypeMatch(dtype, kDLUInt, 8)) { + uint8_t* data = reinterpret_cast(static_cast(X->data) + X->byte_offset); + DumpArray("uint8 tensor:", data, shape); + } else if (tvm::runtime::TypeMatch(dtype, kDLUInt, 16)) { + uint16_t* data = reinterpret_cast(static_cast(X->data) + X->byte_offset); + DumpArray("uint16 tensor:", data, shape); + } else if (tvm::runtime::TypeMatch(dtype, kDLUInt, 32)) { + uint32_t* data = reinterpret_cast(static_cast(X->data) + X->byte_offset); + DumpArray("uint32 tensor:", data, shape); + } else { + MTI_ASSERT(0 && "not implemented!"); + } + }); + +tvm::Array +PrintTVMTensorExtern(const tvm::Tensor& X, + const std::string& name) { + return topi::detail::make_extern( + {X->shape}, + {X->dtype}, + {X}, + [&](tvm::Array ins, tvm::Array outs) { + return topi::detail::call_packed({tvm::Expr("tvm.contrib.onnxruntime.print"), + topi::detail::pack_buffer(ins[0]), + topi::detail::pack_buffer(outs[0])}); + }, + name + "_print", "", {}); +} + +tvm::Tensor PrintImmutable(const tvm::Tensor& X) { + auto outputs = PrintTVMTensorExtern(X, X->op->name + "_print"); + return outputs[0]; +} + +void Print(tvm::Tensor& X) { + X = PrintImmutable(X); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/debug/tvm_print.h b/onnxruntime/core/codegen/mti/debug/tvm_print.h new file mode 100644 index 0000000000000..91a334785a2a4 --- /dev/null +++ b/onnxruntime/core/codegen/mti/debug/tvm_print.h @@ -0,0 +1,19 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Array PrintTVMTensorExtern( + const tvm::Tensor& X, + const std::string& name = "PrintTVM2DTensorExtern"); + +tvm::Tensor PrintImmutable(const tvm::Tensor& X); + +void Print(tvm::Tensor& X); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/binary_ops.cc b/onnxruntime/core/codegen/mti/math/binary_ops.cc new file mode 100644 index 0000000000000..f3048799458f4 --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/binary_ops.cc @@ -0,0 +1,70 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/math/binary_ops.h" + +#include "core/codegen/mti/math/unary_ops.h" +#include "core/codegen/mti/mti_tvm_utils.h" +#include "core/codegen/mti/tensor/cast_ops.h" +#include + +// Using namespace topi for override operator +-*/ +using namespace topi; + +namespace onnxruntime { +namespace tvm_codegen { + +#define TVM_BINARY_OP1(op, expr) \ + tvm::Tensor op(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name) { \ + return Rename(expr, name); \ + } \ + tvm::Tensor op(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name) { \ + return Rename(expr, name); \ + } + +#define TVM_BINARY_OP(op, expr) \ + TVM_BINARY_OP1(op, expr) \ + tvm::Tensor op(const tvm::Expr& lhs, const tvm::Tensor& rhs, const std::string& name) { \ + return Rename(expr, name); \ + } + +TVM_BINARY_OP(Add, lhs + rhs); +TVM_BINARY_OP(Div, lhs / rhs); +TVM_BINARY_OP(Max, maximum(lhs, rhs)); +TVM_BINARY_OP(Min, minimum(lhs, rhs)); +TVM_BINARY_OP(Mul, lhs* rhs); +TVM_BINARY_OP1(PRelu, Relu(lhs) - rhs * Relu(0 - lhs)); +TVM_BINARY_OP(Sub, lhs - rhs); + +tvm::Tensor Equal(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name) { + return topi::equal(lhs, rhs, name); +} +tvm::Tensor Equal(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name) { + return topi::equal(lhs, rhs, name); +} +tvm::Tensor Equal(const tvm::Expr& lhs, const tvm::Tensor& rhs, const std::string& name) { + return topi::equal(lhs, rhs, name); +} + +tvm::Tensor Greater(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name) { + return topi::greater(lhs, rhs, name); +} +tvm::Tensor Greater(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name) { + return topi::greater(lhs, rhs, name); +} +tvm::Tensor Greater(const tvm::Expr& lhs, const tvm::Tensor& rhs, const std::string& name) { + return topi::greater(lhs, rhs, name); +} + +tvm::Tensor Less(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name) { + return topi::less(lhs, rhs, name); +} +tvm::Tensor Less(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name) { + return topi::less(lhs, rhs, name); +} +tvm::Tensor Less(const tvm::Expr& lhs, const tvm::Tensor& rhs, const std::string& name) { + return topi::less(lhs, rhs, name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/binary_ops.h b/onnxruntime/core/codegen/mti/math/binary_ops.h new file mode 100644 index 0000000000000..dd51ce5e7917d --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/binary_ops.h @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Add(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name = "add"); +tvm::Tensor Add(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name = "add"); +tvm::Tensor Add(const tvm::Expr& lhs, const tvm::Tensor& rhs, const std::string& name = "add"); +tvm::Tensor Div(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name = "div"); +tvm::Tensor Div(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name = "div"); +tvm::Tensor Div(const tvm::Expr& lhs, const tvm::Tensor& rhs, const std::string& name = "div"); +tvm::Tensor Equal(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name = "equal"); +tvm::Tensor Equal(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name = "equal"); +tvm::Tensor Equal(const tvm::Expr& lhs, const tvm::Tensor& rhs, const std::string& name = "equal"); +tvm::Tensor Greater(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name = "greater"); +tvm::Tensor Greater(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name = "greater"); +tvm::Tensor Greater(const tvm::Expr& lhs, const tvm::Tensor& rhs, const std::string& name = "greater"); +tvm::Tensor Less(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name = "less"); +tvm::Tensor Less(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name = "less"); +tvm::Tensor Less(const tvm::Expr& lhs, const tvm::Tensor& rhs, const std::string& name = "less"); +tvm::Tensor Max(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name = "max"); +tvm::Tensor Max(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name = "max"); +tvm::Tensor Max(const tvm::Expr& lhs, const tvm::Tensor& rhs, const std::string& name = "max"); +tvm::Tensor Min(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name = "min"); +tvm::Tensor Min(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name = "min"); +tvm::Tensor Min(const tvm::Expr& lhs, const tvm::Tensor& rhs, const std::string& name = "min"); +tvm::Tensor Mul(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name = "mul"); +tvm::Tensor Mul(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name = "mul"); +tvm::Tensor Mul(const tvm::Expr& lhs, const tvm::Tensor& rhs, const std::string& name = "mul"); +tvm::Tensor PRelu(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name = "prelu"); +tvm::Tensor PRelu(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name = "prelu"); +tvm::Tensor Sub(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name = "sub"); +tvm::Tensor Sub(const tvm::Tensor& lhs, const tvm::Expr& rhs, const std::string& name = "sub"); +tvm::Tensor Sub(const tvm::Expr& lhs, const tvm::Tensor& rhs, const std::string& name = "sub"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/gemm.cc b/onnxruntime/core/codegen/mti/math/gemm.cc new file mode 100644 index 0000000000000..b5e5da5301775 --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/gemm.cc @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/math/gemm.h" + +#include "core/codegen/mti/math/matmul_ops.h" +#include "core/codegen/mti/mti_tvm_utils.h" +#include + +// Using namespace topi for override operator +-*/ +using namespace topi; + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Gemm(const tvm::Tensor& A, const tvm::Tensor& B, const tvm::Tensor& C, + bool trans_A, bool trans_B, float alpha, float beta, + const std::string& name) { + auto A_dot_B = MatMul2D(A, B, trans_A, trans_B, name + "_matmul2d"); + if (beta != 0) { + return Rename(alpha * A_dot_B + (beta * C), name); + } else { + return Rename(alpha * A_dot_B, name); + } +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/gemm.h b/onnxruntime/core/codegen/mti/math/gemm.h new file mode 100644 index 0000000000000..3bb205c13fdc9 --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/gemm.h @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Gemm(const tvm::Tensor& p_A, const tvm::Tensor& p_B, const tvm::Tensor& p_C, + bool trans_A, bool trans_B, float alpha, float beta, + const std::string& name = "gemm"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/logsoftmax.cc b/onnxruntime/core/codegen/mti/math/logsoftmax.cc new file mode 100644 index 0000000000000..cd8c2edae6959 --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/logsoftmax.cc @@ -0,0 +1,18 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/math/logsoftmax.h" + +#include "core/codegen/mti/tensor/reshape_ops.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor LogSoftmax(const tvm::Tensor& input, int64_t axis, const std::string& name) { + tvm::Tensor flatten_t = Flatten(input, axis, "logsoftmax_flatten"); + return Reshape(topi::nn::log_softmax(flatten_t, name), input->shape, "logsoftmax_reshape"); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/logsoftmax.h b/onnxruntime/core/codegen/mti/math/logsoftmax.h new file mode 100644 index 0000000000000..606a32806434b --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/logsoftmax.h @@ -0,0 +1,11 @@ +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor LogSoftmax(const tvm::Tensor& input, int64_t axis, const std::string& name = "logsoftmax"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/matmul_ops.cc b/onnxruntime/core/codegen/mti/math/matmul_ops.cc new file mode 100644 index 0000000000000..672aa3a6cf8db --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/matmul_ops.cc @@ -0,0 +1,138 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/math/matmul_ops.h" + +#include "core/codegen/mti/common.h" +#include "core/codegen/mti/mti_tvm_utils.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor MatMul2D(const tvm::Tensor& A, const tvm::Tensor& B, bool trans_a, bool trans_b, const std::string& name) { + return topi::matmul(A, B, trans_a, trans_b, name); +} + +/* + * Generic Matrix Multiplication + * + * If both arguments are 2-D, they are multiplied like conventional matrices. + * + * If either argument is N-D and N > 2, it is treated as a stack of matrices residing in the last two indexes and broadcast accordingly. + * + * If the first argument is 1-D, it is promoted to a matrix by prepending a 1 to its dimensions. + * After matrix multiplication the prepended 1 is removed. + * + * If the second argument is 1-D, it is promoted to a matrix by appending a 1 to its dimensions. + * After matrix multiplication the appended 1 is removed. + */ +tvm::Tensor MatMul(const tvm::Tensor& A, const tvm::Tensor& B, const std::string& name) { + int64_t a_rank = static_cast(A->shape.size()); + int64_t b_rank = static_cast(B->shape.size()); + const auto& A_shape = A->shape; + const auto& B_shape = B->shape; + if (a_rank == 2 && b_rank == 2) { + // 2-D X 2-D + return MatMul2D(A, B); + } else if (a_rank == 1 && b_rank == 1) { + // 1-D X 1-D + auto k = tvm::reduce_axis(tvm::Range(0, A_shape[0]), "k"); + + return tvm::compute( + {}, + [&](const tvm::Array& /*indices*/) { + return tvm::sum(A[k] * B[k], {k}); + }, + name); + } else if (a_rank == 1) { + // 1-D X n-D + auto k = tvm::reduce_axis(tvm::Range(0, A_shape[0]), "k"); + + auto l = [&](const tvm::Array& indices) { + auto ndims = indices.size(); + MTI_ASSERT(ndims >= 1); + tvm::Array b_indices; + for (size_t bi = 0; bi < ndims - 1; ++bi) { + b_indices.push_back(indices[bi]); + } + b_indices.push_back(k); + b_indices.push_back(indices[ndims - 1]); + return tvm::sum(A({k}) * B(b_indices), {k}); + }; + return tvm::compute(ConcatShapes(SliceShapeToDimension(B_shape, -2), SliceShapeFromDimension(B_shape, -1)), l, name); + } else if (b_rank == 1) { + // n-D X 1-D + auto k = tvm::reduce_axis(tvm::Range(0, B_shape[0]), "k"); + + auto l = [&](const tvm::Array& indices) { + tvm::Array a_indices(indices.begin(), indices.end()); + a_indices.push_back(k); + return tvm::sum(A(a_indices) * B({k}), {k}); + }; + return tvm::compute(SliceShapeToDimension(A->shape, -1), l, name); + } else { + // n-D X m-D + MTI_ASSERT(a_rank >= 2 && b_rank >= 2); + auto k = tvm::reduce_axis(tvm::Range(0, A_shape[a_rank - 1]), "k"); + + auto l = [&](const tvm::Array& indices) { + auto ndims = static_cast(indices.size()); + MTI_ASSERT(ndims > 2); + tvm::Array a_indices, b_indices; + + // handle broadcasting + int i = 0, a_idx = 0, b_idx = 0; + bool a_greater = a_rank > b_rank; + for (; i < std::abs(a_rank - b_rank); ++i) { + if (a_greater) { + a_indices.push_back(indices[i]); + a_idx++; + } else { + b_indices.push_back(indices[i]); + b_idx++; + } + } + for (; i < ndims - 2; ++i, ++a_idx, ++b_idx) { + auto tp = indices[i].type(); + if (IsOne(A_shape, a_idx)) { + a_indices.push_back(tvm::make_zero(tp)); + b_indices.push_back(indices[i]); + } else if (IsOne(B_shape, b_idx)) { + b_indices.push_back(tvm::make_zero(tp)); + a_indices.push_back(indices[i]); + } else { + a_indices.push_back(indices[i]); + b_indices.push_back(indices[i]); + } + } + + MTI_ASSERT(a_idx == a_rank - 2 && b_idx == b_rank - 2); + a_indices.push_back(indices[ndims - 2]); + a_indices.push_back(k); + + b_indices.push_back(k); + b_indices.push_back(indices[ndims - 1]); + + return tvm::sum(A(a_indices) * B(b_indices), {k}); + }; + + tvm::Array output_shape; + int64_t output_rank = std::max(a_rank, b_rank); + MTI_ASSERT(tvm::ir::Equal(A_shape[a_rank - 1], B_shape[b_rank - 2])); + for (int64_t i = 0; i < output_rank - 2; i++) { + tvm::Expr broadcasted_dim = tvm::make_const(HalideIR::Int(32), 1); + bool broadcasted = + BroadcastDim(A_shape, i, output_rank, broadcasted_dim) && + BroadcastDim(B_shape, i, output_rank, broadcasted_dim); + MTI_ASSERT(broadcasted); + output_shape.push_back(broadcasted_dim); + } + output_shape.push_back(A_shape[a_rank - 2]); + output_shape.push_back(B_shape[b_rank - 1]); + return tvm::compute(output_shape, l, name); + } +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/matmul_ops.h b/onnxruntime/core/codegen/mti/math/matmul_ops.h new file mode 100644 index 0000000000000..c149486a87fab --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/matmul_ops.h @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor MatMul2D(const tvm::Tensor& A, const tvm::Tensor& B, bool trans_a = false, bool trans_b = false, const std::string& name = "matmul2d"); + +tvm::Tensor MatMul(const tvm::Tensor& A, const tvm::Tensor& B, const std::string& name = "matmul"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/reduce_ops.cc b/onnxruntime/core/codegen/mti/math/reduce_ops.cc new file mode 100644 index 0000000000000..7d179e2b04316 --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/reduce_ops.cc @@ -0,0 +1,90 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/math/reduce_ops.h" + +#include "core/codegen/mti/math/binary_ops.h" +#include "core/codegen/mti/math/unary_ops.h" +#include "core/codegen/mti/mti_tvm_utils.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor ArgMax(const tvm::Tensor& X, int64_t axis, bool keep_dims, const std::string& name) { + return Rename(topi::argmax(X, ToTvmArrayInt({axis}), keep_dims), name); +} + +tvm::Tensor ArgMin(const tvm::Tensor& X, int64_t axis, bool keep_dims, const std::string& name) { + return Rename(topi::argmin(X, ToTvmArrayInt({axis}), keep_dims), name); +} + +tvm::Tensor ReduceL1(const tvm::Tensor& X, const std::vector& axes, bool keep_dims, const std::string& name) { + return ReduceSum(Abs(X), axes, keep_dims, name); +} + +tvm::Tensor ReduceL2(const tvm::Tensor& X, const std::vector& axes, bool keep_dims, const std::string& name) { + return Sqrt(ReduceSumSquare(X, axes, keep_dims), name); +} + +tvm::Tensor ReduceLogSum(const tvm::Tensor& X, const std::vector& axes, bool keep_dims, const std::string& name) { + return Log(ReduceSum(X, axes, keep_dims), name); +} + +tvm::Tensor ReduceLogSumExp(const tvm::Tensor& X, const std::vector& axes, bool keep_dims, const std::string& name) { + tvm::Tensor reduce_max = ReduceMax(X, axes, true); + tvm::Tensor exp_delta = Exp(Sub(X, reduce_max)); + tvm::Tensor reduce_max_keep_dims = ReduceMax(X, axes, keep_dims); + return Add(ReduceLogSum(exp_delta, axes, keep_dims), reduce_max_keep_dims, name); +} + +tvm::Tensor ReduceMax(const tvm::Tensor& X, const std::vector& axes, bool keep_dims, const std::string& name) { + return Rename(topi::max(X, ToTvmArrayInt(axes), keep_dims), name); +} + +tvm::Tensor ReduceMean(const tvm::Tensor& X, const std::vector& axes, bool keep_dims, const std::string& name) { + tvm::Tensor reduce_sum = ReduceSum(X, axes, keep_dims); + tvm::Expr count = tvm::make_const(reduce_sum->dtype, 1.0f); + if (axes.empty()) { + for (const auto& dim : X->shape) + count = count * dim; + } else { + for (int64_t axis : axes) { + int64_t i = HandleNegativeAxis(axis, X->shape.size()); + count = count * X->shape[i]; + } + } + return tvm::compute( + reduce_sum->shape, + [&](const tvm::Array& i) { + return reduce_sum(i) / count; + }, + name); +} + +tvm::Tensor ReduceMin(const tvm::Tensor& X, const std::vector& axes, bool keep_dims, const std::string& name) { + return Rename(topi::min(X, ToTvmArrayInt(axes), keep_dims), name); +} + +tvm::Tensor ReduceProd(const tvm::Tensor& X, const std::vector& axes, bool keep_dims, const std::string& name) { + auto prod = [](tvm::Expr source, tvm::Array rdom) { + tvm::Var x("x", source.type()), y("y", source.type()); + tvm::Expr Rename_element = tvm::make_const(source.type(), 1.0f); + tvm::ir::CommReducer combiner = + tvm::ir::CommReducerNode::make({x}, {y}, {x * y}, {Rename_element}); + return tvm::ir::Reduce::make(combiner, {source}, rdom, tvm::make_const(tvm::Bool(1), true), 0); + }; + + return Rename(topi::CommReduce(X, ToTvmArrayInt(axes), prod, keep_dims, true), name); +} + +tvm::Tensor ReduceSum(const tvm::Tensor& X, const std::vector& axes, bool keep_dims, const std::string& name) { + return Rename(topi::sum(X, ToTvmArrayInt(axes), keep_dims), name); +} + +tvm::Tensor ReduceSumSquare(const tvm::Tensor& X, const std::vector& axes, bool keep_dims, const std::string& name) { + return Rename(topi::sum(Mul(X, X), ToTvmArrayInt(axes), keep_dims), name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/reduce_ops.h b/onnxruntime/core/codegen/mti/math/reduce_ops.h new file mode 100644 index 0000000000000..f782df5e6515f --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/reduce_ops.h @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor ArgMax(const tvm::Tensor& X, + int64_t axis, + bool keep_dims, + const std::string& name = "argmax"); + +tvm::Tensor ArgMin(const tvm::Tensor& X, + int64_t axis, + bool keep_dims, + const std::string& name = "argmin"); + +tvm::Tensor ReduceL1(const tvm::Tensor& X, + const std::vector& axes, + bool keep_dims, + const std::string& name = "reduce_l1"); + +tvm::Tensor ReduceL2(const tvm::Tensor& X, + const std::vector& axes, + bool keep_dims, + const std::string& name = "reduce_l2"); + +tvm::Tensor ReduceLogSum(const tvm::Tensor& X, + const std::vector& axes, + bool keep_dims, + const std::string& name = "reduce_log_sum"); + +tvm::Tensor ReduceLogSumExp(const tvm::Tensor& X, + const std::vector& axes, + bool keep_dims, + const std::string& name = "argmareduce_log_sum_exp"); + +tvm::Tensor ReduceMax(const tvm::Tensor& X, + const std::vector& axes, + bool keep_dims, + const std::string& name = "reduce_max"); + +tvm::Tensor ReduceMean(const tvm::Tensor& X, + const std::vector& axes, + bool keep_dims, + const std::string& name = "reduce_mean"); + +tvm::Tensor ReduceMin(const tvm::Tensor& X, + const std::vector& axes, + bool keep_dims, + const std::string& name = "reduce_min"); + +tvm::Tensor ReduceProd(const tvm::Tensor& X, + const std::vector& axes, + bool keep_dims, + const std::string& name = "reduce_prod"); + +tvm::Tensor ReduceSum(const tvm::Tensor& X, + const std::vector& axes, + bool keep_dims, + const std::string& name = "reduce_sum"); + +tvm::Tensor ReduceSumSquare(const tvm::Tensor& X, + const std::vector& axes, + bool keep_dims, + const std::string& name = "reduce_sum_square"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/softmax.cc b/onnxruntime/core/codegen/mti/math/softmax.cc new file mode 100644 index 0000000000000..d7404137bb873 --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/softmax.cc @@ -0,0 +1,18 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/math/softmax.h" + +#include "core/codegen/mti/tensor/reshape_ops.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Softmax(const tvm::Tensor& input, int64_t axis, const std::string& name) { + tvm::Tensor flatten_t = Flatten(input, axis, "softmax_flatten"); + return Reshape(topi::nn::softmax(flatten_t, 1, name), input->shape, "softmax_reshape"); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/softmax.h b/onnxruntime/core/codegen/mti/math/softmax.h new file mode 100644 index 0000000000000..fb16fbaeb56a2 --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/softmax.h @@ -0,0 +1,11 @@ +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Softmax(const tvm::Tensor& input, int64_t axis, const std::string& name = "softmax"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/unary_ops.cc b/onnxruntime/core/codegen/mti/math/unary_ops.cc new file mode 100644 index 0000000000000..7f45a9115fb0b --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/unary_ops.cc @@ -0,0 +1,142 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/math/unary_ops.h" + +#include "core/codegen/common/settings.h" +#include "core/codegen/mti/mti_tvm_utils.h" +#include +#include +#include +#include + +// Using namespace topi for override operator +-*/ +using namespace topi; + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Abs(const tvm::Tensor& X, const std::string& name) { + return abs(X, name); +} + +tvm::Tensor Affine(const tvm::Tensor& X, float alpha, float beta, const std::string& name) { + return Rename(alpha * X + beta, name); +} + +tvm::Tensor Ceil(const tvm::Tensor& X, const std::string& name) { + return topi::ceil(X, name); +} + +tvm::Tensor Clip(const tvm::Tensor& X, float min_value, float max_value, const std::string& name) { + auto Y = tvm::compute( + X->shape, + [&](const tvm::Array& indices) { + return tvm::min(tvm::max(X(indices), min_value), max_value); + }, + name); + return Y; +} + +tvm::Tensor Elu(const tvm::Tensor& X, float alpha, const std::string& name) { + return Rename(Relu(X) - alpha * Relu(1 - Exp(X)), name); +} + +tvm::Tensor Exp(const tvm::Tensor& X, const std::string& name) { + return tvm::compute( + X->shape, + [&](const tvm::Array& indices) { + return tvm::exp(X(indices)); + }, + name); +} + +tvm::Tensor Floor(const tvm::Tensor& X, const std::string& name) { + return topi::floor(X, name); +} + +tvm::Tensor HardSigmoid(const tvm::Tensor& X, float alpha, float beta, const std::string& name) { + return maximum(0, minimum(1, alpha * X + beta), name); +} + +tvm::Tensor LeakyRelu(const tvm::Tensor& X, float alpha, const std::string& name) { + return Rename(Relu(X) - alpha * Relu(0 - X), name); +} + +tvm::Tensor Log(const tvm::Tensor& X, const std::string& name) { + return tvm::compute( + X->shape, + [&](const tvm::Array& indices) { + return tvm::log(X(indices)); + }, + name); +} + +tvm::Tensor Neg(const tvm::Tensor& X, const std::string& name) { + return negative(X, name); +} + +tvm::Tensor ParametricSoftplus(const tvm::Tensor& X, float alpha, float beta, const std::string& name) { + return Rename(alpha * Softplus(beta * X), name); +} + +tvm::Tensor Reciprocal(const tvm::Tensor& X, const std::string& name) { + return Rename(1 / X, name); +} + +tvm::Tensor Relu(const tvm::Tensor& X, const std::string& name) { + return maximum(X, 0, name); +} + +tvm::Tensor ScaledTanh(const tvm::Tensor& X, float alpha, float beta, const std::string& name) { + return Rename(alpha * Tanh(beta * X), name); +} + +tvm::Tensor Selu(const tvm::Tensor& X, float alpha, float gamma, const std::string& name) { + return Rename(gamma * (-alpha * Relu(1 - Exp(X)) + Relu(X)), name); +} + +tvm::Tensor Sigmoid(const tvm::Tensor& X, const std::string& name) { + return tvm::compute( + X->shape, + [&](const tvm::Array& indices) { + return tvm::ir::Select::make(X(indices) > 0, + 1 / (1 + tvm::exp(-X(indices))), + tvm::exp(X(indices)) / (tvm::exp(X(indices)) + 1)); + }, + name); +} + +tvm::Tensor SignNoZero(const tvm::Tensor& X, const std::string& name) { + return Rename(greater_equal(X, 0) * 2 - 1, name); +} + +tvm::Tensor Softplus(const tvm::Tensor& X, const std::string& name) { + return Rename(Log(1 + Exp(Neg(Abs(X)))) + Relu(X), name); +} + +tvm::Tensor Softsign(const tvm::Tensor& X, const std::string& name) { + return Rename(X / (1 + Abs(X)), name); +} + +tvm::Tensor Sqrt(const tvm::Tensor& X, const std::string& name) { + return sqrt(X, name); +} + +tvm::Tensor Tanh(const tvm::Tensor& X, const std::string& name) { + return tvm::compute( + X->shape, + [&](const tvm::Array& indices) { + return tvm::ir::Select::make(X(indices) < 0, + (tvm::exp(2 * X(indices)) - 1) / (tvm::exp(2 * X(indices)) + 1), + (1 - tvm::exp(-2 * X(indices))) / (1 + tvm::exp(-2 * X(indices)))); + }, + name); +} + +tvm::Tensor ThresholdedRelu(const tvm::Tensor& X, float alpha, const std::string& name) { + return topi::where(greater(X, alpha), X, topi::full_like(X, tvm::make_zero(X->dtype)), name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/math/unary_ops.h b/onnxruntime/core/codegen/mti/math/unary_ops.h new file mode 100644 index 0000000000000..ae1f17099fa7e --- /dev/null +++ b/onnxruntime/core/codegen/mti/math/unary_ops.h @@ -0,0 +1,36 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Abs(const tvm::Tensor& X, const std::string& name = "abs"); +tvm::Tensor Affine(const tvm::Tensor& X, float alpha, float beta, const std::string& name = "affine"); +tvm::Tensor Ceil(const tvm::Tensor& X, const std::string& name = "ceil"); +tvm::Tensor Clip(const tvm::Tensor& X, float min_value, float max_value, const std::string& name = "clip"); +tvm::Tensor Elu(const tvm::Tensor& X, float alpha, const std::string& name = "elu"); +tvm::Tensor Exp(const tvm::Tensor& X, const std::string& name = "exp"); +tvm::Tensor Floor(const tvm::Tensor& X, const std::string& name = "floor"); +tvm::Tensor HardSigmoid(const tvm::Tensor& X, float alpha, float beta, const std::string& name = "hard_sigmoid"); +tvm::Tensor LeakyRelu(const tvm::Tensor& X, float alpha, const std::string& name = "leaky_relu"); +tvm::Tensor Log(const tvm::Tensor& X, const std::string& name = "log"); +tvm::Tensor Neg(const tvm::Tensor& X, const std::string& name = "neg"); +tvm::Tensor ParametricSoftplus(const tvm::Tensor& X, float alpha, float beta, const std::string& name = "parametric_softplus"); +tvm::Tensor Reciprocal(const tvm::Tensor& X, const std::string& name = "reciprocal"); +tvm::Tensor Relu(const tvm::Tensor& X, const std::string& name = "relu"); +tvm::Tensor ScaledTanh(const tvm::Tensor& X, float alpha, float beta, const std::string& name = "scaled_tanh"); +tvm::Tensor Selu(const tvm::Tensor& X, float alpha, float gamma, const std::string& name = "selu"); +tvm::Tensor Sigmoid(const tvm::Tensor& X, const std::string& name = "sigmoid"); +tvm::Tensor SignNoZero(const tvm::Tensor& X, const std::string& name = "sign_no_zero"); +tvm::Tensor Softplus(const tvm::Tensor& X, const std::string& name = "softplus"); +tvm::Tensor Softsign(const tvm::Tensor& X, const std::string& name = "softsign"); +tvm::Tensor Sqrt(const tvm::Tensor& X, const std::string& name = "sqrt"); +tvm::Tensor Tanh(const tvm::Tensor& X, const std::string& name = "tanh"); +tvm::Tensor ThresholdedRelu(const tvm::Tensor& X, float alpha, const std::string& name = "thresholded_relu"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/mti_tvm_utils.cc b/onnxruntime/core/codegen/mti/mti_tvm_utils.cc new file mode 100644 index 0000000000000..e905a34432a6e --- /dev/null +++ b/onnxruntime/core/codegen/mti/mti_tvm_utils.cc @@ -0,0 +1,162 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/mti_tvm_utils.h" + +#include "core/codegen/common/settings.h" +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Array ToTvmArray(const std::vector& shape) { + tvm::Array arr; + for (size_t i = 0; i < shape.size(); ++i) { + arr.push_back(tvm::Expr(static_cast(shape[i]))); + } + return arr; +} + +tvm::Array ToTvmArrayInt(const std::vector& shape) { + tvm::Array arr; + for (size_t i = 0; i < shape.size(); ++i) { + arr.push_back(shape[i]); + } + return arr; +} + +tvm::Expr SizeToDimension(const tvm::Array& shape, int64_t axis) { + tvm::Expr size(1); + auto rank = shape.size(); + if (static_cast(axis) != rank) { + axis = HandleNegativeAxis(axis, rank); + } + for (size_t d = 0; d < std::min(rank, static_cast(axis)); ++d) + size = tvm::ir::Simplify(size * shape[d]); + return size; +} + +tvm::Expr SizeFromDimension(const tvm::Array& shape, int64_t axis) { + tvm::Expr size(1); + auto rank = shape.size(); + if (static_cast(axis) != rank) { + axis = HandleNegativeAxis(axis, rank); + } + for (size_t d = static_cast(axis); d < rank; ++d) + size = tvm::ir::Simplify(size * shape[d]); + return size; +} + +tvm::Expr RoundUp(tvm::Expr value, tvm::Expr alignment) { + return tvm::ir::Simplify((value + alignment - 1) / alignment * alignment); +} + +tvm::Array ConcatShapes( + const tvm::Array& shape1, + const tvm::Array& shape2) { + tvm::Array result; + for (size_t i = 0; i < shape1.size(); i++) + result.push_back(shape1[i]); + for (size_t i = 0; i < shape2.size(); i++) + result.push_back(shape2[i]); + return result; +} + +tvm::Tensor Rename(tvm::Tensor X, const std::string& name) { + const_cast(X->op->name) = name; + return X; +} + +tvm::Array SliceShape(const tvm::Array& shape, const std::vector& axes) { + tvm::Array new_shape; + for (auto axis : axes) { + CHECK(axis < static_cast(shape.size())); + new_shape.push_back(shape[axis]); + } + return new_shape; +} + +tvm::Array SliceShapeFromDimension(const tvm::Array& shape, int64_t axis) { + int64_t rank = static_cast(shape.size()); + axis = HandleNegativeAxis(axis, rank); + std::vector axes; + for (auto i = axis; i < rank; ++i) + axes.push_back(i); + return SliceShape(shape, axes); +} + +tvm::Array SliceShapeToDimension(const tvm::Array& shape, int64_t axis) { + int64_t rank = static_cast(shape.size()); + axis = HandleNegativeAxis(axis, rank); + std::vector axes; + for (auto i = 0; i < axis; ++i) + axes.push_back(i); + return SliceShape(shape, axes); +} + +bool IsOne(const tvm::Array& shape, int64_t axis) { + int64_t rank = static_cast(shape.size()); + axis = HandleNegativeAxis(axis, rank); + const auto& dim = shape[axis]; + auto* p = tvm::as_const_int(dim); + return p != nullptr && *p == 1; +} + +tvm::Tensor Promote(const tvm::Expr& expr, const tvm::Array& shape, const std::string& name) { + return tvm::compute( + shape, + [&](const tvm::Array&) { + return expr; + }, + name); +} + +void DumpTVMModuleToFile(const std::string& filename_prefix, tvm::runtime::Module& module) { + const codegen::CodeGenSettings& settings = codegen::CodeGenSettings::Instance(); + if (!settings.HasOption(codegen::CodeGenSettings::kCodeGenDumpModule)) + return; + + static int dump_module_cnt = 0; + // ISSUE: note that all option values are converted to lower case. It doesn't cause + // any issue currently, because all supported formats (i.e. file exts) are of lower case. + // Just keep in mind that we might have issue if somehow we started to support dump + // formats with upper case, although it's quite unlikely. + std::string format = settings.GetOptionValue(codegen::CodeGenSettings::kCodeGenDumpModule); + std::string module_filename = filename_prefix + "_" + std::to_string(dump_module_cnt++) + "." + format; + module->SaveToFile(module_filename, format); +} + +tvm::Tensor MakeZeroTensor(const tvm::Array& shape, + HalideIR::Type type, + const std::string& name) { + auto l = [&](const tvm::Array& /*indices*/) { + return tvm::make_zero(type); + }; + return tvm::compute(shape, l, name); +} + +bool BroadcastDim(const tvm::Array& shape, size_t i, size_t output_rank, tvm::Expr& dim) { + if (i >= output_rank - shape.size()) { + auto new_dim = shape[shape.size() - output_rank + i]; + if (tvm::ir::Equal(new_dim, dim)) + return true; + + const int64_t* p_new = tvm::as_const_int(new_dim); + if (p_new != nullptr && *p_new == 1) { + return true; + } else { + const int64_t* p_old = tvm::as_const_int(dim); + if (p_old != nullptr && *p_old == 1) { + dim = new_dim; + return true; + } + } + return false; + } + // auto broadcast to outer dims + return true; +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/mti_tvm_utils.h b/onnxruntime/core/codegen/mti/mti_tvm_utils.h new file mode 100644 index 0000000000000..3f65658554f2c --- /dev/null +++ b/onnxruntime/core/codegen/mti/mti_tvm_utils.h @@ -0,0 +1,64 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include "core/codegen/mti/common.h" + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Array ToTvmArray(const std::vector& shape); + +tvm::Array ToTvmArrayInt(const std::vector& shape); + +// Helper function to compute sub shape size to axis (not included) +tvm::Expr SizeToDimension(const tvm::Array& shape, int64_t axis); + +// Helper function to compute sub shape size from axis (included) +tvm::Expr SizeFromDimension(const tvm::Array& shape, int64_t axis); + +// Helper function to align +tvm::Expr RoundUp(tvm::Expr value, tvm::Expr alignment); + +tvm::Array ConcatShapes( + const tvm::Array& shape1, + const tvm::Array& shape2); + +// Helper function to rename tvm::Tensor +tvm::Tensor Rename(tvm::Tensor X, const std::string& name); + +// Helper function to slice TVM shape +tvm::Array SliceShape(const tvm::Array& shape, const std::vector& axes); + +// Helper function to slice TVM shape from axis (inclusive). +// Basically, this function returns the shape of [axis, shape.size()-1] +tvm::Array SliceShapeFromDimension(const tvm::Array& shape, int64_t axis); + +// this function returns the shape of [0, axis-1] +tvm::Array SliceShapeToDimension(const tvm::Array& shape, int64_t axis); + +// check if dimension is 1 +bool IsOne(const tvm::Array& shape, int64_t axis); + +// Helper function to convert tvm::Expr to tvm::Tensor +tvm::Tensor Promote(const tvm::Expr& expr, + const tvm::Array& shape, + const std::string& name = "PromoteExpr"); + +tvm::Tensor MakeZeroTensor(const tvm::Array& shape, HalideIR::Type type, const std::string& name); + +void DumpTVMModuleToFile(const std::string& filename_prefix, tvm::runtime::Module& module); + +bool BroadcastDim(const tvm::Array& shape, size_t i, size_t output_rank, tvm::Expr& dim); + +inline int64_t HandleNegativeAxis(int64_t axis, int64_t rank) { + MTI_ASSERT(axis >= -rank && axis <= rank - 1); + return axis = axis < 0 ? (axis + rank) : axis; +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/nn/conv_ops.cc b/onnxruntime/core/codegen/mti/nn/conv_ops.cc new file mode 100644 index 0000000000000..ca2fd3ee897c6 --- /dev/null +++ b/onnxruntime/core/codegen/mti/nn/conv_ops.cc @@ -0,0 +1,199 @@ +#include "core/codegen/mti/nn/conv_ops.h" + +#include "core/codegen/mti/math/matmul_ops.h" +#include "core/codegen/mti/tensor/pad_ops.h" +#include "core/codegen/mti/tensor/reshape_ops.h" +#include "core/codegen/mti/tensor/transpose.h" + +namespace onnxruntime { +namespace tvm_codegen { + +static tvm::Tensor PadTensor1D(const tvm::Tensor& input, + const tvm::Array& padding, + size_t width_axis, + const std::string& name) { + auto pad_left = padding[0]; + auto pad_right = padding[1]; + + tvm::Array pad_before(std::vector(input->shape.size(), 0)); + pad_before.Set(width_axis, pad_left); + tvm::Array pad_after(std::vector(input->shape.size(), 0)); + pad_after.Set(width_axis, pad_right); + + const int64_t* padding_w0 = tvm::as_const_int(pad_left); + const int64_t* padding_w1 = tvm::as_const_int(pad_right); + + const bool do_pad = ((padding_w0 != nullptr && *padding_w0) || + (padding_w1 != nullptr && *padding_w1)); + + return do_pad ? Pad(input, pad_before, pad_after, + 0, "constant", name + "_input_padded") + : input; +} + +tvm::Tensor Conv1D(const tvm::Tensor& input, + const tvm::Tensor& filter, + const tvm::Array& out_shape, + const tvm::Array& stride, + const tvm::Array& padding, + const std::string& name) { + size_t channel_axis = 1; + size_t width_axis = 2; + + auto stride_width = stride[width_axis - 2]; + + auto input_padded = PadTensor1D(input, padding, width_axis, name); + auto rc = tvm::reduce_axis((tvm::Range(0, filter->shape[1])), "rc"); + auto rx = tvm::reduce_axis((tvm::Range(0, filter->shape[2])), "rx"); + + return tvm::compute( + out_shape, + [&](const tvm::Array& output) { + tvm::Array indices; + for (const tvm::Var& var : output) { + indices.push_back(var); + } + indices.Set(channel_axis, rc); + indices.Set(width_axis, output[width_axis] * stride_width + rx); + + return tvm::sum(input_padded(indices) * filter({output[1], rc, rx}), + {rc, rx}); + }, + name); +} + +tvm::Tensor Conv2D(const tvm::Tensor& input, + const tvm::Tensor& filter, + const tvm::Array& output_shape, + const tvm::Array& stride, + const tvm::Array& padding, + const std::string& name) { + // Gemm Convolution + const int64_t* batch_size = tvm::as_const_int(input->shape[0]); + if (batch_size != nullptr && *batch_size == 1) + return Conv2D_gemm(input, filter, output_shape, stride, padding); + + // Native Convolution + return Conv2D_native(input, filter, output_shape, stride, padding); +} + +static tvm::Tensor PadTensor2D(const tvm::Tensor& input, + const tvm::Array& padding, + size_t height_axis, + size_t width_axis, + const std::string& name) { + auto pad_top = padding[0]; + auto pad_left = padding[1]; + auto pad_bottom = padding[2]; + auto pad_right = padding[3]; + + tvm::Array pad_before(std::vector(input->shape.size(), 0)); + pad_before.Set(height_axis, pad_top); + pad_before.Set(width_axis, pad_left); + + tvm::Array pad_after(std::vector(input->shape.size(), 0)); + pad_after.Set(height_axis, pad_bottom); + pad_after.Set(width_axis, pad_right); + + const int64_t* padding_h0 = tvm::as_const_int(pad_top); + const int64_t* padding_w0 = tvm::as_const_int(pad_left); + const int64_t* padding_h1 = tvm::as_const_int(pad_bottom); + const int64_t* padding_w1 = tvm::as_const_int(pad_right); + + const bool do_pad = ((padding_h0 != nullptr && *padding_h0) || + (padding_w0 != nullptr && *padding_w0)) || + ((padding_h1 != nullptr && *padding_h1) || + (padding_w1 != nullptr && *padding_w1)); + + return do_pad ? Pad(input, pad_before, pad_after, + 0, "constant", name + "_input_padded") + : input; +} + +tvm::Tensor Conv2D_native(const tvm::Tensor& input, + const tvm::Tensor& filter, + const tvm::Array& out_shape, + const tvm::Array& stride, + const tvm::Array& padding, + const std::string& name) { + size_t channel_axis = 1; + size_t height_axis = 2; + size_t width_axis = 3; + + auto stride_height = stride[height_axis - 2]; + auto stride_width = stride[width_axis - 2]; + + auto input_padded = PadTensor2D(input, padding, height_axis, width_axis, name); + + auto rc = tvm::reduce_axis((tvm::Range(0, filter->shape[1])), "rc"); + auto ry = tvm::reduce_axis((tvm::Range(0, filter->shape[2])), "ry"); + auto rx = tvm::reduce_axis((tvm::Range(0, filter->shape[3])), "rx"); + + return tvm::compute( + out_shape, + [&](const tvm::Array& output) { + tvm::Array indices; + for (const tvm::Var& var : output) { + indices.push_back(var); + } + indices.Set(channel_axis, rc); + indices.Set(height_axis, output[height_axis] * stride_height + ry); + indices.Set(width_axis, output[width_axis] * stride_width + rx); + + return tvm::sum(input_padded(indices) * filter({output[1], rc, ry, rx}), + {rc, ry, rx}); + }, + name); +} + +tvm::Tensor Conv2D_gemm(const tvm::Tensor& input, + const tvm::Tensor& filter, + const tvm::Array& out_shape, + const tvm::Array& stride, + const tvm::Array& padding, + const std::string& name) { + size_t height_axis = 2; + size_t width_axis = 3; + + auto stride_height = stride[height_axis - 2]; + auto stride_width = stride[width_axis - 2]; + + auto input_padded = PadTensor2D(input, padding, height_axis, width_axis, name); + + tvm::Array img_col_tmp(std::vector(6, 0)); + img_col_tmp.Set(0, out_shape[0]); + img_col_tmp.Set(1, out_shape[2]); + img_col_tmp.Set(2, out_shape[3]); + img_col_tmp.Set(3, filter->shape[1]); + img_col_tmp.Set(4, filter->shape[2]); + img_col_tmp.Set(5, filter->shape[3]); + + auto img_col = tvm::compute( + img_col_tmp, + [&](const tvm::Array& output) { + tvm::Array indices; + indices.push_back(output[0]); + indices.push_back(output[3]); + indices.push_back(output[1] * stride_height + output[4]); + indices.push_back(output[2] * stride_width + output[5]); + return input_padded(indices); + }, + name); + + tvm::Array input_col_shape(std::vector(2, 0)); + input_col_shape.Set(0, img_col_tmp[1] * img_col_tmp[2]); + input_col_shape.Set(1, img_col_tmp[3] * img_col_tmp[4] * img_col_tmp[5]); + auto input_col = Reshape(img_col, input_col_shape); + + tvm::Array filter_row_shape(std::vector(2, 0)); + filter_row_shape.Set(0, filter->shape[0]); + filter_row_shape.Set(1, filter->shape[1] * filter->shape[2] * filter->shape[3]); + auto filter_row = Reshape(filter, filter_row_shape, name); + + auto Y = MatMul2D(input_col, filter_row, false, true, name); + auto Y_T = Transpose(Y, /*axes=*/{}, name); + return Reshape(Y_T, out_shape, name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/nn/conv_ops.h b/onnxruntime/core/codegen/mti/nn/conv_ops.h new file mode 100644 index 0000000000000..1396c216865a7 --- /dev/null +++ b/onnxruntime/core/codegen/mti/nn/conv_ops.h @@ -0,0 +1,39 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Conv1D(const tvm::Tensor& input, + const tvm::Tensor& filter, + const tvm::Array& output_shape, + const tvm::Array& stride, + const tvm::Array& padding, + const std::string& name = "conv1d"); + +tvm::Tensor Conv2D(const tvm::Tensor& input, + const tvm::Tensor& filter, + const tvm::Array& output_shape, + const tvm::Array& stride, + const tvm::Array& padding, + const std::string& name = "conv2d"); + +tvm::Tensor Conv2D_native(const tvm::Tensor& input, + const tvm::Tensor& filter, + const tvm::Array& output_shape, + const tvm::Array& stride, + const tvm::Array& padding, + const std::string& name = "conv2d_native"); + +tvm::Tensor Conv2D_gemm(const tvm::Tensor& input, + const tvm::Tensor& filter, + const tvm::Array& output_shape, + const tvm::Array& stride, + const tvm::Array& padding, + const std::string& name = "conv2d_gemm"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/nn/lstm.cc b/onnxruntime/core/codegen/mti/nn/lstm.cc new file mode 100644 index 0000000000000..1148b0924e869 --- /dev/null +++ b/onnxruntime/core/codegen/mti/nn/lstm.cc @@ -0,0 +1,140 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/nn/lstm.h" + +#include "core/codegen/mti/math/binary_ops.h" +#include "core/codegen/mti/math/unary_ops.h" +#include "core/codegen/mti/math/matmul_ops.h" +#include "core/codegen/mti/math/reduce_ops.h" +#include "core/codegen/mti/mti_tvm_utils.h" +#include "core/codegen/mti/tensor/reshape_ops.h" +#include "core/codegen/mti/tensor/split.h" + +namespace onnxruntime { +namespace tvm_codegen { + +/* +`X` - input tensor +`i` - input gate +`o` - output gate +`f` - forget gate +`c` - cell gate +`t` - time step (t-1 means previous time step) + +`W[iofc]` - W parameter weight matrix for input, output, forget, and cell gates +`R[iofc]` - R recurrence weight matrix for input, output, forget, and cell gates +`Wb[iofc]` - W bias vectors for input, output, forget, and cell gates +`Rb[iofc]` - R bias vectors for input, output, forget, and cell gates +`P[iof]` - P peephole weight vector for input, output, and forget gates +`WB[iofc]` - W parameter weight matrix for backward input, output, forget, and cell gates +`RB[iofc]` - R recurrence weight matrix for backward input, output, forget, and cell gates +`WBb[iofc]` - W bias vectors for backward input, output, forget, and cell gates +`RBb[iofc]` - R bias vectors for backward input, output, forget, and cell gates +`PB[iof]` - P peephole weight vector for backward input, output, and forget gates + +`H` - Hidden state +`num_directions` - 2 if direction == bidirectional else 1 + +Equations (Default: f=Sigmoid, g=Tanh, h=Tanh): + it = f(Xt*(Wi^T) + Ht-1*(Ri^T) + Pi (.) Ct-1 + Wbi + Rbi) + ft = f(Xt*(Wf^T) + Ht-1*(Rf^T) + Pf (.) Ct-1 + Wbf + Rbf) + ct = g(Xt*(Wc^T) + Ht-1*(Rc^T) + Wbc + Rbc) + Ct = ft (.) Ct-1 + it (.) ct + ot = f(Xt*(Wo^T) + Ht-1*(Ro^T) + Po (.) Ct + Wbo + Rbo) + Ht = ot (.) h(Ct) +*/ + +void LSTM_cell( + const LSTMAttributes& lstm_attrs, + const tvm::Tensor& X, + const tvm::Tensor& W, + const tvm::Tensor& R, + const tvm::Tensor& B, + bool has_B, + const tvm::Tensor& prev_H, + const tvm::Tensor& prev_C, + const tvm::Tensor& P, + bool has_P, + tvm::Tensor& Y_h, + tvm::Tensor& Y_c) { + // Input projection: Xt*(W[iofc]^T) for forward direction or Xt*(WB[iofc]^T) for reverse direction + // (batch_size, input_size) * trans(4 * hidden_size, input_size) => (batch_size, 4 * hidden_size) + tvm::Tensor input_proj = MatMul2D(X, W, /*trans_a*/ false, /*trans_b*/ true); + + // Hidden projection: Ht-1*(R[iofc]^T) for forward direction or Ht-1*(RB[iofc]^T) for reverse direction + // (batch_size, hidden_size) * trans(4 * hidden_size, hidden_size) => (batch_size, 4 * hidden_size) + tvm::Tensor hidden_proj = MatMul2D(prev_H, R, /*trans_a*/ false, /*trans_b*/ true); + + // (batch_size, 4 * hidden_size) + tvm::Tensor sum_proj = Add(input_proj, hidden_proj); + + // Concatenation of [Wb[iofc], Rb[iofc]] or [WBb[iofc], RBb[iofc]] + if (has_B) { + // (8 * hidden_size) -> (2, 4 * hidden_size) -> (1, 4 * hidden_size), should be done in const folding + tvm::Tensor reduce_B = + ReduceSum(Reshape(B, {2, 4 * static_cast(lstm_attrs.hidden_size)}), {0}, /*keep_dims*/ true); + // (batch_size, 4 * hidden_size) via broadcasting reduce_B + sum_proj = Add(sum_proj, reduce_B); + } + + std::vector iofc_sum_split_sizes(4, lstm_attrs.hidden_size); + // Split sum_proj into iofc, where each gate proj is of (batch_size, hidden_size) + tvm::Array iofc_sum_projs = Split(sum_proj, ToTvmArray(iofc_sum_split_sizes), /*axis*/ 1); + MTI_ASSERT(iofc_sum_projs.size() == 4); + tvm::Tensor i_proj = iofc_sum_projs[0], + o_proj = iofc_sum_projs[1], + f_proj = iofc_sum_projs[2], + c_proj = iofc_sum_projs[3]; + + tvm::Tensor P_i, P_o, P_f; + if (has_P) { + std::vector iof_p_split_sizes(3, lstm_attrs.hidden_size); + // Split P into P_i, P_o, P_f, in const pre-processing (P_i, P_f might be merged?) + // where each P_[iof] has the shape of (hidden_size) + tvm::Array iof_P_projs = Split(P, ToTvmArray(iof_p_split_sizes), /*axis*/ 0); + MTI_ASSERT(iof_P_projs.size() == 3); + P_i = iof_P_projs[0], + P_o = iof_P_projs[1], + P_f = iof_P_projs[2]; + + // (batch_size, hidden_size) via broadcasting P_[if] + i_proj = Add(i_proj, Mul(P_i, prev_C)); + f_proj = Add(f_proj, Mul(P_f, prev_C)); + } + + // TODO: handle more general cases for activations f, h, g and activation_alpha and + // activation_beta. We may consider to move some code such as ActivationInfo from deep_cpu_lstm + // into a common header file, because the code can be used here. + + // Note that by default f = Sigmoid, g = Tanh, h = Tanh + + // it = f(Xt*(Wi^T) + Ht-1*(Ri^T) + Pi (.) Ct-1 + Wbi + Rbi) + // shape: (batch_size, hidden_size) + tvm::Tensor i_t = Sigmoid(i_proj); + // ft = f(Xt*(Wf^T) + Ht-1*(Rf^T) + Pf (.) Ct-1 + Wbf + Rbf) + // shape: (batch_size, hidden_size) + tvm::Tensor f_t = Sigmoid(f_proj); + // ct = g(Xt*(Wc^T) + Ht-1*(Rc^T) + Wbc + Rbc) + // shape: (batch_size, hidden_size) + tvm::Tensor c_t = Tanh(c_proj); + + // Ct = ft (.) Ct-1 + it (.) ct + // shape: (batch_size, hidden_size) + Y_c = Add(Mul(f_t, prev_C), Mul(i_t, c_t), Y_c->op->name); + + // ot = f(Xt*(Wo^T) + Ht-1*(Ro^T) + Po (.) Ct + Wbo + Rbo) + // shape: (batch_size, hidden_size) + if (has_P) { + o_proj = Add(o_proj, Mul(P_o, Y_c)); + } + // ot = f(Xt*(Wo^T) + Ht-1*(Ro^T) + Po (.) Ct + Wbo + Rbo) + // shape: (batch_size, hidden_size) + o_proj = Sigmoid(o_proj); + // Ht = ot (.) h(Ct) + // shape: (batch_size, hidden_size) + Y_h = Mul(o_proj, Tanh(Y_c), Y_h->op->name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/nn/lstm.h b/onnxruntime/core/codegen/mti/nn/lstm.h new file mode 100644 index 0000000000000..851fa880c4427 --- /dev/null +++ b/onnxruntime/core/codegen/mti/nn/lstm.h @@ -0,0 +1,35 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +// A bubble now. But don't remove it +// TODO: refactor the LSTMcell building to a tvm function +// and move it here + +namespace onnxruntime { +namespace tvm_codegen { + +struct LSTMAttributes { + LSTMAttributes(int64_t hidden_size_p) : hidden_size(hidden_size_p) {} + int64_t hidden_size; +}; + +void LSTM_cell( + const LSTMAttributes& lstm_attrs, + const tvm::Tensor& X, + const tvm::Tensor& W, + const tvm::Tensor& R, + const tvm::Tensor& B, + bool has_B, + const tvm::Tensor& prev_H, + const tvm::Tensor& prev_C, + const tvm::Tensor& P, + bool has_P, + tvm::Tensor& Y_h, + tvm::Tensor& Y_c); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/nn/pool_ops.cc b/onnxruntime/core/codegen/mti/nn/pool_ops.cc new file mode 100644 index 0000000000000..5af944186c178 --- /dev/null +++ b/onnxruntime/core/codegen/mti/nn/pool_ops.cc @@ -0,0 +1,58 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/nn/pool_ops.h" + +#include + +namespace onnxruntime { +namespace tvm_codegen { + +// TODO: topi only support 2d-pool, MaxPool1d and MaxPool3d will need to be added if necessary. +// only support version < 8 for topi doesn't come with implementation to output index tensor +tvm::Tensor MaxPool( + const tvm::Tensor& input, + const tvm::Array& kernel_size, + const tvm::Array& stride_size, + const tvm::Array& padding_size, + const std::string& layout, + bool count_include_pad) { + return topi::nn::pool(input, kernel_size, stride_size, padding_size, + topi::nn::kMaxPool, + false, + layout, + count_include_pad); +} + +tvm::Tensor AveragePool( + const tvm::Tensor& input, + const tvm::Array& kernel_size, + const tvm::Array& stride_size, + const tvm::Array& padding_size, + const std::string& layout, + bool count_include_pad) { + return topi::nn::pool(input, kernel_size, stride_size, padding_size, + topi::nn::kAvgPool, + false, + layout, + count_include_pad); +} + +tvm::Tensor GlobalMaxPool( + const tvm::Tensor& input, + const std::string& layout) { + return topi::nn::global_pool(input, + topi::nn::kMaxPool, + layout); +} + +tvm::Tensor GlobalAveragePool( + const tvm::Tensor& input, + const std::string& layout) { + return topi::nn::global_pool(input, + topi::nn::kAvgPool, + layout); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/nn/pool_ops.h b/onnxruntime/core/codegen/mti/nn/pool_ops.h new file mode 100644 index 0000000000000..23fbda913e277 --- /dev/null +++ b/onnxruntime/core/codegen/mti/nn/pool_ops.h @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor MaxPool(const tvm::Tensor& input, + const tvm::Array& kernel_size, + const tvm::Array& stride_size, + const tvm::Array& padding_size, + const std::string& layout, + bool count_include_pad); + +tvm::Tensor AveragePool(const tvm::Tensor& input, + const tvm::Array& kernel_size, + const tvm::Array& stride_size, + const tvm::Array& padding_size, + const std::string& layout, + bool count_include_pad); + +tvm::Tensor GlobalMaxPool(const tvm::Tensor& input, + const std::string& layout); + +tvm::Tensor GlobalAveragePool(const tvm::Tensor& input, + const std::string& layout); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/cast_ops.cc b/onnxruntime/core/codegen/mti/tensor/cast_ops.cc new file mode 100644 index 0000000000000..a8fc86488d82b --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/cast_ops.cc @@ -0,0 +1,37 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/tensor/cast_ops.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Cast(const tvm::Tensor& X, tvm::Type type, const std::string& name) { + return topi::cast(X, type, name); +} + +// handle cases where bool is reprented as uint8 (e.g. in ONNX). +tvm::Tensor CastToUInt8Bool(const tvm::Tensor& X, const std::string& name) { + return tvm::compute( + X->shape, + [&](const tvm::Array& indices) { + auto val = X(indices); + // A special cast from float16 to bool, first cast up to float32, + // to workaround a float16 bug in many TVM backends. + // Intel Skylake is one of them. https://github.com/dmlc/tvm/issues/2959 + // TODO: remove it, after TVM is fixed + if (X->dtype == HalideIR::Float(16)) + val = tvm::cast(HalideIR::Float(32), val); + return tvm::ir::Select::make(topi::equal(val, tvm::make_zero(val.type())), + tvm::make_zero(HalideIR::UInt(8)), + tvm::make_const(HalideIR::UInt(8), 1)); + }, + name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/cast_ops.h b/onnxruntime/core/codegen/mti/tensor/cast_ops.h new file mode 100644 index 0000000000000..02f6f9cb1fde7 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/cast_ops.h @@ -0,0 +1,15 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Cast(const tvm::Tensor& X, tvm::Type type, const std::string& name = "cast"); +tvm::Tensor CastToUInt8Bool(const tvm::Tensor& X, const std::string& name = "cast_uint8_bool"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/concat_ops.cc b/onnxruntime/core/codegen/mti/tensor/concat_ops.cc new file mode 100644 index 0000000000000..13e8a4edc320a --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/concat_ops.cc @@ -0,0 +1,83 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/tensor/concat_ops.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include "gsl/gsl_util" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Concat(const tvm::Array& inputs, + int64_t axis, + const std::string& name) { + return ConcatSafe(inputs, axis, name); +} + +// Note topi's implementation requires control flow within iterations to avoid out-of-bound access. +// Therefore, MTI implements a ConcatSafe that does not have out-of-bound access, +// and does not requires control or predicate. +tvm::Tensor ConcatSafe(const tvm::Array& inputs, + int64_t axis, + const std::string& name) { + axis = HandleNegativeAxis(axis, gsl::narrow(inputs[0]->shape.size())); + MTI_ASSERT(axis < gsl::narrow(inputs[0]->shape.size()) && "axis out of bounds"); + + tvm::Array axis_sizes; + for (auto t : inputs) { + axis_sizes.push_back(t->shape[axis]); + } + + tvm::Expr join_size = axis_sizes[0]; + for (size_t i = 1; i < axis_sizes.size(); ++i) { + join_size += axis_sizes[i]; + } + join_size = tvm::ir::Simplify(join_size); + tvm::Array out_shape; + for (size_t i = 0; i < inputs[0]->shape.size(); ++i) { + out_shape.push_back(i == gsl::narrow(axis) ? join_size : inputs[0]->shape[i]); + } + + return tvm::compute( + out_shape, [&](const tvm::Array& ovars) { + tvm::Array indices; + + // preset + tvm::Expr min = 0; + tvm::Expr extent = axis_sizes[0]; + tvm::Expr offset = 0; + tvm::Expr ret; + + //input i = 0 + for (size_t j = 0; j < ovars.size(); ++j) { + if (j == gsl::narrow(axis)) { + tvm::Expr ivar = ovars[j]; + indices.push_back(tvm::max(tvm::min(ivar, min + extent - 1), min)); + } else { + indices.push_back(ovars[j]); + } + } + ret = inputs[0](indices); + + for (size_t i = 1; i < inputs.size(); ++i) { + offset += extent; + tvm::Expr min = 0; + extent = axis_sizes[i]; + auto j = gsl::narrow(axis); + tvm::Expr ivar = ovars[j] - offset; + indices.Set(j, tvm::max(tvm::min(ivar, min + extent - 1), min)); + + ret = tvm::ir::Select::make(ivar >= 0, + inputs[i](indices), + ret); + } + + return ret; + }, + name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/concat_ops.h b/onnxruntime/core/codegen/mti/tensor/concat_ops.h new file mode 100644 index 0000000000000..153afebb44615 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/concat_ops.h @@ -0,0 +1,15 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Concat(const tvm::Array& inputs, int64_t axis, const std::string& name = "concat"); +tvm::Tensor ConcatSafe(const tvm::Array& inputs, int64_t axis, const std::string& name = "concat_safe"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/crop.cc b/onnxruntime/core/codegen/mti/tensor/crop.cc new file mode 100644 index 0000000000000..3fe569100df12 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/crop.cc @@ -0,0 +1,58 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/tensor/crop.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Crop(const tvm::Tensor& t, + const tvm::Array& border, + const tvm::Array& scale, + const std::string& name) { + MTI_ASSERT(t->shape.size() == 4); + tvm::Expr N = t->shape[0]; + tvm::Expr C = t->shape[1]; + tvm::Expr H = t->shape[2]; + tvm::Expr W = t->shape[3]; + + MTI_ASSERT(border.size() == 4); + tvm::Expr leftBorder = border[0]; + tvm::Expr topBorder = border[1]; + tvm::Expr rightBorder = border[2]; + tvm::Expr bottomBorder = border[3]; + + tvm::Expr bottomLimit = H - bottomBorder; + tvm::Expr rightLimit = W - rightBorder; + + if (!scale.empty()) { + CHECK_EQ(scale.size(), 2); + bottomLimit = topBorder + scale[0]; + rightLimit = leftBorder + scale[1]; + } + + tvm::Array output_shape; + output_shape.push_back(tvm::ir::Simplify(N)); + output_shape.push_back(tvm::ir::Simplify(C)); + output_shape.push_back(tvm::ir::Simplify(bottomLimit - topBorder)); + output_shape.push_back(tvm::ir::Simplify(rightLimit - leftBorder)); + + auto l = [&](const tvm::Array& ovars) { + tvm::Array indices; + + indices.push_back(tvm::min(ovars[0], output_shape[0] - 1)); + indices.push_back(tvm::min(ovars[1], output_shape[1] - 1)); + indices.push_back(tvm::min(topBorder + ovars[2], topBorder + output_shape[2] - 1)); + indices.push_back(tvm::min(leftBorder + ovars[3], leftBorder + output_shape[3] - 1)); + + return t(indices); + }; + + return tvm::compute(output_shape, l, name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/crop.h b/onnxruntime/core/codegen/mti/tensor/crop.h new file mode 100644 index 0000000000000..ffb6a05c70504 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/crop.h @@ -0,0 +1,17 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Crop(const tvm::Tensor& t, + const tvm::Array& border, + const tvm::Array& scale = {}, + const std::string& name = "crop"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/gather.cc b/onnxruntime/core/codegen/mti/tensor/gather.cc new file mode 100644 index 0000000000000..283d29c6eaa5d --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/gather.cc @@ -0,0 +1,53 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/tensor/gather.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include "gsl/gsl_util" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Gather(const tvm::Tensor& t, + int64_t axis, + const tvm::Tensor& indices, + const std::string& name) { + // handle negative axis + axis = HandleNegativeAxis(axis, gsl::narrow(t->shape.size())); + size_t axis_t = gsl::narrow(axis); + + tvm::Array output_shape; + for (size_t i = 0; i < axis_t; ++i) + output_shape.push_back(t->shape[i]); + + for (size_t i = 0; i < indices->shape.size(); ++i) + output_shape.push_back(indices->shape[i]); + + for (size_t i = axis_t + 1; i < t->shape.size(); ++i) + output_shape.push_back(t->shape[i]); + + auto l = [&](const tvm::Array& ovars) { + tvm::Array ivars; + for (size_t i = 0; i < t->shape.size(); ++i) { + if (i < axis_t) { + ivars.push_back(ovars[i]); + } else if (i == axis_t) { + tvm::Array idx_vars; + for (size_t d = 0; d < indices->shape.size(); ++d) + idx_vars.push_back(ovars[axis_t + d]); + ivars.push_back(tvm::cast(tvm::Int(32), indices(idx_vars))); // tvm indices must be Int32 + } else { + ivars.push_back(ovars[i - 1 + indices->shape.size()]); + } + } + return tvm::ir::Select::make((ivars[axis_t] >= 0) && (ivars[axis_t] < t->shape[axis_t]), + t(ivars), tvm::make_zero(t->dtype)); + }; + + return tvm::compute(output_shape, l, name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/gather.h b/onnxruntime/core/codegen/mti/tensor/gather.h new file mode 100644 index 0000000000000..a44bf3e4127d5 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/gather.h @@ -0,0 +1,17 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Gather(const tvm::Tensor& t, + int64_t axis, + const tvm::Tensor& indices, + const std::string& name = "gather"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/pad_ops.cc b/onnxruntime/core/codegen/mti/tensor/pad_ops.cc new file mode 100644 index 0000000000000..2f688290d109e --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/pad_ops.cc @@ -0,0 +1,121 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/tensor/pad_ops.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +// Note topi::pad does not support modes {edge, reflect} +// Therefore, MTI implements a generic Pad +tvm::Tensor Pad(const tvm::Tensor& t, + const tvm::Array& pad_before, + const tvm::Array& pad_after, + float pad_value, + const std::string& mode, + const std::string& name) { + MTI_ASSERT(pad_before.size() >= 1); + MTI_ASSERT(pad_before.size() == pad_after.size()); + MTI_ASSERT(pad_before.size() == t->shape.size()); + + tvm::Array output_shape; + for (size_t i = 0; i < t->shape.size(); ++i) { + output_shape.push_back( + tvm::ir::Simplify(t->shape[i] + pad_before[i] + pad_after[i])); + } + + auto l = [&](const tvm::Array& ovars) { + tvm::Array conds; + tvm::Array indices; + tvm::Array coords; + + for (size_t i = 0; i < t->shape.size(); ++i) { + tvm::Expr ivar = ovars[i] - pad_before[i]; + tvm::Expr min = 0; + tvm::Expr extent = t->shape[i]; + + conds.push_back(ivar < min); + conds.push_back(ivar >= min + extent); + indices.push_back(tvm::max(tvm::min(ivar, min + extent - 1), min)); + + if (mode == "reflect") { + // calculate indices for reflect mode + tvm::Expr limit = extent - 1; + tvm::Expr coord = ivar - min; + // Avoid mod zero when tensor shape has 1, + // e.g. input shape is [1, 3, 3] instead of [3, 3] + auto* p_limit = tvm::as_const_int(limit); + if (p_limit != nullptr && *p_limit != 0) + coord = (coord + 2 * limit) % (2 * limit); // avoid negative value + coord = coord - limit; + coord = tvm::abs(coord); + coord = limit - coord; + coord = coord + min; + coords.push_back(coord); + } + } + + if (mode == "reflect") { + return tvm::ir::Select::make(topi::detail::Map(conds, tvm::ir::Or::make), + t(coords), t(indices)); + } else if (mode == "constant") { + return tvm::ir::Select::make(topi::detail::Map(conds, tvm::ir::Or::make), + tvm::make_const(t->dtype, pad_value), t(indices)); + } + + // default mode is edge + return t(indices); + }; + + return tvm::compute(output_shape, l, name); +} + +tvm::Tensor Pad(const tvm::Tensor& t, + const tvm::Array& output_shape, + const tvm::Expr& pad_value, + const std::string& name) { + MTI_ASSERT(t->dtype == pad_value.type()); + + auto l = [&](const tvm::Array& ovars) { + tvm::Array conds; + tvm::Array indices; + + for (size_t i = 0; i < t->shape.size(); ++i) { + tvm::Expr ivar = ovars[i]; + tvm::Expr min = 0; + tvm::Expr extent = t->shape[i]; + + conds.push_back(ivar < min); + conds.push_back(ivar >= min + extent); + indices.push_back(tvm::max(tvm::min(ivar, min + extent - 1), min)); + } + + return tvm::ir::Select::make(topi::detail::Map(conds, tvm::ir::Or::make), + pad_value, t(indices)); + }; + + return tvm::compute(output_shape, l, name); +} + +tvm::Tensor PadLastDim(const tvm::Tensor& t, + const int32_t align_size, + const tvm::Expr& pad_value, + const std::string& name) { + auto input_shape = t->shape; + tvm::Array out_shape; + size_t input_shape_rank = input_shape.size(); + for (size_t i = 0; i < input_shape_rank - 1; ++i) { + out_shape.push_back(input_shape[i]); + } + out_shape.push_back( + (input_shape[input_shape_rank - 1] + align_size - 1) / + align_size * align_size); + + return Pad(t, out_shape, pad_value, name + "_pad"); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/pad_ops.h b/onnxruntime/core/codegen/mti/tensor/pad_ops.h new file mode 100644 index 0000000000000..6e8e350d71e97 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/pad_ops.h @@ -0,0 +1,34 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +// ONNX Pad semantics +tvm::Tensor Pad(const tvm::Tensor& t, + const tvm::Array& pad_before, + const tvm::Array& pad_after, + float pad_value = 0.0f, + const std::string& mode = "constant", + const std::string& name = "pad"); + +// Other common Pad interfaces +// Pad for a given shape +tvm::Tensor Pad(const tvm::Tensor& t, + const tvm::Array& output_shape, + const tvm::Expr& pad_value, + const std::string& name = "pad"); + +// Pad for the last dim only. +// This is widely used for weight layout to guard alignment +tvm::Tensor PadLastDim(const tvm::Tensor& t, + const int32_t align_size, + const tvm::Expr& pad_value, + const std::string& name = "pad_last_dim"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/reshape_ops.cc b/onnxruntime/core/codegen/mti/tensor/reshape_ops.cc new file mode 100644 index 0000000000000..817fb32c2837a --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/reshape_ops.cc @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/tensor/reshape_ops.h" + +#include "core/codegen/mti/common.h" +#include "core/codegen/mti/mti_tvm_utils.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Flatten(const tvm::Tensor& X, int64_t axis, const std::string& name) { + const auto& input_shape = X->shape; + return Reshape(X, {SizeToDimension(input_shape, axis), SizeFromDimension(input_shape, axis)}, name); +} + +tvm::Tensor Identity(const tvm::Tensor& X, const std::string& name) { + return Reshape(X, X->shape, name); +} + +tvm::Tensor Reshape(const tvm::Tensor& X, const tvm::Array& new_shape, const std::string& name) { + if (new_shape.size() > 0) { + auto X_dim = SizeToDimension(X->shape, X->shape.size()); + auto new_dim = SizeToDimension(new_shape, new_shape.size()); + auto* pX_dim = tvm::as_const_int(X_dim); + auto* pNew_dim = tvm::as_const_int(new_dim); + + if (pX_dim != nullptr && pNew_dim != nullptr) { + MTI_ASSERT(*pX_dim == *pNew_dim); + } + return topi::reshape(X, new_shape, name); + } else { + // generate empty dim tensor with origial input data value + tvm::Array tmp_shape; + tmp_shape.push_back(1); + auto tmp_tensor = topi::reshape(X, tmp_shape); + return tvm::compute( + new_shape, + [&](const tvm::Array&) { + return tmp_tensor[0]; + }, + name); + } +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/reshape_ops.h b/onnxruntime/core/codegen/mti/tensor/reshape_ops.h new file mode 100644 index 0000000000000..e23d62e4c57b0 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/reshape_ops.h @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Flatten(const tvm::Tensor& X, int64_t axis, const std::string& name = "flatten"); +tvm::Tensor Identity(const tvm::Tensor& X, const std::string& name = "identity"); +tvm::Tensor Reshape(const tvm::Tensor& X, const tvm::Array& new_shape, const std::string& name = "reshape"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/slice.cc b/onnxruntime/core/codegen/mti/tensor/slice.cc new file mode 100644 index 0000000000000..4caf3946ce6a9 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/slice.cc @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/tensor/slice.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +static const int64_t max_range = INT_MAX; + +tvm::Expr position(const tvm::Expr& dim, const tvm::Integer& offset) { + if (offset->value >= max_range) + return dim; + else if (offset->value < 0) + return dim + offset; + else + return offset; +} + +tvm::Tensor Slice(const tvm::Tensor& X, + const tvm::Array& starts, + const tvm::Array& ends, + const std::string& name) { + tvm::Array output_shape; + for (size_t i = 0; i < X->shape.size(); ++i) { + tvm::Expr start = position(X->shape[i], starts[i]); + tvm::Expr end = position(X->shape[i], ends[i]); + output_shape.push_back(tvm::ir::Simplify(end - start)); + } + return tvm::compute( + output_shape, + [&](const tvm::Array& ovars) { + tvm::Array ivars; + for (size_t i = 0; i < X->shape.size(); ++i) + ivars.push_back(ovars[i] + tvm::ir::Simplify(position(X->shape[i], starts[i]))); + + return X(ivars); + }, + name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/slice.h b/onnxruntime/core/codegen/mti/tensor/slice.h new file mode 100644 index 0000000000000..26f53650b1b6d --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/slice.h @@ -0,0 +1,17 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Slice(const tvm::Tensor& X, + const tvm::Array& starts, + const tvm::Array& ends, + const std::string& name = "slice"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/split.cc b/onnxruntime/core/codegen/mti/tensor/split.cc new file mode 100644 index 0000000000000..7264f94e390b4 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/split.cc @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/tensor/split.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include "gsl/gsl_util" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +// Similar to numpy, topi::split takes split indices rather than the +// sizes of the splits. Thus we implement our own. +tvm::Array Split(const tvm::Tensor& X, + const tvm::Array& split_sizes, + int64_t axis, + const std::string& name) { + MTI_ASSERT(axis < gsl::narrow(X->shape.size())); + size_t axis_t = gsl::narrow(axis); + + tvm::Array> output_shapes; + int num_splits = gsl::narrow(split_sizes.size()); + for (auto& s : split_sizes) { + tvm::Array shape; + for (size_t i = 0; i < axis_t; i++) { + shape.push_back(X->shape[i]); + } + shape.push_back(s); + for (size_t i = axis_t + 1; i < X->shape.size(); i++) { + shape.push_back(X->shape[i]); + } + output_shapes.push_back(shape); + } + + tvm::Array res; + int idx = 0; + for (int i_split = 0; i_split < num_splits; ++i_split) { + tvm::Expr s = split_sizes[i_split]; + auto l = [&](const tvm::Array& indices) { + tvm::Array new_indices; + for (size_t i = 0; i < axis_t; i++) { + new_indices.push_back(indices[i]); + } + new_indices.push_back(indices[axis_t] + idx); + for (size_t i = axis_t + 1; i < X->shape.size(); i++) { + new_indices.push_back(indices[i]); + } + MTI_ASSERT(topi::detail::IsConstInt(s)); + MTI_ASSERT(new_indices.size() == X->shape.size()); + int size = topi::detail::GetConstInt(s); + idx += size; + return X(new_indices); + }; + res.push_back(tvm::compute(output_shapes[i_split], l, name)); + } + + MTI_ASSERT(topi::detail::IsConstInt(X->shape[axis_t])); + int size_of_splitted_axis = static_cast(topi::detail::GetConstInt(X->shape[axis_t])); + MTI_ASSERT(idx == size_of_splitted_axis); + return res; +} + +tvm::Array SplitWithIndices(const tvm::Tensor& X, + const tvm::Array& split_sizes, + int64_t axis, + const std::string& name) { + return topi::split(X, split_sizes, gsl::narrow(axis), name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/split.h b/onnxruntime/core/codegen/mti/tensor/split.h new file mode 100644 index 0000000000000..bcb9c47d936dd --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/split.h @@ -0,0 +1,25 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +// ONNX Split semantics +tvm::Array Split(const tvm::Tensor& X, + const tvm::Array& split_sizes, + int64_t axis, + const std::string& name = "split"); + +// Another common Split interface +// Split with chunck indices +tvm::Array SplitWithIndices(const tvm::Tensor& X, + const tvm::Array& split_sizes, + int64_t axis, + const std::string& name = "split_with_indices"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/tile.cc b/onnxruntime/core/codegen/mti/tensor/tile.cc new file mode 100644 index 0000000000000..57cf7097c5ae7 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/tile.cc @@ -0,0 +1,40 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/tensor/tile.h" +#include "core/codegen/mti/mti_tvm_utils.h" +#include "gsl/gsl_util" + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Tile(const tvm::Tensor& t, + const std::vector& repeats, + const std::string& name) { + MTI_ASSERT(repeats.size() == t->shape.size()); + tvm::Array output_shape; + + bool repeats_zero = false; + for (size_t i = 0; i < t->shape.size(); ++i) { + if (repeats[i] == 0) + repeats_zero = true; + output_shape.push_back(t->shape[i] * gsl::narrow(repeats[i])); + } + + auto l = [&](const tvm::Array& ovars) { + if (repeats_zero) + return tvm::make_zero(t->dtype); + + tvm::Array ivars; + for (size_t i = 0; i < t->shape.size(); ++i) { + tvm::Expr ovar = ovars[i]; + ivars.push_back(ovar % t->shape[i]); + } + return t(ivars); + }; + + return tvm::compute(output_shape, l, name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/tile.h b/onnxruntime/core/codegen/mti/tensor/tile.h new file mode 100644 index 0000000000000..7ce331fb5ea95 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/tile.h @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Tile(const tvm::Tensor& t, + const std::vector& repeats, + const std::string& name = "tile"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/transpose.cc b/onnxruntime/core/codegen/mti/tensor/transpose.cc new file mode 100644 index 0000000000000..873ff8d7f1708 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/transpose.cc @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/tensor/transpose.h" + +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Transpose(const tvm::Tensor& X, const tvm::Array& axes, const std::string& name) { + return topi::transpose(X, axes, name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/transpose.h b/onnxruntime/core/codegen/mti/tensor/transpose.h new file mode 100644 index 0000000000000..a2a98fedf1e79 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/transpose.h @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Transpose(const tvm::Tensor& X, + const tvm::Array& axes, + const std::string& name = "transpose"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/where.cc b/onnxruntime/core/codegen/mti/tensor/where.cc new file mode 100644 index 0000000000000..2bdac3cae7ef5 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/where.cc @@ -0,0 +1,36 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/mti/mti_tvm_utils.h" +#include "core/codegen/mti/tensor/where.h" + +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Where(const tvm::Tensor& B, + const tvm::Tensor& X, + const tvm::Tensor& Y, + const std::string& name) { + size_t rank = std::max(std::max(B->shape.size(), X->shape.size()), Y->shape.size()); + tvm::Array output_shape; + for (size_t i = 0; i < rank; ++i) { + tvm::Expr dim = tvm::make_const(HalideIR::Int(32), 1); + bool broadcasted = + BroadcastDim(B->shape, i, rank, dim) && + BroadcastDim(X->shape, i, rank, dim) && + BroadcastDim(Y->shape, i, rank, dim); + MTI_ASSERT(broadcasted); + output_shape.push_back(dim); + } + + return topi::where(topi::broadcast_to(B, output_shape), + topi::broadcast_to(X, output_shape), + topi::broadcast_to(Y, output_shape), + name); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/where.h b/onnxruntime/core/codegen/mti/tensor/where.h new file mode 100644 index 0000000000000..68c5288eb3580 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/where.h @@ -0,0 +1,17 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Where(const tvm::Tensor& B, + const tvm::Tensor& X, + const tvm::Tensor& Y, + const std::string& name = "where"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/codegen_context.cc b/onnxruntime/core/codegen/target/codegen_context.cc new file mode 100644 index 0000000000000..17bef98c0591c --- /dev/null +++ b/onnxruntime/core/codegen/target/codegen_context.cc @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/codegen_context.h" + +#include "core/codegen/common/common.h" + +namespace onnxruntime { +namespace tvm_codegen { + +CodeGenContext::CodeGenContext( + const codegen::CodeGenHandle* handle) + : handle_(handle), unname_symbol_counter_(0) {} + +tvm::Var CodeGenContext::GetOrCreateDynamicDim(const std::string& name) { + if (dynamic_dims_.count(name) == 0) + dynamic_dims_.emplace(name, tvm::Var(name)); + + return dynamic_dims_.at(name); +} + +std::string CodeGenContext::CreateUnnamedSymbol() { + return "unnamed_" + std::to_string(unname_symbol_counter_++); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/codegen_context.h b/onnxruntime/core/codegen/target/codegen_context.h new file mode 100644 index 0000000000000..641552bd3b2e8 --- /dev/null +++ b/onnxruntime/core/codegen/target/codegen_context.h @@ -0,0 +1,44 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/codegen/common/handle.h" +#include "core/codegen/common/common.h" +#include "core/common/common.h" +#include "core/framework/data_types.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +// CodeGenContext is a data structure involving across passes +// Compiler developers can use it to store meta data +// to support fine-grained control of code generation +class CodeGenContext { + public: + CodeGenContext(const codegen::CodeGenHandle* handle); + + virtual ~CodeGenContext() = default; + + // returns tvm::Var for the dynamic dim + tvm::Var GetOrCreateDynamicDim(const std::string& name); + + const codegen::CodeGenHandle* GetCodeGenHandle() const { + return handle_; + } + + std::string CreateUnnamedSymbol(); + + protected: + std::unordered_map dynamic_dims_; + + const codegen::CodeGenHandle* handle_; + + int unname_symbol_counter_; +}; + +// Add Promote for CodeGenContext +DYNAMIC_PROMOTE(CodeGenContext) + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/all_ops.h b/onnxruntime/core/codegen/target/generic/op_ir_creator/all_ops.h new file mode 100644 index 0000000000000..eeba129a7d597 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/all_ops.h @@ -0,0 +1,47 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/codegen/target/codegen_context.h" +#include "core/codegen/common/op_macro.h" +#include "core/codegen/target/tvm_op_creator.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// This macro declares a TVM IR builder +// based on ORT OP type with postfix DefaultTVM +#define DECLARE_GENERIC_OP_IR_CREATOR_CLASS(OP) \ + DECLARE_OP_IR_CREATOR_CLASS(OP, DefaultTVM) + +// This macro returns a TVM IR builder class name +// based ORT OP type with postfix DefaultTVM +#define GENERIC_OP_IR_CREATOR_CLASS(OP) \ + CREATOR_CLASS(OP, DefaultTVM##IRCreator) + +#define GENERIC_OP_IR_CREATOR_STRING(OP) \ + STRINGIZE(GENERIC_OP_IR_CREATOR_CLASS(OP)) + +// define all ops for DefaultTVM +#define ADD_OP_ITEM(OP) DECLARE_GENERIC_OP_IR_CREATOR_CLASS(OP) +#define BINARY_OP(OP) ADD_OP_ITEM(OP) +#define BINARY_CMP_OP(OP) ADD_OP_ITEM(OP) +#define POOL_OP(OP) ADD_OP_ITEM(OP) +#define UNARY_OP(OP) ADD_OP_ITEM(OP) +#define VARIADIC_OP(OP) ADD_OP_ITEM(OP) +#define REDUCE_INDEXED_OP(OP) ADD_OP_ITEM(OP) +#define REDUCE_OP(OP) ADD_OP_ITEM(OP) + +LIST_ALL_GENERIC_OPS() + +#undef ADD_OP_ITEM +#undef BINARY_OP +#undef BINARY_CMP_OP +#undef POOL_OP +#undef REDUCE_OP +#undef REDUCE_INDEXED_OP +#undef UNARY_OP +#undef VARIADIC_OP + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/binary_ops.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/binary_ops.cc new file mode 100644 index 0000000000000..771b2f18d1dc8 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/binary_ops.cc @@ -0,0 +1,46 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/common/op_macro.h" +#include "core/codegen/mti/math/binary_ops.h" +#include "core/codegen/mti/tensor/cast_ops.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// helper local macro defines Evaluate of BINARY_OP OpIRCreators +#define BINARY_OP(name) \ + Status GENERIC_OP_IR_CREATOR_CLASS(name)::Evaluate( \ + const tvm::Array& inputs, \ + const Node& node, \ + CodeGenContext&, \ + tvm::Array& outputs) { \ + tvm::Tensor Y = name(inputs[0], inputs[1], node.Name()); \ + outputs.push_back(Y); \ + return Status::OK(); \ + } + +LIST_BINARY_OPS() + +#undef BINARY_OP + +// helper local macro defines Evaluate of BINARY_CMP_OP OpIRCreators +#define BINARY_CMP_OP(name) \ + Status GENERIC_OP_IR_CREATOR_CLASS(name)::Evaluate( \ + const tvm::Array& inputs, \ + const Node& node, \ + CodeGenContext&, \ + tvm::Array& outputs) { \ + tvm::Tensor Y = Cast(name(inputs[0], inputs[1], node.Name()), HalideIR::UInt(8), "cast_bool_" #name); \ + outputs.push_back(Y); \ + return Status::OK(); \ + } + +LIST_BINARY_CMP_OPS() + +#undef BINARY_CMP_OP + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/clip.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/clip.cc new file mode 100644 index 0000000000000..3551af6682828 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/clip.cc @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/math/unary_ops.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of Clip OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Clip)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext&, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper info(&ctx); + + float max_value, min_value; + ORT_RETURN_IF_ERROR(info.GetAttr("max", &max_value)); + ORT_RETURN_IF_ERROR(info.GetAttr("min", &min_value)); + + tvm::Tensor Y = Clip(inputs[0], min_value, max_value, node.Name() + "_Clip"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/gemm.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/gemm.cc new file mode 100644 index 0000000000000..b8628511e6308 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/gemm.cc @@ -0,0 +1,39 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/math/gemm.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of Gemm OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Gemm)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext& /*ctx_codegen*/, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper attrs(&ctx); + + tvm::Tensor A = inputs[0]; + tvm::Tensor B = inputs[1]; + tvm::Tensor C = inputs[2]; + + int64_t trans_A, trans_B; + ORT_RETURN_IF_ERROR(attrs.GetAttr("transA", &trans_A)); + ORT_RETURN_IF_ERROR(attrs.GetAttr("transB", &trans_B)); + + float alpha, beta; + ORT_ENFORCE(attrs.GetAttr("alpha", &alpha).IsOK()); + ORT_ENFORCE(attrs.GetAttr("beta", &beta).IsOK()); + + tvm::Tensor Y = Gemm(A, B, C, trans_A != 0, trans_B != 0, alpha, beta, node.Name() + "_Gemm"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/logsoftmax.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/logsoftmax.cc new file mode 100644 index 0000000000000..4543e7b6f6f20 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/logsoftmax.cc @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/math/logsoftmax.h" +#include "core/framework/op_kernel_info.h" +#include "core/providers/common.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of LogSoftmax OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(LogSoftmax)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext& ctx_codegen, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper info(&ctx); + + int64_t axis_i64; + ORT_RETURN_IF_ERROR(info.GetAttr("axis", &axis_i64)); + axis_i64 = HandleNegativeAxis(axis_i64, gsl::narrow_cast(inputs[0]->shape.size())); + + tvm::Tensor Y = LogSoftmax(inputs[0], axis_i64, node.Name() + "_LogSoftmax"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/matmul.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/matmul.cc new file mode 100644 index 0000000000000..6dd2bb840e067 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/matmul.cc @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/math/matmul_ops.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of MatMul OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(MatMul)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext&, + tvm::Array& outputs) { + tvm::Tensor Y = MatMul(inputs[0], inputs[1], node.Name() + "_MatMul"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/quantize/matmul_integer.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/quantize/matmul_integer.cc new file mode 100644 index 0000000000000..df67c3cb8ca7a --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/quantize/matmul_integer.cc @@ -0,0 +1,37 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/math/binary_ops.h" +#include "core/codegen/mti/math/matmul_ops.h" +#include "core/codegen/mti/tensor/cast_ops.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of MatMulInteger OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(MatMulInteger)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext& ctx_codegen, + tvm::Array& outputs) { + const auto& lhs_tensor = inputs[0]; + const auto& rhs_tensor = inputs[1]; + auto& name = node.Name(); + + // A generic path, cast to int32 + // Support skipped trailing inputs + auto lhs = (node.InputDefs().size() >= 3 && node.InputDefs()[2]->Exists()) + ? Sub(Cast(lhs_tensor, HalideIR::Int(32)), Cast(inputs[2], HalideIR::Int(32))) + : Cast(lhs_tensor, HalideIR::Int(32)); + auto rhs = (node.InputDefs().size() >= 4 && node.InputDefs()[3]->Exists()) + ? Sub(Cast(rhs_tensor, HalideIR::Int(32)), Cast(inputs[3], HalideIR::Int(32))) + : Cast(rhs_tensor, HalideIR::Int(32)); + tvm::Tensor Y = MatMul(lhs, rhs, name + "_MatMulInteger"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/reduce_ops.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/reduce_ops.cc new file mode 100644 index 0000000000000..1773dd11fbcc3 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/reduce_ops.cc @@ -0,0 +1,111 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/common/op_macro.h" +#include "core/codegen/mti/math/reduce_ops.h" +#include "core/codegen/mti/tensor/cast_ops.h" +#include "core/codegen/mti/tensor/reshape_ops.h" +#include "core/framework/op_kernel_info.h" +#include "core/providers/common.h" + +namespace onnxruntime { +namespace tvm_codegen { + +using ReduceIndexedFunc = tvm::Tensor (*)(const tvm::Tensor& X, int64_t axis, bool keep_dims, const std::string& name); +using ReduceFunc = tvm::Tensor (*)(const tvm::Tensor& X, const std::vector& axes, bool keep_dims, const std::string& name); + +// helper class for for REDUCE_INDEXED_OP +class FuncReduceIndexed { + public: + FuncReduceIndexed(const Node& node, ReduceIndexedFunc func, const std::string& name) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper info(&ctx); + axis_ = info.GetAttrOrDefault("axis", 0); + int64_t keepdims_i = 1; + ORT_ENFORCE(info.GetAttr("keepdims", &keepdims_i).IsOK()); + keep_dims_ = (keepdims_i == 1); + func_ = func; + name_ = name; + } + + tvm::Tensor operator()(const tvm::Tensor& X) const { + auto axis = HandleNegativeAxis(axis_, gsl::narrow_cast(X->shape.size())); + tvm::Tensor index32 = func_(X, axis, keep_dims_, name_); + return Cast(index32, tvm::Int(64)); + } + + private: + int64_t axis_; + bool keep_dims_; + ReduceIndexedFunc func_; + std::string name_; +}; + +// helper class for REDUCE_OP +class FuncReduce { + public: + FuncReduce(const Node& node, ReduceFunc func, const std::string& name) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper info(&ctx); + axes_ = info.GetAttrsOrDefault("axes"); + int64_t keepdims_i = 1; + ORT_ENFORCE(info.GetAttr("keepdims", &keepdims_i).IsOK()); + keep_dims_ = (keepdims_i == 1); + func_ = func; + name_ = name; + } + + tvm::Tensor operator()(const tvm::Tensor& X) const { + std::vector axes; + for (auto i : axes_) + axes.push_back(HandleNegativeAxis(i, gsl::narrow_cast(X->shape.size()))); + + return func_(X, axes, keep_dims_, name_); + } + + private: + std::vector axes_; + bool keep_dims_; + ReduceFunc func_; + std::string name_; +}; + +// helper macro defines Evaluate of REDUCE_OP OpIRCreators +#define REDUCE_OP(name) \ + Status GENERIC_OP_IR_CREATOR_CLASS(name)::Evaluate( \ + const tvm::Array& inputs, \ + const Node& node, \ + CodeGenContext&, \ + tvm::Array& outputs) { \ + tvm::Tensor Y; \ + if (ShapeRank(node.OutputDefs()[0]) == 0) { \ + tvm::Tensor temp = FuncReduce(node, &name, #name)(inputs[0]); \ + Y = Reshape(temp, {}); \ + } else { \ + Y = FuncReduce(node, &name, #name)(inputs[0]); \ + } \ + outputs.push_back(Y); \ + return Status::OK(); \ + } + +// helper macro defines Evaluate of REDUCE_INDEXED_OP OpIRCreators +#define REDUCE_INDEXED_OP(name) \ + Status GENERIC_OP_IR_CREATOR_CLASS(name)::Evaluate( \ + const tvm::Array& inputs, \ + const Node& node, \ + CodeGenContext&, \ + tvm::Array& outputs) { \ + tvm::Tensor Y = FuncReduceIndexed(node, &name, #name)(inputs[0]); \ + outputs.push_back(Y); \ + return Status::OK(); \ + } + +LIST_REDUCE_OPS() + +#undef REDUCE_OP +#undef REDUCE_INDEXED_OP + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/softmax.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/softmax.cc new file mode 100644 index 0000000000000..b8ce32ca30f6e --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/softmax.cc @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/math/softmax.h" +#include "core/framework/op_kernel_info.h" +#include "core/providers/common.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of Softmax OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Softmax)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext& ctx_codegen, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper info(&ctx); + + int64_t axis_i64; + ORT_RETURN_IF_ERROR(info.GetAttr("axis", &axis_i64)); + + axis_i64 = HandleNegativeAxis(axis_i64, gsl::narrow_cast(inputs[0]->shape.size())); + tvm::Tensor Y = Softmax(inputs[0], axis_i64, node.Name() + "_Softmax"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/unary_ops.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/unary_ops.cc new file mode 100644 index 0000000000000..fc31097fc5453 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/unary_ops.cc @@ -0,0 +1,136 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/common/op_macro.h" +#include "core/codegen/mti/math/unary_ops.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// helper class for unary_ops with alpha +class FuncWithAlpha { + public: + FuncWithAlpha(const Node& node) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper attrs(&ctx); + ORT_ENFORCE(attrs.GetAttr("alpha", &alpha_).IsOK()); + } + + protected: + float alpha_; +}; + +// helper class for unary_ops with alpha and beta +class FuncWithAlphaBeta { + public: + FuncWithAlphaBeta(const Node& node) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper attrs(&ctx); + ORT_ENFORCE(attrs.GetAttr("alpha", &alpha_).IsOK()); + ORT_ENFORCE(attrs.GetAttr("beta", &beta_).IsOK()); + } + + protected: + float alpha_; + float beta_; +}; + +// helper class for unary_ops with alpha and gamma +class FuncWithAlphaGamma { + public: + FuncWithAlphaGamma(const Node& node) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper attrs(&ctx); + ORT_ENFORCE(attrs.GetAttr("alpha", &alpha_).IsOK()); + ORT_ENFORCE(attrs.GetAttr("gamma", &gamma_).IsOK()); + } + + protected: + float alpha_; + float gamma_; +}; + +// helper macro declares unary_ops helper class without attribute +#define FuncClass(name) \ + class Func##name { \ + public: \ + Func##name(const Node&) {} \ + tvm::Tensor operator()(const tvm::Tensor& X) const { \ + return name(X); \ + } \ + } + +// helper macro declares unary_ops helper class with alpha +#define FuncClassAlpha(name) \ + class Func##name : public FuncWithAlpha { \ + public: \ + Func##name(const Node& node) : FuncWithAlpha(node) {} \ + tvm::Tensor operator()(const tvm::Tensor& X) const { \ + return name(X, alpha_); \ + } \ + } + +// helper macro declares unary_ops helper class with alpha and beta +#define FuncClassAlphaBeta(name) \ + class Func##name : public FuncWithAlphaBeta { \ + public: \ + Func##name(const Node& node) : FuncWithAlphaBeta(node) {} \ + tvm::Tensor operator()(const tvm::Tensor& X) const { \ + return name(X, alpha_, beta_); \ + } \ + } + +// helper macro declares unary_ops helper class with alpha and gamma +#define FuncClassAlphaGamma(name) \ + class Func##name : public FuncWithAlphaGamma { \ + public: \ + Func##name(const Node& node) : FuncWithAlphaGamma(node) {} \ + tvm::Tensor operator()(const tvm::Tensor& X) const { \ + return name(X, alpha_, gamma_); \ + } \ + } + +FuncClass(Abs); +FuncClassAlphaBeta(Affine); +FuncClass(Ceil); +FuncClassAlpha(Elu); +FuncClass(Exp); +FuncClass(Floor); +FuncClassAlphaBeta(HardSigmoid); +FuncClassAlpha(LeakyRelu); +FuncClass(Log); +FuncClass(Neg); +FuncClassAlphaBeta(ParametricSoftplus); +FuncClass(Reciprocal); +FuncClass(Relu); +FuncClassAlphaBeta(ScaledTanh); +FuncClassAlphaGamma(Selu); +FuncClass(Sigmoid); +FuncClass(Softplus); +FuncClass(Softsign); +FuncClass(Sqrt); +FuncClass(Tanh); +FuncClassAlpha(ThresholdedRelu); + +// helper macro defines Evaluate of UNARY_OP OpIRCreators +#define UNARY_OP(name) \ + Status GENERIC_OP_IR_CREATOR_CLASS(name)::Evaluate( \ + const tvm::Array& inputs, \ + const Node& node, \ + CodeGenContext&, \ + tvm::Array& outputs) { \ + tvm::Tensor Y = Func##name(node)(inputs[0]); \ + outputs.push_back(Y); \ + return Status::OK(); \ + } + +// helper local macros to replace some calls in LIST_UNARY_OPS +LIST_UNARY_OPS() + +#undef UNARY_OP + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/variadic_ops.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/variadic_ops.cc new file mode 100644 index 0000000000000..ff719b82f85e5 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/math/variadic_ops.cc @@ -0,0 +1,36 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/math/binary_ops.h" +#include "core/codegen/mti/tensor/reshape_ops.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Sum(const tvm::Tensor& lhs, const tvm::Tensor& rhs, const std::string& name) { + return Add(lhs, rhs, name); +} + +// helper local macro defines Evaluate of BINARY_OP OpIRCreators +#define VARIADIC_OP(name) \ + Status GENERIC_OP_IR_CREATOR_CLASS(name)::Evaluate( \ + const tvm::Array& inputs, \ + const Node& node, \ + CodeGenContext&, \ + tvm::Array& outputs) { \ + tvm::Tensor Y = Identity(inputs[0], node.Name() + "0"); \ + for (size_t i = 1; i < inputs.size(); ++i) \ + Y = name(Y, inputs[i], node.Name() + std::to_string(i)); \ + outputs.push_back(Y); \ + return Status::OK(); \ + } + +LIST_VARIADIC_OPS() + +#undef VARIADIC_OP + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/conv.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/conv.cc new file mode 100644 index 0000000000000..31ca33df8f335 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/conv.cc @@ -0,0 +1,131 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/nn/conv_ops.h" +#include "core/codegen/mti/mti_tvm_utils.h" +#include "core/codegen/mti/tensor/concat_ops.h" +#include "core/codegen/mti/tensor/split.h" +#include "core/codegen/target/ort_tvm_utils.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +Status GENERIC_OP_IR_CREATOR_CLASS(Conv)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext& ctx_codegen, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper info(&ctx); + + // Attributes + int64_t group; + std::string auto_pad; + std::vector kernel_shape, strides, dilations, pads; + + info.GetAttrOrDefault("group", &group, 1); + info.GetAttrOrDefault("auto_pad", &auto_pad, "NOTSET"); + + ORT_ENFORCE(info.GetAttrs("kernel_shape", kernel_shape).IsOK()); + ORT_ENFORCE(kernel_shape.size() <= 2, "Only support 1D/2D convolution currently!"); + ORT_ENFORCE(info.GetAttrs("strides", strides).IsOK()); + + dilations = info.GetAttrs("dilations", dilations).IsOK() ? dilations : std::vector(kernel_shape.size(), 1); + ORT_ENFORCE(dilations == std::vector(kernel_shape.size(), 1), "Only support dilation is 1 currently"); + + pads = info.GetAttrs("pads", pads).IsOK() ? pads : std::vector(kernel_shape.size() * 2, 0); + + // auto_pad + if (auto_pad != "NOTSET") { + auto rank = inputs[0]->shape.size() - 2; + ORT_ENFORCE(rank > 0); + for (uint64_t i = 0; i < rank; i++) { + if (auto_pad == "VALID") { + pads[i] = 0; + pads[i + rank] = 0; + } else if (auto_pad == "SAME_UPPER" || auto_pad == "SAME_LOWER") { + // TODO: handle symbolic dim + ORT_ENFORCE(ShapeHasValue(node.InputDefs()[0], 2 + i)); + + int64_t input_dim_value = ShapeValue(node.InputDefs()[0], 2 + i); + int64_t output_dim_value = (input_dim_value + strides[i] - 1) / strides[i]; + int64_t pad_needed = (output_dim_value - 1) * strides[i] + kernel_shape[i] - input_dim_value; + + pads[i] = auto_pad == "SAME_LOWER" ? (pad_needed + 1) / 2 : pad_needed / 2; + pads[i + rank] = pad_needed - pads[i]; + } else { + ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unknown auto_pad value"); + } + } + } + + // Inputs + tvm::Tensor X = inputs[0]; + tvm::Tensor W = inputs[1]; + // Outputs + tvm::Tensor Y; + tvm::Array Y_shape = ShapeToTvmArray(node.OutputDefs()[0], ctx_codegen); + + // 1-D convolution + if (kernel_shape.size() == 1) { + Y = Conv1D(X, W, Y_shape, ToTvmArray(strides), ToTvmArray(pads), node.Name() + "_Conv1D"); + } + // 2-D convolution + else if (kernel_shape.size() == 2) { + if (group == 1) { + Y = Conv2D(X, W, Y_shape, ToTvmArray(strides), ToTvmArray(pads), node.Name() + "_Conv2D"); + } else { + int64_t channel_out = ShapeValue(node.InputDefs()[1], 0); + int64_t channel_in = ShapeValue(node.InputDefs()[1], 1); + ORT_ENFORCE(channel_out % group == 0); + + int64_t cout_group = channel_out / group; + Y_shape.Set(1, Y_shape[1] / gsl::narrow_cast(group)); + + tvm::Array split_index0; + tvm::Array split_index1; + + for (int i = 1; i < group; i++) { + split_index0.push_back(i * channel_in); + split_index1.push_back(i * cout_group); + } + + auto input_groups = SplitWithIndices(X, split_index0, 1); + auto weight_groups = SplitWithIndices(W, split_index1, 0); + + // FIXME: This will trigger a llvm buffer overflow when group is too large + // TODO: fix this change it to batched gemm/conv + tvm::Array output_tensors; + for (int i = 0; i < group; i++) { + auto output_tensor = Conv2D(input_groups[i], + weight_groups[i], + Y_shape, + ToTvmArray(strides), + ToTvmArray(pads), + node.Name() + "_Conv2D"); + output_tensors.push_back(output_tensor); + } + Y = Concat(output_tensors, 1); + } + } + + // Add bias if provided + // Support skipped trailing inputs + if (node.InputDefs().size() > 2 && node.InputDefs()[2]->Exists()) { + tvm::Tensor B = inputs[2]; + Y = tvm::compute( + Y_shape, + [&](const tvm::Array& indices) { + return Y(indices) + B(indices[1]); + }); + } + + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/lstm.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/lstm.cc new file mode 100644 index 0000000000000..0f244a69a3205 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/lstm.cc @@ -0,0 +1,64 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/nn/lstm.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// In the cell computation, we don't have the "direction" dimention and sequence dimension, +// which have been processed outside of the cell. +// Here we implement an LTSM cell. +// For those args (inputs/outputs) of hidden states we put AFTER regular args (inputs/outputs) +// with a pre-defined order +// In a LSTM, the order is H and then C. +// Ouputs of LSTM is Y_h and then Y_c +Status GENERIC_OP_IR_CREATOR_CLASS(LSTM)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext& ctx_codegen, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper attrs(&ctx); + + std::string direction_attr; + ORT_RETURN_IF_ERROR(attrs.GetAttr("direction", &direction_attr)); + int64_t hidden_size; + ORT_RETURN_IF_ERROR(attrs.GetAttr("hidden_size", &hidden_size)); + + // input tensor with shape [seq_length, batch_size, input_size] + const tvm::Tensor& X = inputs[0]; // input tensor with shape [seq_length, batch_size, input_size] + const tvm::Tensor& W = inputs[1]; // weights tensor with shape [4*hidden_size, input_size] + const tvm::Tensor& R = inputs[2]; // recurrence tensor with shape [4*hidden_size, hidden_size] + const tvm::Tensor& B = inputs[3]; // optional bias tensor with shape [8*hidden_size] + bool has_B = node.InputDefs()[3]->Exists(); + + // Unsupported the 4th inputs + // optional tensor specifying sequence lengths in a batch, shape: [batch_size] + // const tvm::Tensor* seq_len = inputs[4] ? &inputs[4]->tensor : nullptr; + + const tvm::Tensor& prev_H = inputs[5]; // optional initial H, shape: [batch_size, hidden_size] + const tvm::Tensor& prev_C = inputs[6]; // optional initial C, shape: [batch_size, hidden_size] + + const tvm::Tensor& P = inputs[7]; // optional peepholes tensor with shape [3*hidde_size] + bool has_P = node.InputDefs()[7]->Exists(); + + tvm::Tensor Y_h; // shape: [batch_size, hidden_size] + tvm::Tensor Y_c; // shape: [batch_size, hidden_size] + LSTMAttributes lstm_attrs(hidden_size); + LSTM_cell(lstm_attrs, X, W, R, B, has_B, prev_H, prev_C, P, has_P, Y_h, Y_c); + + // Since we only generate lstm cell, lstm's states need to be always outputs, + // regardless whethere they are skipped or not. + // The skipped trailing outputs need to be handled by Execution + outputs.push_back(Y_h); + outputs.push_back(Y_c); + + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/pool_ops.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/pool_ops.cc new file mode 100644 index 0000000000000..e14a3f1bfb9f8 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/pool_ops.cc @@ -0,0 +1,93 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include "core/codegen/mti/nn/pool_ops.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// helper class for pool_ops with arguments +class FuncWithPoolingArgument { + public: + FuncWithPoolingArgument(const Node& node, const std::string& op_name) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper info(&ctx); + int64_t storage_order{0}; // MaxPool_8 only. 0 is row major, and 1 is column major. Default is 0. + + ORT_ENFORCE(info.GetAttrs("kernel_shape", kernel_shape_).IsOK(), "No kernel shape is set."); + if (kernel_shape_.size() != 2) + ORT_NOT_IMPLEMENTED(kernel_shape_.size(), "d pooling is not implementated"); + if (!info.GetAttrs("pads", pads_).IsOK() || pads_.empty()) { + pads_.resize(kernel_shape_.size() * 2, 0); + } + if (!info.GetAttrs("strides", strides_).IsOK() || strides_.empty()) { + strides_.resize(kernel_shape_.size(), 1); + } + if (op_name == "AveragePool") { + int64_t temp; + ORT_ENFORCE(info.GetAttr("count_include_pad", &temp).IsOK()); + count_include_pad_ = (temp != 0); + } + + if (op_name == "MaxPool") { + // TODO: add version check or not? remove version check since only after version 8 would have storage_order, otherwise, it would be zero + storage_order = info.GetAttrOrDefault("storage_order", 0 /*default_value*/); + if (storage_order != 1) { + layout_ = "NCWH"; + } + } + } + + std::vector kernel_shape_; + std::vector pads_; + std::vector strides_; + std::string layout_ = "NCHW"; + bool count_include_pad_ = false; +}; + +// A local macro to create Pool Ops + +// helper macro defines Evaluate of of POOL_OP OpIRCreators +#define POOL_OP(name) \ + Status GENERIC_OP_IR_CREATOR_CLASS(name)::Evaluate( \ + const tvm::Array& inputs, \ + const Node& node, \ + CodeGenContext&, \ + tvm::Array& outputs) { \ + if (outputs.size() > 1) ORT_NOT_IMPLEMENTED("output size = 2 is not implementated"); \ + FuncWithPoolingArgument argment(node, #name); \ + tvm::Tensor Y = name(inputs[0], ToTvmArray(argment.kernel_shape_), ToTvmArray(argment.strides_), ToTvmArray(argment.pads_), argment.layout_, argment.count_include_pad_); \ + outputs.push_back(Y); \ + return Status::OK(); \ + } // namespace tvm_codegen + +POOL_OP(MaxPool) +POOL_OP(AveragePool) + +#undef POOL_OP + +// helper macro defines Evaluate of of GlobalPOOL_OP OpIRCreators +#define POOL_OP(name) \ + Status GENERIC_OP_IR_CREATOR_CLASS(name)::Evaluate( \ + const tvm::Array& inputs, \ + const Node& node, \ + CodeGenContext&, \ + tvm::Array& outputs) { \ + if (inputs[0]->shape.size() != 4) \ + ORT_NOT_IMPLEMENTED(gsl::narrow_cast(inputs[0]->shape.size()) - 2, "d global pooling is not implementated"); \ + tvm::Tensor Y = name(inputs[0], "NCHW"); \ + outputs.push_back(Y); \ + return Status::OK(); \ + } + +POOL_OP(GlobalMaxPool) +POOL_OP(GlobalAveragePool) + +#undef POOL_OP + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/cast.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/cast.cc new file mode 100644 index 0000000000000..a3b9f5305d984 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/cast.cc @@ -0,0 +1,40 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/tensor/cast_ops.h" +#include "core/codegen/target/ort_tvm_utils.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of Cast OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Cast)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext&, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper attrs(&ctx); + + int64_t to; + ORT_RETURN_IF_ERROR(attrs.GetAttr("to", &to)); + auto to_type_proto = gsl::narrow_cast(to); + + tvm::Tensor X = inputs[0]; + tvm::Tensor Y; + if (to_type_proto == ONNX_NAMESPACE::TensorProto_DataType_BOOL) { + // special case for bool as ONNX bool is uint8, while in tvm it's uint1 + Y = CastToUInt8Bool(X, node.Name() + "_Cast"); + } else { + Y = Cast(X, ToTvmType(to_type_proto), node.Name() + "_Cast"); + } + + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/concat.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/concat.cc new file mode 100644 index 0000000000000..c04603266eaa2 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/concat.cc @@ -0,0 +1,30 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/tensor/concat_ops.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of Concat OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Concat)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext&, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper info(&ctx); + + int64_t axis; + ORT_RETURN_IF_ERROR(info.GetAttr("axis", &axis)); + + tvm::Tensor Y = Concat(inputs, axis, node.Name() + "_Concat"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/crop.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/crop.cc new file mode 100644 index 0000000000000..a14bba5f9f8a8 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/crop.cc @@ -0,0 +1,45 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include "core/codegen/mti/tensor/crop.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of Crop OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Crop)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext&, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper attrs(&ctx); + + if (inputs[0]->shape.size() != 4) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input is expected to have four dimensions corresponding to [N,C,H,W]"); + } + + std::vector border; + std::vector scale; + + ORT_ENFORCE(attrs.GetAttrs("border", border).IsOK()); + // scale is optional and status is false when omit + attrs.GetAttrs("scale", scale); + + if (border.size() != 4) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Attribute border needs to be specified with four border elements"); + } + + tvm::Tensor Y = Crop(inputs[0], ToTvmArray(border), ToTvmArray(scale), node.Name() + "_Crop"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/gather.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/gather.cc new file mode 100644 index 0000000000000..7a314b572b49e --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/gather.cc @@ -0,0 +1,30 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/tensor/gather.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of Gather OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Gather)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext&, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper attrs(&ctx); + + int64_t axis; + ORT_ENFORCE(attrs.GetAttr("axis", &axis).IsOK()); + + tvm::Tensor Y = Gather(inputs[0], axis, inputs[1], node.Name() + "_Gather"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/pad.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/pad.cc new file mode 100644 index 0000000000000..159825b0dc088 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/pad.cc @@ -0,0 +1,49 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include "core/codegen/mti/tensor/pad_ops.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of Pad OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Pad)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext&, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper attrs(&ctx); + + std::string mode; + std::vector pads; + float value; + + ORT_ENFORCE(attrs.GetAttr("mode", &mode).IsOK()); + ORT_ENFORCE(attrs.GetAttrs("pads", pads).IsOK()); + ORT_ENFORCE(attrs.GetAttr("value", &value).IsOK()); + + if (mode != "constant" && mode != "edge" && mode != "reflect") + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Pad: Unsupported padding mode!"); + + if (pads.size() != 2 * inputs[0]->shape.size()) + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Pad: pads rank does not match inputs rank!"); + + std::vector pad_before, pad_after; + size_t offset = pads.size() / 2; + for (size_t i = 0; i < offset; i++) { + pad_before.push_back(pads[i]); + pad_after.push_back(pads[i + offset]); + } + + tvm::Tensor Y = Pad(inputs[0], ToTvmArray(pad_before), ToTvmArray(pad_after), value, mode, node.Name() + "_Pad"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/reshape_ops.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/reshape_ops.cc new file mode 100644 index 0000000000000..a45d82475f2a7 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/reshape_ops.cc @@ -0,0 +1,96 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include "core/codegen/mti/tensor/reshape_ops.h" +#include "core/codegen/target/ort_tvm_utils.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of Dropout OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Dropout)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext& ctx_codegen, + tvm::Array& outputs) { + tvm::Tensor Y = Identity(inputs[0]); + outputs.push_back(Y); + + // optional mask + // Support skipped trailing outputs + if (node.OutputDefs().size() > 1 && node.OutputDefs()[1]->Exists()) { + // A fake mask with all zeros + tvm::Tensor mask = MakeZeroTensor(inputs[0]->shape, inputs[0]->dtype, "mask"); + outputs.push_back(mask); + } + + return Status::OK(); +} + +// Evaluate of Flatten OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Flatten)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext&, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper attrs(&ctx); + + int64_t axis; + ORT_RETURN_IF_ERROR(attrs.GetAttr("axis", &axis)); + + tvm::Tensor Y = Flatten(inputs[0], axis, node.Name() + "_Flatten"); + outputs.push_back(Y); + return Status::OK(); +} + +// Evaluate of Identity OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Identity)::Evaluate( + const tvm::Array& inputs, + const Node&, + CodeGenContext&, + tvm::Array& outputs) { + tvm::Tensor Y = Identity(inputs[0]); + outputs.push_back(Y); + return Status::OK(); +} + +// Evaluate of Reshape OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Reshape)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext& ctx_codegen, + tvm::Array& outputs) { + tvm::Tensor Y = Reshape(inputs[0], ShapeToTvmArray(node.OutputDefs()[0], ctx_codegen), node.Name() + "_Reshape"); + outputs.push_back(Y); + return Status::OK(); +} + +// Evaluate of Squeeze OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Squeeze)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext& ctx_codegen, + tvm::Array& outputs) { + tvm::Tensor Y = Reshape(inputs[0], ShapeToTvmArray(node.OutputDefs()[0], ctx_codegen), node.Name() + "_Squeeze"); + outputs.push_back(Y); + return Status::OK(); +} + +// Evaluate of Unsqueeze OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Unsqueeze)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext& ctx_codegen, + tvm::Array& outputs) { + tvm::Tensor Y = Reshape(inputs[0], ShapeToTvmArray(node.OutputDefs()[0], ctx_codegen), node.Name() + "_Unsqueeze"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/slice.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/slice.cc new file mode 100644 index 0000000000000..2915dcc701cd1 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/slice.cc @@ -0,0 +1,107 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/target/ort_tvm_utils.h" +#include "core/codegen/mti/mti_tvm_utils.h" +#include "core/codegen/mti/tensor/slice.h" +#include "core/framework/op_kernel_info.h" + +#include + +namespace onnxruntime { +namespace tvm_codegen { + +// local constexpr for INT_MAX +constexpr int64_t max_range = INT_MAX; + +Status SliceCommon(const tvm::Array& inputs, + const Node& node, + tvm::Array& outputs, + const std::vector& starts, + const std::vector& ends, + const std::vector& axes1) { + ORT_RETURN_IF_NOT(nullptr != node.InputDefs()[0]); + const ONNX_NAMESPACE::TensorShapeProto* shape_proto = node.InputDefs()[0]->Shape(); + + std::vector axes; + if (axes1.size() == 0) { + for (size_t i = 0; i < starts.size(); ++i) { + axes.push_back(gsl::narrow_cast(i)); + } + } else { + axes = axes1; + } + + tvm::Array tvm_starts, tvm_ends; + bool empty = false; + + for (int dim = 0; dim < shape_proto->dim_size(); ++dim) { + auto axes_iter = std::find(axes.begin(), axes.end(), dim); + const ONNX_NAMESPACE::TensorShapeProto_Dimension& proto_dim = shape_proto->dim(dim); + bool found_in_axes = (axes_iter != axes.end()); + if (!found_in_axes) { + tvm_starts.push_back(0); + if (proto_dim.has_dim_value()) { + tvm_ends.push_back(proto_dim.dim_value()); + } else { + tvm_ends.push_back(max_range); + } + } else { + auto axes_index = axes_iter - axes.begin(); + int64_t start = starts[axes_index]; + int64_t end = ends[axes_index]; + if (proto_dim.has_dim_value()) { + int64_t dim_max = proto_dim.dim_value(); + if (start < 0) start += dim_max; + if (end < 0) end += dim_max; + start = std::min(dim_max, std::max(static_cast(0), start)); + end = std::min(dim_max, std::max(start, end)); + } + tvm_starts.push_back(start); + tvm_ends.push_back(end); + empty = empty || (start == end); + } + } + + tvm::Tensor Y; + if (empty) { + tvm::Array shape; + for (size_t dim = 0; dim < gsl::narrow_cast(shape_proto->dim_size()); ++dim) { + shape.push_back(tvm::ir::Simplify(tvm_ends[dim] - tvm_starts[dim])); + } + Y = MakeZeroTensor(shape, inputs[0]->dtype, node.Name() + "_zeros"); + } else { + Y = Slice(inputs[0], tvm_starts, tvm_ends, node.Name() + "_Slice"); + } + + outputs.push_back(Y); + return Status::OK(); +} + +// Evaluate of Slice OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Slice)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext& ctx_codegen, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper info(&ctx); + + // NOTE that in opset 10, Slice has changed starts/ends/axes from attribute to input + // which may lead to dynamic output shape. + int version = ctx_codegen.GetCodeGenHandle()->domain_version_lookup_func(node.Domain()); + ORT_RETURN_IF_NOT(version <= 9, "Dynamic Slice is not supported yet"); + + std::vector starts, ends; + ORT_RETURN_IF_ERROR(info.GetAttrs("starts", starts)); + ORT_RETURN_IF_ERROR(info.GetAttrs("ends", ends)); + ORT_RETURN_IF_NOT(starts.size() == ends.size()); + + auto axes = info.GetAttrsOrDefault("axes"); + + return SliceCommon(inputs, node, outputs, starts, ends, axes); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/split.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/split.cc new file mode 100644 index 0000000000000..5c82ad405e17a --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/split.cc @@ -0,0 +1,65 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include "core/codegen/mti/tensor/split.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of Split OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Split)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext&, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper info(&ctx); + + int64_t axis; + ORT_RETURN_IF_ERROR(info.GetAttr("axis", &axis)); + axis = HandleNegativeAxis(axis, gsl::narrow_cast(inputs[0]->shape.size())); + std::vector split_sizes; + + int64_t split_size_sum = 0; + if (info.GetAttrs("split", split_sizes).IsOK()) { + // optional + split_size_sum = std::accumulate(split_sizes.cbegin(), split_sizes.cend(), 0LL); + ORT_RETURN_IF_NOT(std::all_of(split_sizes.cbegin(), split_sizes.cend(), [](int64_t value) { return value > 0; }), + "Invalid value in 'split' attribute. All values must be > 0"); + + // check split sizes + for (size_t i = 0; i < node.OutputDefs().size(); ++i) { + ORT_RETURN_IF_NOT(split_sizes[i] == ShapeValue(node.OutputDefs()[i], gsl::narrow(axis))); + } + + } else { + for (size_t i = 0; i < node.OutputDefs().size(); ++i) { + split_sizes.push_back(ShapeValue(node.OutputDefs()[i], gsl::narrow(axis))); + split_size_sum += split_sizes[i]; + } + } + + // check total size + if (ShapeHasValue(node.InputDefs()[0], axis)) { + int64_t input_axis_dim = ShapeValue(node.InputDefs()[0], axis); + if (split_size_sum != input_axis_dim) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Cannot split using values in 'split' attribute. Axis=", axis, + " Dim being splitted=", input_axis_dim, + " Sum of sizes in 'split' (must equal size of selected axis) was ", split_size_sum); + } + } + + tvm::Array output_tensors = Split(inputs[0], ToTvmArray(split_sizes), axis, node.Name() + "_Split"); + for (size_t i = 0; i < node.OutputDefs().size(); ++i) { + outputs.push_back(output_tensors[i]); + } + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/transpose.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/transpose.cc new file mode 100644 index 0000000000000..d47ea30084aec --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/transpose.cc @@ -0,0 +1,46 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include "core/codegen/mti/tensor/transpose.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of Transpose OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Transpose)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext&, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper attrs(&ctx); + + size_t input_0_shape_rank = inputs[0]->shape.size(); + std::vector permute; + attrs.GetAttrs("perm", permute); + if (permute.size() != 0 && permute.size() != input_0_shape_rank) + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Transpose: Incorrect permute size"); + + std::vector default_permute; + const std::vector* perm; + if (permute.size() > 0) { + perm = &permute; + } else { + default_permute.resize(input_0_shape_rank); + for (size_t i = 0; i < input_0_shape_rank; ++i) { + default_permute[i] = gsl::narrow(input_0_shape_rank - 1 - i); + } + perm = &default_permute; + } + + tvm::Tensor Y = Transpose(inputs[0], ToTvmArrayInt(*perm), node.Name() + "_Transpose"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/where.cc b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/where.cc new file mode 100644 index 0000000000000..e20df5f3b0382 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/where.cc @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" + +#include "core/codegen/mti/mti_tvm_utils.h" +#include "core/codegen/mti/tensor/where.h" +#include "core/framework/op_kernel_info.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Evaluate of Transpose OpIRCreator +Status GENERIC_OP_IR_CREATOR_CLASS(Where)::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext&, + tvm::Array& outputs) { + ProtoHelperNodeContext ctx(node); + OpNodeProtoHelper attrs(&ctx); + + tvm::Tensor Y = Where(inputs[0], inputs[1], inputs[2], node.Name() + "_Where"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/scheduler/all_schedules.h b/onnxruntime/core/codegen/target/generic/scheduler/all_schedules.h new file mode 100644 index 0000000000000..54177009c6b72 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/scheduler/all_schedules.h @@ -0,0 +1,20 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/codegen/target/tvm_scheduler.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// AlwaysRoot is for debug purpose +DECLARE_TVM_SCHEDULER_CLASS(AlwaysRoot, GenericTVMRule) +// Create schedule for TVM Rule +DECLARE_TVM_SCHEDULER_CLASS(Extern, GenericTVMRule) +DECLARE_TVM_SCHEDULER_CLASS(Reduce, GenericTVMRule) + +//Crete scheduler for ORT OpType, Softmax +DECLARE_TVM_SCHEDULER_CLASS(Softmax, GenericOrtOpType) + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/scheduler/ort_type_schedule.cc b/onnxruntime/core/codegen/target/generic/scheduler/ort_type_schedule.cc new file mode 100644 index 0000000000000..4cd9efb313fc1 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/scheduler/ort_type_schedule.cc @@ -0,0 +1,22 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/scheduler/all_schedules.h" + +#include "core/codegen/target/generic/scheduler/schedule_utils.h" + +namespace onnxruntime { +namespace tvm_codegen { + +bool TVM_SCHEDULER_CLASS(Softmax, GenericOrtOpType)::Evaluate( + const tvm::Tensor& tensor, + const Node*, + CodeGenContext&, + ScheduleContext& ctx_sched) { + // compute root the exp since it is reused more than once + auto& tensor_exp = tensor->op->InputTensors()[0]; + return InsertRootSchedule(tensor_exp, ctx_sched); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.cc b/onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.cc new file mode 100644 index 0000000000000..319bdfed92ea4 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.cc @@ -0,0 +1,164 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/scheduler/schedule_utils.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// Check the schedule of tensor +// If it has no compute_root, Insert compute_root to tensor, and record it to ctx.scheduled_tensors +bool InsertRootSchedule( + const tvm::Tensor& tensor, + ScheduleContext& ctx) { + auto it = ctx.scheduled_tensors.find(tensor->op.get()); + if (it != ctx.scheduled_tensors.end()) { + if (it->second == ScheduleType::ScheduleClosure || + it->second == ScheduleType::ScheduleRoot) { + return false; + } + it->second = ScheduleType::ScheduleRoot; + } else { + ctx.scheduled_tensors.insert(std::make_pair(tensor->op.get(), ScheduleType::ScheduleRoot)); + } + ctx.schedule[tensor->op].compute_root(); + return true; +} + +// Check the schedule of tensor +// If it is not labeled as closure, lable it. +bool InsertClosure(const tvm::Tensor& tensor, + ScheduleContext& ctx) { + auto it = ctx.scheduled_tensors.find(tensor->op.get()); + if (it != ctx.scheduled_tensors.end()) { + if (it->second == ScheduleType::ScheduleClosure) + return false; + it->second = ScheduleType::ScheduleClosure; + } else { + ctx.scheduled_tensors.insert(std::make_pair(tensor->op.get(), ScheduleType::ScheduleClosure)); + } + return true; +} + +// Combination of InsertRootSchedule and InsertClosure +bool InsertRootScheduleAndClosure( + const tvm::Tensor& tensor, + ScheduleContext& ctx) { + auto it = ctx.scheduled_tensors.find(tensor->op.get()); + if (it != ctx.scheduled_tensors.end()) { + if (it->second == ScheduleType::ScheduleClosure) { + return false; + } + it->second = ScheduleType::ScheduleClosure; + } else { + ctx.scheduled_tensors.insert(std::make_pair(tensor->op.get(), ScheduleType::ScheduleClosure)); + } + ctx.schedule[tensor->op].compute_root(); + return true; +} + +// Check the schedule of tensor +// If it is not scheduled, try to vectorize it. +// Note TryVectorization has to use with compute_root. +// Therefore, there is a safty check of tensor's schedule +bool TryVectorization( + const tvm::Tensor& tensor, + int64_t natural_vector_size, + ScheduleContext& ctx) { + auto it = ctx.scheduled_tensors.find(tensor->op.get()); + if (it != ctx.scheduled_tensors.end()) { + if (it->second > ScheduleType::ScheduleInline) { + return false; + } + } + + auto shape = tensor->shape; + auto rank = shape.size(); + if (rank < 1) { + return false; + } + const int64_t* tail_dim = as_const_int(shape[rank - 1]); + + if (nullptr != tail_dim) { + auto extern_op = tensor->op.as(); + if (nullptr != extern_op) { + return false; + } + + auto compute_op = tensor->op.as(); + + if (nullptr != compute_op) { + auto axis = compute_op->axis; + tvm::IterVar x = axis[rank - 1]; + if ((*tail_dim) > natural_vector_size) { + if ((*tail_dim) % natural_vector_size == 0) { + tvm::IterVar xi, xo; + ctx.schedule[tensor->op].split(x, static_cast(natural_vector_size), &xo, &xi); + ctx.schedule[tensor->op].vectorize(xi); + return true; + } + } else if (*tail_dim > 0) { + // don't vectorize if dim is 0 + ctx.schedule[tensor->op].vectorize(x); + return true; + } + } + } + return false; +} + +// Check the schedule of tensor +// If it is not scheduled, try to add compute_inline on it. +// Note TryInlineSchedule cannot be used with compute_root. +// Therefore, there is a safty check of tensor's schedule. +bool TryInlineSchedule( + const tvm::Tensor& tensor, + ScheduleContext& ctx) { + auto it = ctx.scheduled_tensors.find(tensor->op.get()); + if (it != ctx.scheduled_tensors.end()) { + if ((int)it->second < (int)ScheduleType::ScheduleInline) { + ctx.schedule[tensor->op].compute_inline(); + it->second = ScheduleType::ScheduleInline; + return true; + } else { + return false; + } + } + ctx.schedule[tensor->op].compute_inline(); + ctx.scheduled_tensors.insert(std::make_pair(tensor->op.get(), ScheduleType::ScheduleInline)); + return true; +} + +// Check the schedule of tensor's inputs, and call InsertRootSchedule for each of them +bool InputRootSchedule( + const tvm::Tensor& tensor, + ScheduleContext& ctx) { + bool status = false; + for (auto& t : tensor->op->InputTensors()) { + if (t->op->InputTensors().size() > 0) { + bool status_root = InsertRootSchedule(t, ctx); + status = status || status_root; + } + } + return status; +} + +// Check the schedule of tensor's inputs, +// and call InsertRootSchedule and TryVectorization for each of them +bool InputRootScheduleWithVectorization( + const tvm::Tensor& tensor, + int64_t natural_vector_size, + ScheduleContext& ctx) { + bool status = false; + for (auto& t : tensor->op->InputTensors()) { + if (t->op->InputTensors().size() > 0) { + bool status_vec = TryVectorization(t, natural_vector_size, ctx); + bool status_root = InsertRootSchedule(t, ctx); + status = status || status_root || status_vec; + } + } + return status; +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.h b/onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.h new file mode 100644 index 0000000000000..f6e1bdb5c4657 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.h @@ -0,0 +1,60 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +// Check the schedule of tensor +// If it has no compute_root, Insert compute_root to tensor, +// and record it to ctx.scheduled_tensors +bool InsertRootSchedule( + const tvm::Tensor& tensor, + ScheduleContext& ctx); + +// Check the schedule of tensor +// If it is not labeled as closure, lable it. +bool InsertClosure( + const tvm::Tensor& tensor, + ScheduleContext& ctx); + +// Combination of InsertRootSchedule and InsertClosure +bool InsertRootScheduleAndClosure( + const tvm::Tensor& tensor, + ScheduleContext& ctx); + +// Check the schedule of tensor +// If it is not scheduled, try to vectorize it. +// Note TryVectorization has to use with compute_root. +// Therefore, there is a safty check of tensor's schedule +bool TryVectorization( + const tvm::Tensor& tensor, + int64_t natural_vector_size, + ScheduleContext& ctx); + +// Check the schedule of tensor +// If it is not scheduled, try to add compute_inline on it. +// Note TryInlineSchedule cannot be used with compute_root. +// Therefore, there is a safty check of tensor's schedule. +bool TryInlineSchedule( + const tvm::Tensor& tensor, + ScheduleContext& ctx); + +// Check the schedule of tensor's inputs, +// and call InsertRootSchedule for each of them +bool InputRootSchedule( + const tvm::Tensor& tensor, + ScheduleContext& ctx); + +// Check the schedule of tensor's inputs, +// and call InsertRootSchedule and TryVectorization for each of them +bool InputRootScheduleWithVectorization( + const tvm::Tensor& tensor, + int64_t natural_vector_size, + ScheduleContext& ctx); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/scheduler/tvm_rule_schedule.cc b/onnxruntime/core/codegen/target/generic/scheduler/tvm_rule_schedule.cc new file mode 100644 index 0000000000000..9f882e6c0b5ad --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/scheduler/tvm_rule_schedule.cc @@ -0,0 +1,41 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/scheduler/all_schedules.h" + +#include "core/codegen/target/generic/scheduler/schedule_utils.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// This is for debug +bool TVM_SCHEDULER_CLASS(AlwaysRoot, GenericTVMRule)::Evaluate( + const tvm::Tensor& tensor, + const Node*, + CodeGenContext&, + ScheduleContext& ctx_sched) { + return InsertRootSchedule(tensor, ctx_sched); +} + +// For External tvm::Tensor +bool TVM_SCHEDULER_CLASS(Extern, GenericTVMRule)::Evaluate( + const tvm::Tensor& tensor, + const Node*, + CodeGenContext&, + ScheduleContext& ctx_sched) { + bool status = InsertRootScheduleAndClosure(tensor, ctx_sched); + bool status_input = InputRootSchedule(tensor, ctx_sched); + return status || status_input; +} + +// For Reduce Compute tvm::Tensor +bool TVM_SCHEDULER_CLASS(Reduce, GenericTVMRule)::Evaluate( + const tvm::Tensor& tensor, + const Node*, + CodeGenContext&, + ScheduleContext& ctx_sched) { + return InsertRootScheduleAndClosure(tensor, ctx_sched); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.cc b/onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.cc new file mode 100644 index 0000000000000..8508fa3e40aa8 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.cc @@ -0,0 +1,65 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/weight_layout/transpose_2d.h" + +#include "core/codegen/target/codegen_context.h" + +namespace onnxruntime { +namespace tvm_codegen { + +constexpr auto local_layout_name = "transpose_2d"; + +const std::string WeightLayoutTranspose2D::GetKey( + ONNX_NAMESPACE::TensorProto_DataType proto_type) { + return WeightLayout::GetKey(local_layout_name, proto_type, 2, 0.0f); +} + +WeightLayoutTranspose2D::WeightLayoutTranspose2D( + ONNX_NAMESPACE::TensorProto_DataType proto_type) + : WeightLayout(local_layout_name, proto_type, 2, 0.0f) {} + +CoordTransFunc WeightLayoutTranspose2D::ToActual(const tvm::Tensor& /*X*/) const { + return [&](const tvm::Array& nominal_coord) { + ORT_ENFORCE(nominal_coord.size() == 2); + const auto& y = nominal_coord[0]; + const auto& x = nominal_coord[1]; + return tvm::Array{ + x, + y}; + }; +} + +CoordTransFunc WeightLayoutTranspose2D::ToNominal(const tvm::Tensor& /*X*/) const { + return [&](const tvm::Array& actual_coord) { + ORT_ENFORCE(actual_coord.size() == 2); + const auto& y = actual_coord[0]; + const auto& x = actual_coord[1]; + return tvm::Array{ + x, + y}; + }; +} + +tvm::Array WeightLayoutTranspose2D::ToActualShape(const tvm::Tensor& X) const { + tvm::Array new_shape = { + X->shape[1], + X->shape[0]}; + return new_shape; +} + +std::vector WeightLayoutTranspose2D::ToActualShape(const Tensor* X) const { + ORT_ENFORCE(X != nullptr); + auto old_shape = X->Shape().GetDims(); + + ORT_ENFORCE(old_shape.size() == 2); + + std::vector new_shape = { + old_shape[1], + old_shape[0]}; + + return new_shape; +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.h b/onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.h new file mode 100644 index 0000000000000..bfa192d1f432c --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.h @@ -0,0 +1,33 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/codegen/target/weight_layout.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +// WeightLayoutTranspose2D for transposing a 2D weight +// [W, H] => [H, W] +class WeightLayoutTranspose2D : public WeightLayout { + public: + static const std::string GetKey(ONNX_NAMESPACE::TensorProto_DataType proto_type); + + public: + WeightLayoutTranspose2D(ONNX_NAMESPACE::TensorProto_DataType proto_type); + + ~WeightLayoutTranspose2D() = default; + + CoordTransFunc ToNominal(const tvm::Tensor& X) const override; + CoordTransFunc ToActual(const tvm::Tensor& X) const override; + tvm::Array ToActualShape(const tvm::Tensor& X) const override; + std::vector ToActualShape(const Tensor* X) const override; + + private: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(WeightLayoutTranspose2D); +}; + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.cc b/onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.cc new file mode 100644 index 0000000000000..8792f753652aa --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.cc @@ -0,0 +1,77 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/generic/weight_layout/vertical_stripes_2d.h" + +#include "core/codegen/target/codegen_context.h" + +namespace onnxruntime { +namespace tvm_codegen { + +constexpr auto local_name_prefix = "vertical_stripe_2d_"; + +const std::string WeightLayoutVerticalStripe2D::GetKey( + ONNX_NAMESPACE::TensorProto_DataType proto_type, + int stripe_width) { + return WeightLayout::GetKey( + local_name_prefix + std::to_string(stripe_width), + proto_type, 2, 0.0f); +} + +WeightLayoutVerticalStripe2D::WeightLayoutVerticalStripe2D( + ONNX_NAMESPACE::TensorProto_DataType proto_type, + int stripe_width) + : WeightLayout( + local_name_prefix + std::to_string(stripe_width), + proto_type, 2, 0.0f), + stripe_width_(stripe_width) { +} + +CoordTransFunc WeightLayoutVerticalStripe2D::ToActual(const tvm::Tensor& /*X*/) const { + return [&](const tvm::Array& nominal_coord) { + ORT_ENFORCE(nominal_coord.size() == 2); + const auto& y = nominal_coord[0]; + const auto& x = nominal_coord[1]; + return tvm::Array{ + x / stripe_width_, + y, + x % stripe_width_}; + }; +} + +CoordTransFunc WeightLayoutVerticalStripe2D::ToNominal(const tvm::Tensor& /*X*/) const { + return [&](const tvm::Array& actual_coord) { + ORT_ENFORCE(actual_coord.size() == 3); + const auto& z = actual_coord[0]; + const auto& y = actual_coord[1]; + const auto& x = actual_coord[2]; + return tvm::Array{ + y, + x + stripe_width_ * z}; + }; +} + +tvm::Array WeightLayoutVerticalStripe2D::ToActualShape(const tvm::Tensor& X) const { + tvm::Array new_shape = { + (X->shape[1] + stripe_width_ - 1) / stripe_width_, + X->shape[0], + stripe_width_}; + return new_shape; +} + +std::vector WeightLayoutVerticalStripe2D::ToActualShape(const Tensor* X) const { + ORT_ENFORCE(X != nullptr); + auto old_shape = X->Shape().GetDims(); + + ORT_ENFORCE(old_shape.size() == 2); + + std::vector new_shape = { + (old_shape[1] + stripe_width_ - 1) / stripe_width_, + old_shape[0], + stripe_width_}; + + return new_shape; +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.h b/onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.h new file mode 100644 index 0000000000000..d9b6f3299aa52 --- /dev/null +++ b/onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.h @@ -0,0 +1,40 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/codegen/common/common.h" +#include "core/codegen/target/weight_layout.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +// WeightLayoutVerticalStripe2D for making a 2D weight to 3D, by tiling the lowest (verteical) dimension +// [W, H] => [H/stripe, W, stripe] +class WeightLayoutVerticalStripe2D : public WeightLayout { + public: + static const std::string GetKey( + ONNX_NAMESPACE::TensorProto_DataType proto_type, + int stripe_width); + + public: + WeightLayoutVerticalStripe2D( + ONNX_NAMESPACE::TensorProto_DataType proto_type, + int stripe_width); + + ~WeightLayoutVerticalStripe2D() = default; + + virtual CoordTransFunc ToNominal(const tvm::Tensor& X) const override; + virtual CoordTransFunc ToActual(const tvm::Tensor& X) const override; + tvm::Array ToActualShape(const tvm::Tensor& X) const override; + std::vector ToActualShape(const Tensor* X) const override; + + private: + int stripe_width_; + + private: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(WeightLayoutVerticalStripe2D); +}; + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/ort_tvm_utils.cc b/onnxruntime/core/codegen/target/ort_tvm_utils.cc new file mode 100644 index 0000000000000..88aa987a371c2 --- /dev/null +++ b/onnxruntime/core/codegen/target/ort_tvm_utils.cc @@ -0,0 +1,186 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/ort_tvm_utils.h" + +#include "core/codegen/common/profile.h" +#include "core/codegen/target/codegen_context.h" +#include "core/providers/common.h" +#include "gsl/gsl_util" + +#include + +namespace onnxruntime { +namespace tvm_codegen { + +#define RETURN_DLDATATYPE_IF_MATCH(type, type_code) \ + if (ml_type == DataTypeImpl::GetType()) { \ + return {type_code, sizeof(type) * 8, 1}; \ + } + +// DLDataType: {DLDataTypeCode, bits, lanes} +DLDataType ToTvmDLDataType(MLDataType ml_type) { + if (ml_type->IsTensorType()) + ml_type = static_cast(ml_type)->GetElementType(); + + RETURN_DLDATATYPE_IF_MATCH(int8_t, kDLInt); + RETURN_DLDATATYPE_IF_MATCH(uint8_t, kDLUInt); + RETURN_DLDATATYPE_IF_MATCH(int16_t, kDLInt); + RETURN_DLDATATYPE_IF_MATCH(uint16_t, kDLUInt); + RETURN_DLDATATYPE_IF_MATCH(int32_t, kDLInt); + RETURN_DLDATATYPE_IF_MATCH(uint32_t, kDLUInt); + RETURN_DLDATATYPE_IF_MATCH(int64_t, kDLInt); + RETURN_DLDATATYPE_IF_MATCH(uint64_t, kDLUInt); + RETURN_DLDATATYPE_IF_MATCH(bool, kDLUInt); + + RETURN_DLDATATYPE_IF_MATCH(float, kDLFloat); + RETURN_DLDATATYPE_IF_MATCH(double, kDLFloat); + RETURN_DLDATATYPE_IF_MATCH(MLFloat16, kDLFloat); + + ORT_NOT_IMPLEMENTED("converting MLDataType ", ml_type, " to tvm DLDataType is not implemented"); +} + +tvm::Type ToTvmType(ONNX_NAMESPACE::TensorProto_DataType proto_type) { + switch (proto_type) { + // Note that bool is uint1 in tvm, but uint8 in ONNX, so it always require special handling + //case ONNX_NAMESPACE::TensorProto_DataType_BOOL: + // return tvm::UInt(1); /*break;*/ + case ONNX_NAMESPACE::TensorProto_DataType_INT16: + return tvm::Int(16); /*break;*/ + case ONNX_NAMESPACE::TensorProto_DataType_INT32: + return tvm::Int(32); /*break;*/ + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + return tvm::Int(64); /*break;*/ + case ONNX_NAMESPACE::TensorProto_DataType_UINT8: + return tvm::UInt(8); /*break;*/ + case ONNX_NAMESPACE::TensorProto_DataType_UINT16: + return tvm::UInt(16); /*break;*/ + case ONNX_NAMESPACE::TensorProto_DataType_UINT32: + return tvm::UInt(32); /*break;*/ + case ONNX_NAMESPACE::TensorProto_DataType_UINT64: + return tvm::UInt(64); /*break;*/ + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: + return tvm::Float(32); /*break;*/ + case ONNX_NAMESPACE::TensorProto_DataType_DOUBLE: + return tvm::Float(64); /*break;*/ + case ONNX_NAMESPACE::TensorProto_DataType_INT8: + return tvm::Int(8); /*break;*/ + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: + return tvm::Float(16); /*break;*/ + case ONNX_NAMESPACE::TensorProto_DataType_STRING: + ORT_THROW("Casting to and from strings is not supported yet."); /*break;*/ + case ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED: + ORT_THROW("Cast op must have 'to' argument of type DataType"); /*break;*/ + default: + ORT_THROW("Unexpected 'to' argument value: ", proto_type); + } +} + +tvm::Array ShapeToTvmArray(const NodeArg* def, CodeGenContext& ctx) { + ORT_ENFORCE(nullptr != def); + const ONNX_NAMESPACE::TensorShapeProto* shape_proto = def->Shape(); + ORT_ENFORCE(nullptr != shape_proto); + + tvm::Array arr; + for (int i = 0; i < shape_proto->dim_size(); ++i) { + arr.push_back(ShapeDimToTvmDim(shape_proto->dim(i), ctx)); + } + return arr; +} + +tvm::Expr ShapeDimToTvmDim(const ONNX_NAMESPACE::TensorShapeProto_Dimension& dim, CodeGenContext& ctx) { + if (dim.has_dim_param()) { + return ctx.GetOrCreateDynamicDim(dim.dim_param()); + } else if (dim.has_dim_value()) { + return tvm::Expr(gsl::narrow_cast(dim.dim_value())); + } + return ctx.GetOrCreateDynamicDim(ctx.CreateUnnamedSymbol()); +} + +#ifdef CODEGEN_ENABLE_PROFILER +struct event_in_bracket_and_id { + bool in_bracket; + int id; +}; +std::unordered_map g_codegen_profiler_event_ids; +std::vector> g_codegen_profiler_events(1024); + +TVM_REGISTER_GLOBAL("tvm.contrib.onnxruntime.profile_event") + .set_body([](tvm::TVMArgs args, tvm::TVMRetValue* ret) { + DLTensor* X = args[0]; + DLTensor* Y = args[1]; + int event_id = args[2]; + bool is_begin = args[3]; + if (!is_begin) { + DCHECK(event_id < g_codegen_profiler_event_ids.size()); + profiling::Profiler::Instance().EndTimeAndRecordEvent( + profiling::EventCategory::NODE_EVENT, + g_codegen_profiler_events[event_id].first, + g_codegen_profiler_events[event_id].second); + } + + { + CODEGEN_PROFILER_EVENT(profile_stub); + int64_t elem_count = 1; + for (int i = 0; i < X->ndim; ++i) { + elem_count *= X->shape[i]; + } + // there's overhead in this copy, so put begin after copy and end before copy + memcpy(static_cast(Y->data) + Y->byte_offset, + static_cast(X->data) + X->byte_offset, + elem_count * X->dtype.bits / 8); + } + + if (is_begin) { + DCHECK(g_codegen_profiler_events.size() > event_id); + DCHECK(!g_codegen_profiler_events[event_id].first.empty()); + DCHECK(g_codegen_profiler_event_ids[g_codegen_profiler_events[event_id].first].id == event_id); + g_codegen_profiler_events[event_id].second = + profiling::Profiler::Instance().StartTime(); + } + }); + +tvm::Tensor ProfileBegin(tvm::Tensor X, const std::string& event_name) { + int event_id; + if (g_codegen_profiler_event_ids.count(event_name) == 0) { + event_id = g_codegen_profiler_event_ids.size(); + ORT_ENFORCE(event_id < g_codegen_profiler_events.size()); + } else { + ORT_ENFORCE(!g_codegen_profiler_event_ids[event_name].in_bracket); + event_id = g_codegen_profiler_event_ids[event_name].id; + } + g_codegen_profiler_event_ids[event_name] = {true, event_id}; + g_codegen_profiler_events[event_id].first = event_name; + return topi::detail::make_extern( + {X->shape}, {X->dtype}, {X}, + [&](tvm::Array ins, tvm::Array outs) { + return topi::detail::call_packed({tvm::Expr("tvm.contrib.onnxruntime.profile_event"), + topi::detail::pack_buffer(ins[0]), + topi::detail::pack_buffer(outs[0]), + event_id, + true}); + }, + event_name + "_begin", "", {})[0]; +} + +tvm::Tensor ProfileEnd(tvm::Tensor X, const std::string& event_name) { + ORT_ENFORCE(g_codegen_profiler_event_ids.at(event_name).in_bracket); + g_codegen_profiler_event_ids.at(event_name).in_bracket = false; + int event_id = g_codegen_profiler_event_ids.at(event_name).id; + ORT_ENFORCE(event_id < g_codegen_profiler_events.size()); + ORT_ENFORCE(g_codegen_profiler_events[event_id].first == event_name); + return topi::detail::make_extern( + {X->shape}, {X->dtype}, {X}, + [&](tvm::Array ins, tvm::Array outs) { + return topi::detail::call_packed({tvm::Expr("tvm.contrib.onnxruntime.profile_event"), + topi::detail::pack_buffer(ins[0]), + topi::detail::pack_buffer(outs[0]), + event_id, + false}); + }, + event_name + "_end", "", {})[0]; +} +#endif + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/ort_tvm_utils.h b/onnxruntime/core/codegen/target/ort_tvm_utils.h new file mode 100644 index 0000000000000..f13e91a2d5cea --- /dev/null +++ b/onnxruntime/core/codegen/target/ort_tvm_utils.h @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/codegen/common/common.h" +#include "core/framework/data_types.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +class CodeGenContext; + +// Helper function that converts a onnxruntime MLDataType to TVM DLDataType +DLDataType ToTvmDLDataType(MLDataType ml_type); + +tvm::Type ToTvmType(ONNX_NAMESPACE::TensorProto_DataType proto_type); + +tvm::Array ShapeToTvmArray(const NodeArg* def, CodeGenContext& ctx); + +tvm::Expr ShapeDimToTvmDim(const ONNX_NAMESPACE::TensorShapeProto_Dimension& dim, CodeGenContext& ctx); + +#ifdef CODEGEN_ENABLE_PROFILER +// Helper functions to inspect into lowered function +tvm::Tensor ProfileBegin(tvm::Tensor X, const std::string& event_name); + +tvm::Tensor ProfileEnd(tvm::Tensor X, const std::string& event_name); +#endif + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/tvm_ir_builder.cc b/onnxruntime/core/codegen/target/tvm_ir_builder.cc new file mode 100644 index 0000000000000..c1531c8cfacc2 --- /dev/null +++ b/onnxruntime/core/codegen/target/tvm_ir_builder.cc @@ -0,0 +1,125 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/tvm_ir_builder.h" + +#include "core/codegen/common/op_macro.h" +#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/common/common.h" + +namespace onnxruntime { +namespace tvm_codegen { + +TVMIRBuilder::TVMIRBuilder(const std::string& name) + : name_(name) {} + +const std::string& TVMIRBuilder::Name() const { + return name_; +} + +void TVMIRBuilder::InsertDispatcher(std::unique_ptr&& ptr) { + dispatchers_.push_back(std::move(ptr)); +} + +void TVMIRBuilder::ClearAllDispatchers() { + dispatchers_.clear(); +} + +void TVMIRBuilder::DumpAllOpIRCreators() const { + int count = 0; + for (auto& d : dispatchers_) { + std::cout << "************ TVM OpIRDispatcher " + << count << " : " + << d->Name() + << " ************" << std::endl; + + d->ForEach([](const std::string& key, OpIRCreator* builder) { + std::cout << "Key " << key + << ", Creator " << builder->Name() << std::endl; + }); + + ++count; + } +} + +// Evaluate finds ONE proper OpIRCreator and build the corresponding OpIR +// If a TVMIRBuilder has more than one OpIRCreator for an ORT Op, +// the first one will be used. +// Please adjust registration order and dispatcher in TVMIRBuilder +// to make sure the proper OpIRCreator is called. +Status TVMIRBuilder::Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext& ctx_codegen, + tvm::Array& outputs) { + OpIRCreator* candidate = nullptr; + for (auto& d : dispatchers_) { + candidate = d->Find(node); + if (nullptr != candidate) + break; + } + + if (nullptr == candidate) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Not implemented: ", node.OpType()); + } + + ORT_RETURN_IF_ERROR(candidate->Evaluate(inputs, node, ctx_codegen, outputs)); + + return Status::OK(); +} + +// BEGIN: Generic IR creator classes +#define ADD_OP_ITEM(name) \ + op_ir_registry->Register(std::move(std::make_unique())); + +#define BINARY_OP(name) ADD_OP_ITEM(name) +#define BINARY_CMP_OP(name) ADD_OP_ITEM(name) +#define POOL_OP(name) ADD_OP_ITEM(name) +#define REDUCE_OP(name) ADD_OP_ITEM(name) +#define REDUCE_INDEXED_OP(name) ADD_OP_ITEM(name) +#define UNARY_OP(name) ADD_OP_ITEM(name) +#define VARIADIC_OP(name) ADD_OP_ITEM(name) + +void RegisterAllGenericOpIRCreators(OpIRRegistry* op_ir_registry) { + LIST_ALL_GENERIC_OPS(); +} + +#undef ADD_OP_ITEM +#undef BINARY_OP +#undef BINARY_CMP_OP +#undef POOL_OP +#undef REDUCE_OP +#undef REDUCE_INDEXED_OP +#undef UNARY_OP +#undef VARIADIC_OP + +// BEGIN: Plugin Generic IR creator classes +#define ADD_OP_ITEM(name) \ + dispatcher->Register(#name, registry->Get(GENERIC_OP_IR_CREATOR_STRING(name))); + +#define BINARY_OP(name) ADD_OP_ITEM(name) +#define BINARY_CMP_OP(name) ADD_OP_ITEM(name) +#define POOL_OP(name) ADD_OP_ITEM(name) +#define REDUCE_OP(name) ADD_OP_ITEM(name) +#define REDUCE_INDEXED_OP(name) ADD_OP_ITEM(name) +#define UNARY_OP(name) ADD_OP_ITEM(name) +#define VARIADIC_OP(name) ADD_OP_ITEM(name) + +void RegisterGenericOrtOpTypeDispatcher(const std::shared_ptr& builder, + const OpIRRegistry* registry) { + auto dispatcher = std::make_unique("GenericOrtOpTypeOpIRCreators"); + LIST_ALL_GENERIC_OPS() + builder->InsertDispatcher(std::move(dispatcher)); +} + +#undef ADD_OP_ITEM +#undef BINARY_OP +#undef BINARY_CMP_OP +#undef POOL_OP +#undef REDUCE_OP +#undef REDUCE_INDEXED_OP +#undef UNARY_OP +// END: Generic IR creators classes + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/tvm_ir_builder.h b/onnxruntime/core/codegen/target/tvm_ir_builder.h new file mode 100644 index 0000000000000..f7f69647e573b --- /dev/null +++ b/onnxruntime/core/codegen/target/tvm_ir_builder.h @@ -0,0 +1,64 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/codegen/target/codegen_context.h" +#include "core/codegen/target/tvm_op_creator.h" +#include "core/common/common.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// TVMIRBuilder contains all applicable TVM OpIRCreators +// OpIRCreators are stored in multiple dispatchers +// that check different conditions of an ORT Node. + +// If an ORT Node satisfies more than one OpIRCreators, +// the first dispatched pass will be applied. + +class TVMIRBuilder { + public: + TVMIRBuilder(const std::string& name); + ~TVMIRBuilder() = default; + + // A debug dumps all existing in this TVMIRBuilders + void DumpAllOpIRCreators() const; + + // Evaluates an OpIRCreator that first satisfies condtions of all dispatchers + Status Evaluate( + const tvm::Array& inputs, + const Node& node, + CodeGenContext& ctx, + tvm::Array& outputs); + + // Inserts a dispatcher and move its ownership to this TVMIRBuilder + void InsertDispatcher(std::unique_ptr&& ptr); + + // Clears all dispatchers in this TVMIRBuilder + void ClearAllDispatchers(); + + // Dumps the name of this TVMIRBuilder + const std::string& Name() const; + + private: + std::vector> dispatchers_; + std::string name_; + + private: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TVMIRBuilder); +}; + +// Utility function to register all builtin generic OpIRCreators into an OpIRRegistry. +// It creates instances of all generic OpIRCreators +// and registers them to op_ir_registry +void RegisterAllGenericOpIRCreators(OpIRRegistry* op_ir_registry); + +// Utility function to bind all builtin generic OpIRCreators to a TVMIRBuilder. +// It creates an instance of a Dispatcher that contains all generic OpIRCreators created above +// and uses OrtOpType to dispatch OpIRCreators. +// Then, it registers the created Dispatcher to a TVMIRBuilder, builder. +void RegisterGenericOrtOpTypeDispatcher(const std::shared_ptr& builder, + const OpIRRegistry* registry); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/tvm_op_creator.cc b/onnxruntime/core/codegen/target/tvm_op_creator.cc new file mode 100644 index 0000000000000..ef2da19c6d8de --- /dev/null +++ b/onnxruntime/core/codegen/target/tvm_op_creator.cc @@ -0,0 +1,37 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/tvm_op_creator.h" + +#include "core/codegen/common/common.h" +#include "core/codegen/common/dispatcher.h" +#include "core/codegen/target/codegen_context.h" + +namespace onnxruntime { +namespace codegen { +// Explicit instantiation for OpIRCreator +template class CreatorBase&, + const Node&, + tvm_codegen::CodeGenContext&, + tvm::Array&, + Status>; + +// Explicit instantiation for OpIRCreators' dispatcher +template class DispatcherBase; + +} // namespace codegen + +namespace tvm_codegen { + +// One dispatcher is based on ORT OpType +OpIRCreator* OP_IR_DISPATCHER_CLASS(OpType)::Find(const Node& node) { + return DispatcherBase::Get(node.OpType()); +} + +// Another dispatcher is based ORT NodeArg name (GetKey) +OpIRCreator* OP_IR_DISPATCHER_CLASS(NodeName)::Find(const Node& node) { + return DispatcherBase::Get(GetKey(&node)); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/tvm_op_creator.h b/onnxruntime/core/codegen/target/tvm_op_creator.h new file mode 100644 index 0000000000000..fe2648462e4f5 --- /dev/null +++ b/onnxruntime/core/codegen/target/tvm_op_creator.h @@ -0,0 +1,84 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/codegen/common/creator.h" +#include "core/codegen/common/dispatcher.h" +#include "core/codegen/common/registry.h" +#include "core/graph/graph.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +class CodeGenContext; + +// OpIRCreator lowers an Ort Node to its corresponding TVM IRs +using OpIRCreator = codegen::CreatorBase< + const tvm::Array&, + const Node&, + CodeGenContext&, + tvm::Array&, + Status>; + +// OpIRDispatcher is the base dispatcher for TVM IR Builder +// It checks whether an Ort Node satisfying a criteria (in Find) +// and dispatches a corresponding OpIRCreator. +class OpIRDispatcher : public codegen::DispatcherBase { + public: + OpIRDispatcher(const std::string& name) + : DispatcherBase(name) {} + + ~OpIRDispatcher() = default; + + virtual OpIRCreator* Find(const Node&) = 0; + + private: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(OpIRDispatcher); +}; + +// Macro returns an OpIRCreators' dispatcher's name +#define OP_IR_DISPATCHER_CLASS(OP) \ + TVM##OP##IRCreator + +// Macro declares an OpIRCreators' dispatcher +#define DECLARE_OP_IR_DISPATCHER_CLASS(OP) \ + class OP_IR_DISPATCHER_CLASS(OP) : public OpIRDispatcher { \ + public: \ + TVM##OP##IRCreator(const std::string& name) \ + : OpIRDispatcher(name) {} \ + ~TVM##OP##IRCreator() = default; \ + OpIRCreator* Find(const Node&) override; \ + \ + private: \ + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(OP_IR_DISPATCHER_CLASS(OP)); \ + }; + +// Declare two common dispatchers for TVM Op IR builders +// One dispatcher is based on Ort OpType +DECLARE_OP_IR_DISPATCHER_CLASS(OpType) +// Another dispatcher is based Ort NodeArg name +DECLARE_OP_IR_DISPATCHER_CLASS(NodeName) + +// OpIRCreator Registry is a registry holds all OpIRCreators +using OpIRRegistry = codegen::RegistryBase; + +// Macro declares an OpIRCreator +#define DECLARE_OP_IR_CREATOR_CLASS(OP, PREFIX) \ + DECLARE_CREATOR_CLASS(OP, PREFIX##IRCreator, \ + const tvm::Array&, \ + const Node&, \ + tvm_codegen::CodeGenContext&, \ + tvm::Array&, \ + Status) + +// Macro returns an OpIRCreator's name with prefix +#define OP_IR_CREATOR_CLASS_EX(OP, PREFIX, ARCH) \ + CREATOR_CLASS(OP, PREFIX##ARCH##IRCreator) + +// Macro declares an OpIRCreator with prefix and arch +#define DECLARE_OP_IR_CREATOR_CLASS_EX(OP, PREFIX, ARCH) \ + DECLARE_OP_IR_CREATOR_CLASS(OP, PREFIX##ARCH) + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/tvm_schedule_builder.cc b/onnxruntime/core/codegen/target/tvm_schedule_builder.cc new file mode 100644 index 0000000000000..be0d6ab44530f --- /dev/null +++ b/onnxruntime/core/codegen/target/tvm_schedule_builder.cc @@ -0,0 +1,104 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/tvm_schedule_builder.h" + +#include "core/codegen/common/op_macro.h" +#include "core/codegen/common/settings.h" +#include "core/common/common.h" +#include "core/common/logging/logging.h" + +namespace onnxruntime { +namespace tvm_codegen { + +TVMScheduleBuilder::TVMScheduleBuilder(const std::string& name) + : name_(name) { +} + +const std::string& TVMScheduleBuilder::Name() const { + return name_; +} + +void TVMScheduleBuilder::InsertDispatcher(std::unique_ptr&& ptr) { + dispatchers_.push_back(std::move(ptr)); +} + +void TVMScheduleBuilder::ClearDispatcher() { + dispatchers_.clear(); +} + +void TVMScheduleBuilder::DumpAllSchedulers() const { + std::ostringstream stream; + int count = 0; + stream << "[CODEGEN_DUMP_SCHEDULE]" << std::endl; + for (auto& d : dispatchers_) { + stream << "************ TVM Scheduler Dispatcher " + << count << " : " + << d->Name() + << " ************" << std::endl; + + d->ForEach([&stream](const std::string& key, Scheduler* op) { + stream << "Key " << key + << ", Creater " << op->Name() << std::endl; + }); + + ++count; + } + + LOGS_DEFAULT(CODEGEN_SETTINGS_LOG_LEVEL) << stream.str(); +} + +Status TVMScheduleBuilder::Evaluate( + const tvm::Tensor& tensor, + const Node* node, + CodeGenContext& ctx_codegen, + ScheduleContext& sched) { + Scheduler* candidate = nullptr; + + for (auto& d : dispatchers_) { + candidate = d->Find(tensor, node, ctx_codegen); + if (nullptr != candidate) + break; + } + + bool enable_dump_schedule = codegen::CodeGenSettings::Instance().HasOption(codegen::CodeGenSettings::kCodeGenDumpSchedule); + + if (nullptr == candidate) { + if (nullptr != node) + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Not implemented: ", node->OpType()); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Not implemented an internal tvm::Tensor: ", tensor->op->name); + } + + bool status = candidate->Evaluate(tensor, node, ctx_codegen, sched); + + if (enable_dump_schedule) { + std::ostringstream stream; + if (nullptr != node) { + stream << std::endl; + stream << "[CODEGEN_DUMP_SCHEDULE] " + << "Schedule Node: " << node->Name() << std::endl; + } else { + stream << std::endl; + } + + if (status) { + stream << "[CODEGEN_DUMP_SCHEDULE] " + << "Schedule tvm::Tesnor " + << tensor->op->name + << " with " + << candidate->Name() << std::endl; + } else { + stream << "[CODEGEN_DUMP_SCHEDULE] " + << "Schedule tvm::Tesnor " + << tensor->op->name + << " is suppressed " << std::endl; + } + + LOGS_DEFAULT(CODEGEN_SETTINGS_LOG_LEVEL) << stream.str(); + } + + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/tvm_schedule_builder.h b/onnxruntime/core/codegen/target/tvm_schedule_builder.h new file mode 100644 index 0000000000000..e125b1c5db00a --- /dev/null +++ b/onnxruntime/core/codegen/target/tvm_schedule_builder.h @@ -0,0 +1,46 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/codegen/target/tvm_scheduler.h" +#include "core/common/common.h" + +namespace onnxruntime { +namespace tvm_codegen { + +// TVMScheduleBuilder contains all applicable TVM scheduler passes. +// Scheduler passes are stored in multiple dispatchers +// that check different conditions of a tvm::Tensor. + +// If a tvm::Tensor satisfies more than one TVM scheduler passes, +// the first dispatched pass will be applied. + +class TVMScheduleBuilder { + public: + // TODO: add more parameter in consructor to support different target + TVMScheduleBuilder(const std::string& name); + ~TVMScheduleBuilder() = default; + + void DumpAllSchedulers() const; + + Status Evaluate( + const tvm::Tensor& tensor, + const Node* node, + CodeGenContext& ctx, + ScheduleContext& sched); + + void InsertDispatcher(std::unique_ptr&& ptr); + void ClearDispatcher(); + + const std::string& Name() const; + + private: + std::vector> dispatchers_; + std::string name_; + + private: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TVMScheduleBuilder); +}; + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/tvm_scheduler.cc b/onnxruntime/core/codegen/target/tvm_scheduler.cc new file mode 100644 index 0000000000000..697a334f469de --- /dev/null +++ b/onnxruntime/core/codegen/target/tvm_scheduler.cc @@ -0,0 +1,79 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/tvm_scheduler.h" + +#include "core/codegen/common/common.h" +#include "core/codegen/common/dispatcher.h" +#include "core/codegen/target/codegen_context.h" + +namespace onnxruntime { +namespace codegen { +// explicit instantiation +template class CreatorBase; + +template class DispatcherBase; + +} // namespace codegen + +namespace tvm_codegen { + +static const std::string TMVOpRuleKey_Extern("TVMOpRule_Extern"); +static const std::string TMVOpRuleKey_ComputeReduce("TVMOpRule_ComputeReduce"); +static const std::string TMVOpRuleKey_ComputeRegular("TVMOpRule_ComputeRegular"); +static const std::string TMVOpRuleKey_AlwaysRoot("TMVOpRuleKey_AlwaysRoot"); +static const std::string TMVOpRuleKey_NoRule("TVMOpRule_NoRule"); + +const std::string& GetTVMOpRule(TVMOpRuleType rule) { + if (rule == TVMOpRuleType::Extern) { + return TMVOpRuleKey_Extern; + } else if (rule == TVMOpRuleType::ComputeReduce) { + return TMVOpRuleKey_ComputeReduce; + } else if (rule == TVMOpRuleType::AlwaysRoot) { + return TMVOpRuleKey_AlwaysRoot; + } + return TMVOpRuleKey_NoRule; +} + +const std::string& GetTVMOpRule(const tvm::Tensor& tensor) { + auto extern_op = tensor->op.as(); + + if (nullptr != extern_op) { + return TMVOpRuleKey_Extern; + } + + auto compute_op = tensor->op.as(); + if (nullptr != compute_op) { + if (compute_op->reduce_axis.size() > 0) { + return TMVOpRuleKey_ComputeReduce; + } + } + + return TMVOpRuleKey_NoRule; +} + +Scheduler* SCHEDULE_DISPATCHER_CLASS(OrtOpType):: + Find(const tvm::Tensor&, const Node* node, tvm_codegen::CodeGenContext&) { + if (nullptr == node) + return nullptr; + return DispatcherBase::Get(node->OpType()); +} + +Scheduler* SCHEDULE_DISPATCHER_CLASS(TVMOpRule):: + Find(const tvm::Tensor& tensor, const Node*, tvm_codegen::CodeGenContext&) { + return DispatcherBase::Get(GetTVMOpRule(tensor)); +} + +Scheduler* SCHEDULE_DISPATCHER_CLASS(OrtOpName):: + Find(const tvm::Tensor&, const Node* node, tvm_codegen::CodeGenContext&) { + if (nullptr == node) + return nullptr; + return DispatcherBase::Get(GetKey(node)); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/tvm_scheduler.h b/onnxruntime/core/codegen/target/tvm_scheduler.h new file mode 100644 index 0000000000000..9c2dbfdd327c4 --- /dev/null +++ b/onnxruntime/core/codegen/target/tvm_scheduler.h @@ -0,0 +1,128 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/common/common.h" +#include "core/codegen/common/creator.h" +#include "core/codegen/common/registry.h" +#include "core/codegen/target/codegen_context.h" +#include "core/graph/graph.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +// These are current generic TVMOpRule we used. +enum class TVMOpRuleType : int { + Extern = 0, + ComputeReduce = 1, + ComputeRegular = 2, + AlwaysRoot = 3, // for debug + NoRule, +}; + +const std::string& GetTVMOpRule(const tvm::Tensor& tensor); +const std::string& GetTVMOpRule(TVMOpRuleType rule); + +// These are current generic ScheduleType in tvm_codegen +enum class ScheduleType : int { + ScheduleNone = 0, + ScheduleInline = 1, + ScheduleAt = 2, + ScheduleRoot = 3, + ScheduleClosure = 4, +}; + +// Data struct to bundle tvm::Schedule and scheduled tensor +struct ScheduleContext { + ScheduleContext(const tvm::Array& ops) { + schedule = tvm::create_schedule(ops); + } + tvm::Schedule schedule; + std::map scheduled_tensors; +}; + +// Scheduler inserts a tvm::Schedule content to a tvm::Tensor +using Scheduler = codegen::CreatorBase< + const tvm::Tensor&, + const Node*, + tvm_codegen::CodeGenContext&, + ScheduleContext&, + bool>; + +// TVMScheduleDispatcher is the base dispatcher for TVM Schedule Builder +// It checks whether a pair of {tvm::Tensor, Ort Node} satisfying a criteria (in Find) +// and dispatches a corresponding Scheduler. +class TVMScheduleDispatcher : public codegen::DispatcherBase { + public: + TVMScheduleDispatcher(const std::string& name) + : DispatcherBase(name) {} + + ~TVMScheduleDispatcher() = default; + + virtual Scheduler* Find(const tvm::Tensor&, + const Node*, + tvm_codegen::CodeGenContext&) = 0; + + private: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TVMScheduleDispatcher); +}; + +// Macro returns an Schedulers' dispatcher's name +#define SCHEDULE_DISPATCHER_CLASS(TYPE) \ + TVM##TYPE##Schedulers + +// Macro declares an Schedulers' dispatcher +#define DECLARE_SCHEDULE_DISPATCHER_CLASS(TYPE) \ + class SCHEDULE_DISPATCHER_CLASS(TYPE) : public tvm_codegen::TVMScheduleDispatcher { \ + public: \ + TVM##TYPE##Schedulers(const std::string& name) \ + : TVMScheduleDispatcher(name) {} \ + ~TVM##TYPE##Schedulers() = default; \ + tvm_codegen::Scheduler* Find(const tvm::Tensor&, \ + const Node*, \ + tvm_codegen::CodeGenContext&) override; \ + \ + private: \ + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TVM##TYPE##Schedulers); \ + }; + +// Common dispatchers are listed here +// For a special pattern, it can be created later. +// One dispatcher is based on Ort OpType +DECLARE_SCHEDULE_DISPATCHER_CLASS(OrtOpType) +// One dispatcher is based on TVMOpRule +DECLARE_SCHEDULE_DISPATCHER_CLASS(TVMOpRule) +// One dispatcher is based Ort NodeArg name +DECLARE_SCHEDULE_DISPATCHER_CLASS(OrtOpName) + +// Scheduler Registry is a registry holds all Schedulers +using TVMScheduleRegistry = codegen::RegistryBase; + +// Macro declares TVM scheduler class +#define DECLARE_TVM_SCHEDULER_CLASS(OP, PRETFIX) \ + DECLARE_CREATOR_CLASS(OP, PRETFIX##Scheduler, \ + const tvm::Tensor&, \ + const Node*, \ + tvm_codegen::CodeGenContext&, \ + tvm_codegen::ScheduleContext&, \ + bool) + +// Macro returns TVM scheduler's name with prefix +#define TVM_SCHEDULER_CLASS(OP, PREFIX) \ + CREATOR_CLASS(OP, PREFIX##Scheduler) + +// Macro returns TVM scheduler's name as string +#define TVM_SCHEDULER_STRING(OP, PREFIX) \ + STRINGIZE(TVM_SCHEDULER_CLASS(OP, PREFIX)) + +// Macro returns TVM scheduler's name with prefix and arch +#define TVM_SCHEDULER_CLASS_EX(OP, PREFIX, ARCH) \ + CREATOR_CLASS(OP, PREFIX##ARCH##Scheduler) + +// Macro declares TVM scheduler class with prefix and arch +#define DECLARE_TVM_SCHEDULER_CLASS_EX(OP, PREFIX, ARCH) \ + DECLARE_TVM_SCHEDULER_CLASS(OP, PREFIX##ARCH) + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/weight_layout.cc b/onnxruntime/core/codegen/target/weight_layout.cc new file mode 100644 index 0000000000000..f74c6ad0727e4 --- /dev/null +++ b/onnxruntime/core/codegen/target/weight_layout.cc @@ -0,0 +1,92 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/codegen/target/weight_layout.h" + +#include "core/codegen/common/common.h" +#include "core/codegen/common/utils.h" +#include "core/codegen/mti/mti_tvm_utils.h" +#include "core/codegen/target/ort_tvm_utils.h" + +namespace onnxruntime { +namespace tvm_codegen { + +static tvm::Tensor CreateTVMPlaceholder( + const std::string& name, + HalideIR::Type type, + int dim) { + tvm::Array shape; + if (dim > 0) { + for (int i = 0; i < dim; ++i) { + shape.push_back(tvm::Var(name + "_v" + std::to_string(i))); + } + } else { + shape.push_back(1); + } + return tvm::placeholder(shape, type, name + "_placeholder"); +} + +const std::string WeightLayout::GetKey( + const std::string& name, + ONNX_NAMESPACE::TensorProto_DataType proto_type, + int input_dim, + float pad_zero) { + std::string key = name; + key += "_type_" + std::to_string(static_cast(proto_type)); + key += "_dim_" + input_dim; + key += "_pad_zero_" + std::to_string(pad_zero); + key = NormalizeCppName(key); + return key; +} + +WeightLayout::WeightLayout( + const std::string& name, + ONNX_NAMESPACE::TensorProto_DataType proto_type, + int input_dim, + float pad_zero) + : name_(GetKey(name, proto_type, input_dim, pad_zero)), + proto_type_(proto_type), + input_dim_(input_dim), + pad_zero_(pad_zero) {} + +const std::string& WeightLayout::Name() const { + return name_; +} + +void WeightLayout::CreateLayoutMarshallingTVMOp(tvm::Array& inputs, + tvm::Array& outputs) const { + HalideIR::Type halide_type = ToTvmType(proto_type_); + + tvm::Tensor placeholder = CreateTVMPlaceholder(name_, halide_type, input_dim_); + inputs.push_back(placeholder); + + tvm::Array new_shape = ToActualShape(placeholder); + CoordTransFunc new_coord_to_old_coord_func = ToNominal(placeholder); + tvm::Expr pad_zero_expr = tvm::make_const(halide_type, pad_zero_); + + tvm::Tensor output = tvm::compute( + new_shape, + [&](const tvm::Array& output_coord) { + tvm::Array output_coord1; + for (const auto& coord : output_coord) + output_coord1.push_back(coord); + auto input_coord = new_coord_to_old_coord_func(output_coord1); + ORT_ENFORCE(input_coord.size() == placeholder->shape.size()); + + if (input_coord.size() > 0) { + auto in_range = (input_coord[0] >= 0) && (input_coord[0] < placeholder->shape[0]); + for (size_t dim = 1; dim < input_coord.size(); ++dim) + in_range = in_range && (input_coord[dim] >= 0) && (input_coord[dim] < placeholder->shape[dim]); + + return tvm::ir::Select::make(in_range, placeholder(input_coord), pad_zero_expr); + } else { + // scalar + return placeholder(input_coord); + } + }); + + outputs.push_back(output); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/target/weight_layout.h b/onnxruntime/core/codegen/target/weight_layout.h new file mode 100644 index 0000000000000..bcd9b229b5a3d --- /dev/null +++ b/onnxruntime/core/codegen/target/weight_layout.h @@ -0,0 +1,68 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/codegen/common/common.h" +#include "core/codegen/common/registry.h" +#include "core/common/common.h" +#include "core/framework/tensor.h" +#include + +namespace onnxruntime { +namespace tvm_codegen { + +using CoordTransFunc = std::function(const tvm::Array&)>; + +// WeightLayout is data layout trasnformer for weight/initializer +class WeightLayout { + public: + // Static function to return unique string as a key + static const std::string GetKey( + const std::string& name, + ONNX_NAMESPACE::TensorProto_DataType proto_type, + int input_dim, + float pad_zero); + + public: + WeightLayout( + const std::string& name, + ONNX_NAMESPACE::TensorProto_DataType proto_type, + int input_dim, + float pad_zero); + + ~WeightLayout() = default; + + // Return a CoordTransFunc from actual (transformed) coordinate to normial (original) coordinate + virtual CoordTransFunc ToNominal(const tvm::Tensor& X) const = 0; + + // Return a CoordTransFunc from normial (original) coordinate to actual (transformed) coordinate + virtual CoordTransFunc ToActual(const tvm::Tensor& X) const = 0; + + // Return actual (transformed) shape in tvm::Array (tvm_codegen) + virtual tvm::Array ToActualShape(const tvm::Tensor& X) const = 0; + + // Return actual (transformed) shape in vector (ort) + virtual std::vector ToActualShape(const Tensor* X) const = 0; + + // Create Layout Marshalling op in outputs + void CreateLayoutMarshallingTVMOp(tvm::Array& inputs, + tvm::Array& outputs) const; + + // Layout name + const std::string& Name() const; + + protected: + std::string name_; + ONNX_NAMESPACE::TensorProto_DataType proto_type_; + int input_dim_; + float pad_zero_; + + private: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(WeightLayout); +}; + +// Weight Layout Registry is a registry holds all WeightLayout +using WeightLayoutRegistry = codegen::RegistryBase; + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/tvm/tvm_compiler.cc b/onnxruntime/core/codegen/tvm/tvm_compiler.cc deleted file mode 100644 index a3ae548f70363..0000000000000 --- a/onnxruntime/core/codegen/tvm/tvm_compiler.cc +++ /dev/null @@ -1,93 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include -#include -#include "core/codegen/tvm/tvm_compiler.h" -namespace onnxruntime { - -TVMGraph::TensorDescriptor::TensorDescriptor(MLDataType type, onnxruntime::ProviderType execution_provider_type, tvm::Tensor tvm_tensor) : tvm_tensor_(tvm_tensor) { - if (execution_provider_type == onnxruntime::kCpuExecutionProvider) { - ctx_.device_type = DLDeviceType::kDLCPU; - ctx_.device_id = 0; - } else { - ORT_NOT_IMPLEMENTED("Non-cpu execution provider not supported on TVM now."); - } - - if (DataTypeImpl::GetTensorType() == type) { - dtype_.code = kDLFloat; - dtype_.bits = 64; - dtype_.lanes = 1; - } else { - ORT_NOT_IMPLEMENTED("Non-double type not supported on TVM now."); - } -} - -class IdGenerator { - public: - IdGenerator() {} - int GetNext() { - return cur_++; - } - - private: - int cur_{0}; -}; - -// This is a special compiler step for the test case that sum two 1-D tensors -static void Compile1DAddToTVM(const onnxruntime::Node& node, std::unordered_map& tvm_tensors, onnxruntime::ProviderType execution_provider_type, IdGenerator& generator) { - ORT_ENFORCE(node.OpType() == "Add"); - tvm::Array shape; - shape.push_back(tvm::var("n1")); - - tvm::Tensor t1; - tvm::Tensor t2; - auto it = tvm_tensors.find(node.InputDefs()[0]->Name()); - if (it == tvm_tensors.end()) { - tvm_tensors[node.InputDefs()[0]->Name()] = TVMGraph::TensorDescriptor( - DataTypeImpl::TypeFromProto(*node.InputDefs()[0]->TypeAsProto()), - execution_provider_type, - tvm::placeholder(shape, tvm::Float(64), "T" + std::to_string(generator.GetNext()))); - } - t1 = tvm_tensors[node.InputDefs()[0]->Name()].tvm_tensor_; - it = tvm_tensors.find(node.InputDefs()[1]->Name()); - if (it == tvm_tensors.end()) { - tvm_tensors[node.InputDefs()[1]->Name()] = TVMGraph::TensorDescriptor( - DataTypeImpl::TypeFromProto(*node.InputDefs()[1]->TypeAsProto()), - execution_provider_type, - tvm::placeholder(shape, tvm::Float(64), "T" + std::to_string(generator.GetNext()))); - } - t2 = tvm_tensors[node.InputDefs()[1]->Name()].tvm_tensor_; - - tvm_tensors[node.OutputDefs()[0]->Name()] = TVMGraph::TensorDescriptor( - DataTypeImpl::TypeFromProto(*node.InputDefs()[1]->TypeAsProto()), - execution_provider_type, - tvm::compute( - t1->shape, [&t1, &t2](tvm::Expr i) { - return t1[i] + t2[i]; - }, - "T" + std::to_string(generator.GetNext()))); -} - -TVMGraph CompileToTVM(const onnxruntime::Graph& graph, onnxruntime::ProviderType execution_provider_type) { - TVMGraph result; - std::unordered_map tvm_tensors; - IdGenerator generator; - for (auto& node : graph.Nodes()) { - Compile1DAddToTVM(node, tvm_tensors, execution_provider_type, generator); - } - - for (auto& input : graph.GetInputs()) { - result.inputs_.push_back(tvm_tensors[input->Name()]); - } - - // check initializer - for (auto& initializer : graph.GetAllInitializedTensors()) { - result.inputs_.push_back(tvm_tensors[initializer.first]); - } - - auto& output = graph.GetOutputs()[0]; - result.outputs_.push_back(tvm_tensors[output->Name()]); - return result; -} -} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/tvm/tvm_compiler.h b/onnxruntime/core/codegen/tvm/tvm_compiler.h deleted file mode 100644 index e4eed0dc80d94..0000000000000 --- a/onnxruntime/core/codegen/tvm/tvm_compiler.h +++ /dev/null @@ -1,36 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once -#include -#include -#include "core/common/common.h" -#include "core/framework/data_types.h" -#include "core/graph/function.h" -#include "core/graph/constants.h" -#include "core/graph/graph_viewer.h" -namespace onnxruntime { - -//TODO: this is just initial design to represent TVM Graph, to make the basic test work. -//We may need to revisit it later to finialize it. -struct TVMGraph { - struct TensorDescriptor { - tvm::Tensor tvm_tensor_; - DLContext ctx_; - DLDataType dtype_; - - public: - TensorDescriptor(MLDataType type, onnxruntime::ProviderType execution_provider_type, tvm::Tensor tvm_tensor); - - TensorDescriptor() = default; - }; - std::vector inputs_; - std::vector outputs_; -}; - -//TODO: compile a onnxruntime graph to tvm's tensor expression is a common logic for all hardwares -//onnxruntime framework should provide this functionality to executionp providers. -//We will need to register how to compiler it for each node. A detail design is needed. -//Here for testing we just provide the functionality that compile add 1D tensors. -TVMGraph CompileToTVM(const onnxruntime::Graph& graph, onnxruntime::ProviderType execution_provider_type); -} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/tvm/tvm_kernel.h b/onnxruntime/core/codegen/tvm/tvm_kernel.h deleted file mode 100644 index 59d997e02d2f1..0000000000000 --- a/onnxruntime/core/codegen/tvm/tvm_kernel.h +++ /dev/null @@ -1,126 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once -#include "core/codegen/tvm/tvm_compiler.h" -#include "core/graph/function.h" -#include "core/framework/op_kernel.h" - -namespace onnxruntime { - -// TVMScheduleCreator is the function that create a tvm schedule based on given TVM graph. -// Different hardware may have different schedule strategy. -typedef tvm::Schedule (*TVMScheduleCreator)(const TVMGraph&); -// TVMModuleBuilder is the function that build a tvm module, given a schedule and args. -// Different tvm kernel may chose different way to build the module, like target to LLVM or other backend. -typedef tvm::runtime::Module (*TVMModuleBuilder)(tvm::Schedule schedule, tvm::BuildConfig config, tvm::Array args, std::vector& target_func_names); - -template -class TVMKernel : public OpKernel { - public: - explicit TVMKernel(const OpKernelInfo& info) : OpKernel(info), tvm_values_(nullptr), dl_tensors_(nullptr), tvm_type_codes_(nullptr) { - auto& node = info.node(); - ORT_ENFORCE(node.NodeType() == Node::Type::Fused); - auto func = node.GetFunctionBody(); - const onnxruntime::Graph& func_body = func->Body(); - //1. compile the onnxruntime Graph to tvm graph. This step is common for all hardware, and provided by onnxruntime framework. - tvm_graph_ = CompileToTVM(func_body, node.GetExecutionProviderType()); - //2. create schedule for tvm graph, this step is depends on the execution provider/hardware. - auto s = S(tvm_graph_); - //3. Build module - std::vector tvm_args; - for (auto& t : tvm_graph_.inputs_) { - tvm_args.push_back(t.tvm_tensor_); - } - for (auto& t : tvm_graph_.outputs_) { - tvm_args.push_back(t.tvm_tensor_); - } - - std::vector func_names; - tvm_module_ = M(s, tvm::build_config(), tvm_args, func_names); - //TODO: do we have case that need more than 1 evaluation function? - evaluate_func_ = tvm_module_.GetFunction(func_names[0]); - //4. prepare args according to the type - n_args_ = tvm_args.size(); - tvm_values_ = new TVMValue[n_args_]; - tvm_type_codes_ = new int[n_args_]; - dl_tensors_ = new DLTensor[n_args_]; - int i = 0; - for (auto& tensor : tvm_graph_.inputs_) { - tvm_type_codes_[i] = kNDArrayContainer; - dl_tensors_[i].ctx = tensor.ctx_; - dl_tensors_[i].dtype = tensor.dtype_; - dl_tensors_[i].strides = nullptr; - dl_tensors_[i].byte_offset = 0; - tvm_values_[i].v_handle = &dl_tensors_[i]; - i++; - } - - for (auto& tensor : tvm_graph_.outputs_) { - tvm_type_codes_[i] = kNDArrayContainer; - dl_tensors_[i].ctx = tensor.ctx_; - dl_tensors_[i].dtype = tensor.dtype_; - dl_tensors_[i].strides = nullptr; - dl_tensors_[i].byte_offset = 0; - tvm_values_[i].v_handle = &dl_tensors_[i]; - i++; - } - ORT_ENFORCE(i == n_args_); - } - - virtual ~TVMKernel() { - if (!tvm_values_) - delete[] tvm_values_; - if (!tvm_type_codes_) - delete[] tvm_type_codes_; - if (!dl_tensors_) - delete[] dl_tensors_; - } - - virtual Status Compute(OpKernelContext* context) const override { - for (int i = 0; i < tvm_graph_.inputs_.size(); ++i) { - auto t = context->Input(i); - dl_tensors_[i].data = const_cast(t)->MutableDataRaw(); - dl_tensors_[i].ndim = static_cast(t->Shape().NumDimensions()); - dl_tensors_[i].shape = dl_tensors_[i].ndim > 0 ? const_cast(&(t->Shape().GetDims()[0])) : nullptr; - } - - int num_inputs = static_cast(tvm_graph_.inputs_.size()); - - for (int i = 0; i < tvm_graph_.outputs_.size(); ++i) { - //TODO: we need to have a shape inference function that could calculate output shape based on the symbolic formular in tvm - //We could build that function as part of tvm module, or reuse the shape inference in onnx funciton. - //Here for testing purpose, assume the output shape is same to input shape. - auto t = context->Output(i, GetOutputShape(context, i)); - dl_tensors_[num_inputs + i].data = t->MutableDataRaw(); - dl_tensors_[num_inputs + i].ndim = static_cast(t->Shape().NumDimensions()); - dl_tensors_[num_inputs + i].shape = dl_tensors_[i].ndim > 0 ? const_cast(&(t->Shape().GetDims()[0])) : nullptr; - } - - tvm::TVMArgs tvm_args(&tvm_values_[0], &tvm_type_codes_[0], static_cast(n_args_)); - tvm::TVMRetValue rvalue; - try { - evaluate_func_.CallPacked(tvm_args, &rvalue); - } catch (std::exception& ex) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "TVM run failed:", ex.what()); - } - if (rvalue.type_code() != kNull) { - return Status(onnxruntime::common::ONNXRUNTIME, onnxruntime::common::FAIL, "TVM return not null"); // TODO: get error code. - } else { - return Status::OK(); - } - } - - protected: - virtual const TensorShape& GetOutputShape(OpKernelContext* context, int i) const = 0; - - TVMGraph tvm_graph_; - tvm::runtime::Module tvm_module_; - tvm::PackedFunc evaluate_func_; - - size_t n_args_; - TVMValue* tvm_values_; - DLTensor* dl_tensors_; - int* tvm_type_codes_; -}; -} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/tvm/tvm_utils.cc b/onnxruntime/core/codegen/tvm/tvm_utils.cc deleted file mode 100644 index d1980ca1d10b4..0000000000000 --- a/onnxruntime/core/codegen/tvm/tvm_utils.cc +++ /dev/null @@ -1,32 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "tvm_utils.h" - -namespace onnxruntime { -namespace tvm_codegen { - -#define RETURN_DLDATATYPE_IF_MATCH(type, type_code) \ - if (ml_type == DataTypeImpl::GetType()) { \ - return {type_code, sizeof(type) * 8, 1}; \ - } - -// DLDataType: {DLDataTypeCode, bits, lanes} -DLDataType ToTvmDLDataType(MLDataType ml_type) { - RETURN_DLDATATYPE_IF_MATCH(int8_t, kDLInt); - RETURN_DLDATATYPE_IF_MATCH(uint8_t, kDLInt); - RETURN_DLDATATYPE_IF_MATCH(int16_t, kDLInt); - RETURN_DLDATATYPE_IF_MATCH(uint16_t, kDLInt); - RETURN_DLDATATYPE_IF_MATCH(int32_t, kDLInt); - RETURN_DLDATATYPE_IF_MATCH(uint32_t, kDLInt); - RETURN_DLDATATYPE_IF_MATCH(int64_t, kDLInt); - RETURN_DLDATATYPE_IF_MATCH(uint64_t, kDLInt); - - RETURN_DLDATATYPE_IF_MATCH(float, kDLFloat); - RETURN_DLDATATYPE_IF_MATCH(double, kDLFloat); - - ORT_NOT_IMPLEMENTED("converting MLDataType ", ml_type, " to tvm DLDataType is not implemented"); -} - -} // namespace tvm_codegen -} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/tvm/tvm_utils.h b/onnxruntime/core/codegen/tvm/tvm_utils.h deleted file mode 100644 index abe35753a07e0..0000000000000 --- a/onnxruntime/core/codegen/tvm/tvm_utils.h +++ /dev/null @@ -1,17 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once -#include - -#include "core/framework/data_types.h" - -namespace onnxruntime { - -constexpr const char* TVM_STACKVM = "TvmStackVm"; - -namespace tvm_codegen { - // Helper function that converts a onnxruntime MLDataType to TVM DLDataType - DLDataType ToTvmDLDataType(MLDataType ml_type); -} // namespace tvm -} // namespace onnxruntime diff --git a/onnxruntime/test/testdata/fuse_add_1.pb b/onnxruntime/test/testdata/fuse_add_1.pb deleted file mode 100644 index 801147fa4e277339c51bfcb7d9c5c014aa5c2868..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 407 zcmWe)OU_8mD@n9!W9E|JVu~>2VhT4DVv01BVs%VOi8oYY22voINRW&XP{s%%V}y_~ z0?HTzWsD&*#t0c>po|Gn#snf`0+tct$SuuLC`m0Y(X!%T6X13N+8AWO00tlIAvA*n zgyw+K0#I56N=rDv^%_y2*O&smCZy|)5&(s;5EmB*2e%Ll7ZV3t5>O=41X;w1g^NLe F9RM4MFMI$1 diff --git a/onnxruntime/test/testdata/fuse_mul_1.pb b/onnxruntime/test/testdata/fuse_mul_1.pb new file mode 100644 index 0000000000000000000000000000000000000000..a638551249f640c3a16a3ce6ad7269e77c4f326e GIT binary patch literal 169 zcmWe)OU_8mD@nAP!N?`S#S~!(!a_`uhElA)r8)72O3XkCBohgfF@ne#A!Ll8dW<15 z#t0c>sEi3j#sn-Q#F1N?qfnAsToNS+mgnZ; -#include "core/codegen/tvm/tvm_kernel.h" -#include "core/framework/execution_provider.h" +#include "core/common/logging/logging.h" #include "core/framework/compute_capability.h" +#include "core/framework/execution_provider.h" +#include "core/framework/kernel_registry.h" +#include "core/framework/op_kernel.h" #include "core/graph/graph_viewer.h" #include "core/providers/cpu/cpu_execution_provider.h" #include "core/session/inference_session.h" #include "core/session/onnxruntime_cxx_api.h" -#include "core/common/logging/logging.h" #include "test/framework/test_utils.h" #include "test/test_environment.h" -#include "core/framework/op_kernel.h" -#include "core/framework/kernel_registry.h" +#include "test/tvm/tvm_demo/demo_compiler.h" -namespace onnxruntime { +#include -tvm::Schedule DefaultTVMScheduleGenerator(const TVMGraph& tvm_graph) { - std::vector args; - for (auto& tensor : tvm_graph.outputs_) - args.push_back(tensor.tvm_tensor_->op); - return tvm::create_schedule(args); -} +namespace onnxruntime { -tvm::runtime::Module BuildStackVMDefaultModule(tvm::Schedule schedule, - tvm::BuildConfig config, - tvm::Array tvm_args, - std::vector& target_func_names) { - auto target = tvm::target::stackvm(); - std::string func_name = "func"; - auto args = tvm::Array(tvm_args); - std::unordered_map binds; - auto lowered = lower(schedule, args, "func", binds, config); - target_func_names.push_back(func_name); - return build(lowered, target, tvm::Target(), config); -} +using namespace tvm_demo; -template -class TVMFuseAddKernels : public TVMKernel { +class TVMDemoKernel : public OpKernel { public: - explicit TVMFuseAddKernels(const OpKernelInfo& info) : TVMKernel(info) {} + explicit TVMDemoKernel(const OpKernelInfo& info) : OpKernel(info) {} protected: - virtual const TensorShape& GetOutputShape(OpKernelContext* context, int /*i*/) const override { + const TensorShape& GetOutputShape(OpKernelContext* context, int /*i*/) const { return context->Input(0)->Shape(); } }; @@ -100,28 +82,28 @@ class FuseExecutionProviderX : public CPUExecutionProvider { GetCapability(const onnxruntime::GraphViewer& graph_viewer, const std::vector& /*kernel_registries*/) const override { std::vector> result; - std::vector add_nodes; + std::vector fused_nodes; for (auto& node : graph_viewer.Nodes()) { - if (node.OpType() == "Add") { - add_nodes.push_back(node.Index()); + if (node.OpType() == "Mul") { + fused_nodes.push_back(node.Index()); } } - UnionSet set(static_cast(add_nodes.size())); - for (int i = 0; i < add_nodes.size(); ++i) { - auto node = graph_viewer.GetNode(add_nodes[i]); + UnionSet set(static_cast(fused_nodes.size())); + for (int i = 0; i < fused_nodes.size(); ++i) { + auto node = graph_viewer.GetNode(fused_nodes[i]); for (auto it = node->InputNodesBegin(); it != node->InputNodesEnd(); ++it) { - auto index_it = std::find(add_nodes.begin(), add_nodes.end(), (*it).Index()); - if (index_it != add_nodes.end()) { - set.merge(i, static_cast(index_it - add_nodes.begin())); + auto index_it = std::find(fused_nodes.begin(), fused_nodes.end(), (*it).Index()); + if (index_it != fused_nodes.end()) { + set.merge(i, static_cast(index_it - fused_nodes.begin())); } } } std::vector> groups; - groups.resize(add_nodes.size()); + groups.resize(fused_nodes.size()); for (int i = 0; i < set.farthers_.size(); ++i) { - groups[set.get(i)].push_back(add_nodes[i]); + groups[set.get(i)].push_back(fused_nodes[i]); } for (auto& group : groups) { @@ -150,7 +132,7 @@ class FuseExecutionProviderX : public CPUExecutionProvider { } auto meta_def = std::make_unique<::onnxruntime::IndexedSubGraph::MetaDef>(); - meta_def->name = "TVMFuseAdd"; + meta_def->name = "TVMFuse"; meta_def->domain = "FuseTest"; for (auto input : fused_inputs) { meta_def->inputs.push_back(input->Name()); @@ -177,22 +159,22 @@ class FuseExecutionProviderX : public CPUExecutionProvider { auto func_body = fused_node->GetFunctionBody(); if (!func_body) return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "Function body is empty"); - //1. compile the onnxruntime Graph to tvm graph. - auto tvm_graph_ = CompileToTVM(func_body->Body(), kCpuExecutionProvider); - //2. create schedule for tvm graph, this step is depends on the execution provider/hardware. - auto s = DefaultTVMScheduleGenerator(tvm_graph_); - //3. Build module + //1. Build tvm IR based on the Ort graph + auto demo_tvm_tensor_ctx = BuildTVMIR(func_body->Body()); + //2. Create schedule for the built tvm IRs + auto s = CreateSchedule(demo_tvm_tensor_ctx); + //3. Build tvm module std::vector tvm_args; - for (auto& t : tvm_graph_.inputs_) { - tvm_args.push_back(t.tvm_tensor_); + for (auto& t : demo_tvm_tensor_ctx.inputs) { + tvm_args.push_back(t); } - for (auto& t : tvm_graph_.outputs_) { - tvm_args.push_back(t.tvm_tensor_); + for (auto& t : demo_tvm_tensor_ctx.outputs) { + tvm_args.push_back(t); } std::vector func_names; auto module_ptr = std::make_shared(); - *module_ptr = BuildStackVMDefaultModule(s, tvm::build_config(), tvm_args, func_names); + *module_ptr = BuildStackVMModule(s, tvm::build_config(), tvm_args, func_names); modules_[fused_node->Name()] = module_ptr; NodeComputeInfo compute_info; @@ -307,7 +289,9 @@ static void RunSession(InferenceSession& session_object, // Now run common::Status st = session_object.Run(run_options, feeds, output_names, &fetches); - std::cout << "Run returned status: " << st.ErrorMessage() << std::endl; + if (!st.IsOK()) { + std::cout << "Run returned status: " << st.ErrorMessage() << std::endl; + } EXPECT_TRUE(st.IsOK()); ASSERT_EQ(1, fetches.size()); auto& rtensor = fetches.front().Get(); @@ -319,9 +303,9 @@ static void RunSession(InferenceSession& session_object, ASSERT_EQ(found[i], values_y[i]); } -static const std::string MODEL_URI = "testdata/fuse_add_1.pb"; +static const std::string MODEL_URI = "testdata/fuse_mul_1.pb"; -TEST(TVMTest, Fuse_Add_Test) { +TEST(TVMTest, CodeGen_Demo_for_Fuse_Mul) { SessionOptions so; so.session_logid = "InferenceSessionTests.NoTimeout"; @@ -346,8 +330,8 @@ TEST(TVMTest, Fuse_Add_Test) { std::vector expected_dims_y = { 6, }; - // now the expected value should be Add's result. - std::vector expected_values_y = {5.0, 10.0, 15.0, 20.0, 25.0, 30.0}; + // now the expected value should be Mul's result. + std::vector expected_values_y = {1.0, 32.0, 243.0, 1024.0, 3125.0, 7776.0}; // Now run RunSession(session_object, run_options, dims_x, values_x, expected_dims_y, expected_values_y); @@ -356,7 +340,7 @@ TEST(TVMTest, Fuse_Add_Test) { } // namespace onnxruntime -TEST(TVMTest, Basic) { +TEST(TVMTest, Native_TVM) { using namespace tvm; auto n = var("n"); Array shape; diff --git a/onnxruntime/test/tvm/tvm_demo/demo_compiler.cc b/onnxruntime/test/tvm/tvm_demo/demo_compiler.cc new file mode 100644 index 0000000000000..a7c2a47ef405d --- /dev/null +++ b/onnxruntime/test/tvm/tvm_demo/demo_compiler.cc @@ -0,0 +1,226 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "test/tvm/tvm_demo/demo_compiler.h" + +#include "core/codegen/target/generic/scheduler/schedule_utils.h" +#include "core/codegen/target/ort_tvm_utils.h" +#include "core/codegen/target/tvm_ir_builder.h" +#include "core/codegen/target/tvm_scheduler.h" +#include "core/codegen/target/tvm_schedule_builder.h" + +#include +#include + +namespace onnxruntime { +namespace tvm_demo { + +// Create a dummy demo handle +static codegen::CodeGenHandle demo_handle; +// Create a dummy demo codegen context +static tvm_codegen::CodeGenContext demo_codegen_ctx(&demo_handle); + +// Translate an Ort graph into tvm IR +// Note this function is specific for this demo. +// This function uses specific way for graph traversal or constructing tvm placeholders. +// It may or may not work for a universal Ort graph. +// For a more general example, please check nuphar provider. +DemoTVMTensorCtx BuildTVMIR(const onnxruntime::Graph& graph) { + // Create OpIRRegistry that holds all OpIRCreators + std::unique_ptr op_ir_registry = + std::make_unique(); + + // Register all generic OpIRCreators + tvm_codegen::RegisterAllGenericOpIRCreators(op_ir_registry.get()); + + // Create OpIRBuilder + std::shared_ptr op_ir_builder = + std::make_shared("Demo_Op_IR_Builder"); + + // Attach all generic OpIRCreators from op_ir_registry to op_ir_builder + tvm_codegen::RegisterGenericOrtOpTypeDispatcher(op_ir_builder, op_ir_registry.get()); + + // Create DemoTVMTensorCtx holdings tvm IR + DemoTVMTensorCtx result; + + // Local lookup from name to tvm::Tensor + std::unordered_map tvm_tensors; + + // Note this is a simplified traversal that works specifically for this demo + // but may or may not work for an univerisal model. + // For more general traversal, please check nuphar provider. + for (auto& node : graph.Nodes()) { + tvm::Array inputs; + tvm::Array outputs; + + // Get inputs + for (auto& def : node.InputDefs()) { + const std::string& name = def->Name(); + auto iter = tvm_tensors.find(name); + // Always create placeholder when not finding a tensor + // Note it is for this demo. + // It may or may not work for a universal graph. + if (iter == tvm_tensors.end()) { + tvm_tensors[name] = + tvm::placeholder(ShapeToTvmArray(def, demo_codegen_ctx), + tvm_codegen::ToTvmType(TensorProtoDataType(def)), + name + "_placeholder"); + } + inputs.push_back(tvm_tensors[name]); + } + + // call OpIBuilder's Evaluate to build tvm IR + op_ir_builder->Evaluate(inputs, node, demo_codegen_ctx, outputs); + + // Store outputs + for (size_t def_id = 0; def_id < node.OutputDefs().size(); ++def_id) { + const NodeArg* def = node.OutputDefs()[def_id]; + tvm_tensors[def->Name()] = outputs[def_id]; + } + } + + // put inputs to DemoTVMTensorCtx + for (auto& input : graph.GetInputs()) { + result.inputs.push_back(tvm_tensors[input->Name()]); + } + + // check initializer + for (auto& initializer : graph.GetAllInitializedTensors()) { + result.inputs.push_back(tvm_tensors[initializer.first]); + } + + // Only one output in this demo + auto& output = graph.GetOutputs()[0]; + result.outputs.push_back(tvm_tensors[output->Name()]); + return result; +} + +// Declare a Demo scheduler that always inserts compute_inline +DECLARE_TVM_SCHEDULER_CLASS(AlwaysInline, DemoTVM) + +// Define a Demo scheduler's Evaluate that always inserts compute_inline +bool TVM_SCHEDULER_CLASS(AlwaysInline, DemoTVM)::Evaluate( + const tvm::Tensor& tensor, + const Node*, + tvm_codegen::CodeGenContext&, + tvm_codegen::ScheduleContext& ctx_sched) { + return TryInlineSchedule(tensor, ctx_sched); +} + +// Register the always inline Scheduler to sched_registry +static void RegisterAlwaysInlineScheduler(tvm_codegen::TVMScheduleRegistry* sched_registry) { + sched_registry->Register( + std::move(std::make_unique())); +} + +// Declare a schedule dispatcher that always dispatches the always inline Scheduler +DECLARE_SCHEDULE_DISPATCHER_CLASS(DemoTVM) + +// Use a predefined key as DemoKey to dispatch the scheduler +constexpr auto predefined_key = "DemoKey"; + +// Define the schedule dispatcher's Find function +// that always dispatches the always inline Scheduler +// Note this dispatcher always returning a predefined_key is only for demo purpose. +// In practice, a dispatcher returns a key by checking tvm::Tensor, Node, +// or even meta data stored in CodeGenContext. +// Derived CodeGenContext allows compiler developers to store their specific meta data. +// For more detailed example, please check nuphar provider. +tvm_codegen::Scheduler* SCHEDULE_DISPATCHER_CLASS(DemoTVM)::Find( + const tvm::Tensor&, const Node*, tvm_codegen::CodeGenContext&) { + return DispatcherBase::Get(predefined_key); +} + +// Attach the always inline Scheduler to the above dispatcher +// and then attach the dispatcher to the scheduler builder +static void AttachAlwaysInlineScheduler(const std::shared_ptr& builder, + const tvm_codegen::TVMScheduleRegistry* registry) { + auto dispatcher = std::make_unique("DemoSchedulers"); + + // Using a predefined_key + dispatcher->Register(predefined_key, + registry->Get(TVM_SCHEDULER_STRING(AlwaysInline, DemoTVM))); + + builder->InsertDispatcher(std::move(dispatcher)); +} + +// Traverse tvm::Tensor and then schedule them +// Note this traversal is simplified and specific for this demo. +// For a more general traversal, please check nuphar provider. +static void TraverseAndSchedule( + std::shared_ptr& schedule_builder, + const tvm::Tensor& tensor, + tvm_codegen::ScheduleContext& ctx_schedule) { + schedule_builder->Evaluate(tensor, nullptr, demo_codegen_ctx, ctx_schedule); + + // Traverse tensor's children (inputs) + for (auto& t : tensor->op->InputTensors()) { + // check whether it is a non-trivial tensor by checking its input size + if (t->op->InputTensors().size() > 0) { + TraverseAndSchedule(schedule_builder, t, ctx_schedule); + } + } +} + +// Create a TVM schedule by always inserting tvm's compute_inline. +// Note this schedule is specific for this demo. +// In practice, always inline might lead to bad performance +// or even illegal loop transformation for some backends. +// For a more general example, please check nuphar provider. +tvm::Schedule CreateSchedule(const DemoTVMTensorCtx& ctx) { + // Create TVMScheduleRegistry that holds all Scheduler + std::unique_ptr schedule_registry = + std::make_unique(); + + // Register the always inline Scheduler to schedule_registry + RegisterAlwaysInlineScheduler(schedule_registry.get()); + + // Create a DemoScheduleBuilder + std::shared_ptr schedule_builder = + std::make_shared("Demo_Schedule_Builder"); + + // Attach the demo inline scheduler to the schedule_builder + AttachAlwaysInlineScheduler(schedule_builder, schedule_registry.get()); + + // Create scheudule object + tvm::Array out_ops; + for (auto& t : ctx.outputs) { + out_ops.push_back(t->op); + } + + // Create scheudule context + tvm_codegen::ScheduleContext ctx_schedule(out_ops); + + // Traverse tvm::Tensor in a DFS way, and then schedule + for (auto& t : ctx.outputs) { + TraverseAndSchedule(schedule_builder, t, ctx_schedule); + } + + // Make sure all outputs compute_root (tvm's requirement) + for (auto& t : ctx.outputs) { + tvm_codegen::InsertRootSchedule(t, ctx_schedule); + } + + return ctx_schedule.schedule; +} + +// Build TVM Module with a schedule using tvm's stackvm. +// Note in real practice, please change stackvm to other backends. +// For a more detailed example, please check nuphar provider. +tvm::runtime::Module BuildStackVMModule(tvm::Schedule schedule, + tvm::BuildConfig config, + tvm::Array tvm_args, + std::vector& target_func_names) { + auto target = tvm::target::stackvm(); + std::string func_name = "func"; + auto args = tvm::Array(tvm_args); + std::unordered_map binds; + auto lowered = lower(schedule, args, "func", binds, config); + // Uncomment the following line to dump lowered func + // std::cout << "Dumping lowered func: " << lowered[0]->body; + target_func_names.push_back(func_name); + return build(lowered, target, tvm::Target(), config); +} + +} // namespace tvm_demo +} // namespace onnxruntime diff --git a/onnxruntime/test/tvm/tvm_demo/demo_compiler.h b/onnxruntime/test/tvm/tvm_demo/demo_compiler.h new file mode 100644 index 0000000000000..3905995293152 --- /dev/null +++ b/onnxruntime/test/tvm/tvm_demo/demo_compiler.h @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/common/common.h" +#include "core/graph/graph_viewer.h" +#include +#include + +namespace onnxruntime { +namespace tvm_demo { +// A Demo data structure to hold tvm IR and context +struct DemoTVMTensorCtx { + tvm::Array inputs; + tvm::Array outputs; +}; + +// Translate an Ort graph into tvm IR +DemoTVMTensorCtx BuildTVMIR(const onnxruntime::Graph& graph); + +// Create a demo schedule for the tvm IR +tvm::Schedule CreateSchedule(const DemoTVMTensorCtx& ctx); + +// Build a demo tvm module with the tvm IR and schedule +tvm::runtime::Module BuildStackVMModule(tvm::Schedule schedule, + tvm::BuildConfig config, + tvm::Array tvm_args, + std::vector& target_func_names); + +} // namespace tvm_demo +} // namespace onnxruntime From 68d24431aa232c73323a3fba0141f306bb7184d7 Mon Sep 17 00:00:00 2001 From: KeDengMS Date: Fri, 28 Jun 2019 22:15:26 -0700 Subject: [PATCH 2/2] Address CR --- cmake/onnxruntime_codegen.cmake | 4 ++-- onnxruntime/core/codegen/mti/nn/conv_ops.cc | 6 ------ .../{target/generic => passes}/op_ir_creator/all_ops.h | 4 ++-- .../op_ir_creator/math/binary_ops.cc | 2 +- .../generic => passes}/op_ir_creator/math/clip.cc | 2 +- .../generic => passes}/op_ir_creator/math/gemm.cc | 2 +- .../op_ir_creator/math/logsoftmax.cc | 2 +- .../generic => passes}/op_ir_creator/math/matmul.cc | 2 +- .../op_ir_creator/math/quantize/matmul_integer.cc | 2 +- .../op_ir_creator/math/reduce_ops.cc | 2 +- .../generic => passes}/op_ir_creator/math/softmax.cc | 2 +- .../generic => passes}/op_ir_creator/math/unary_ops.cc | 2 +- .../op_ir_creator/math/variadic_ops.cc | 2 +- .../generic => passes}/op_ir_creator/nn/conv.cc | 4 ++-- .../generic => passes}/op_ir_creator/nn/lstm.cc | 2 +- .../generic => passes}/op_ir_creator/nn/pool_ops.cc | 2 +- .../generic => passes}/op_ir_creator/tensor/cast.cc | 4 ++-- .../generic => passes}/op_ir_creator/tensor/concat.cc | 2 +- .../generic => passes}/op_ir_creator/tensor/crop.cc | 2 +- .../generic => passes}/op_ir_creator/tensor/gather.cc | 2 +- .../generic => passes}/op_ir_creator/tensor/pad.cc | 2 +- .../op_ir_creator/tensor/reshape_ops.cc | 4 ++-- .../generic => passes}/op_ir_creator/tensor/slice.cc | 4 ++-- .../generic => passes}/op_ir_creator/tensor/split.cc | 2 +- .../op_ir_creator/tensor/transpose.cc | 2 +- .../generic => passes}/op_ir_creator/tensor/where.cc | 2 +- .../{target => passes/op_ir_creator}/tvm_ir_builder.cc | 4 ++-- .../{target => passes/op_ir_creator}/tvm_ir_builder.h | 4 ++-- .../{target => passes/op_ir_creator}/tvm_op_creator.cc | 4 ++-- .../{target => passes/op_ir_creator}/tvm_op_creator.h | 0 .../generic => passes}/scheduler/all_schedules.h | 2 +- .../generic => passes}/scheduler/ort_type_schedule.cc | 4 ++-- .../generic => passes}/scheduler/schedule_utils.cc | 2 +- .../generic => passes}/scheduler/schedule_utils.h | 2 +- .../generic => passes}/scheduler/tvm_rule_schedule.cc | 4 ++-- .../scheduler}/tvm_schedule_builder.cc | 2 +- .../scheduler}/tvm_schedule_builder.h | 2 +- .../{target => passes/scheduler}/tvm_scheduler.cc | 4 ++-- .../{target => passes/scheduler}/tvm_scheduler.h | 2 +- .../{target => passes/utils}/codegen_context.cc | 2 +- .../codegen/{target => passes/utils}/codegen_context.h | 0 .../codegen/{target => passes/utils}/ort_tvm_utils.cc | 4 ++-- .../codegen/{target => passes/utils}/ort_tvm_utils.h | 0 .../generic => passes}/weight_layout/transpose_2d.cc | 4 ++-- .../generic => passes}/weight_layout/transpose_2d.h | 2 +- .../weight_layout/vertical_stripes_2d.cc | 4 ++-- .../weight_layout/vertical_stripes_2d.h | 2 +- .../{target => passes/weight_layout}/weight_layout.cc | 4 ++-- .../{target => passes/weight_layout}/weight_layout.h | 0 onnxruntime/test/tvm/tvm_demo/demo_compiler.cc | 10 +++++----- 50 files changed, 65 insertions(+), 71 deletions(-) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/all_ops.h (91%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/math/binary_ops.cc (97%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/math/clip.cc (93%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/math/gemm.cc (94%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/math/logsoftmax.cc (93%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/math/matmul.cc (90%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/math/quantize/matmul_integer.cc (95%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/math/reduce_ops.cc (98%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/math/softmax.cc (93%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/math/unary_ops.cc (98%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/math/variadic_ops.cc (95%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/nn/conv.cc (97%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/nn/lstm.cc (97%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/nn/pool_ops.cc (98%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/tensor/cast.cc (91%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/tensor/concat.cc (92%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/tensor/crop.cc (95%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/tensor/gather.cc (92%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/tensor/pad.cc (96%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/tensor/reshape_ops.cc (96%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/tensor/slice.cc (97%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/tensor/split.cc (97%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/tensor/transpose.cc (95%) rename onnxruntime/core/codegen/{target/generic => passes}/op_ir_creator/tensor/where.cc (92%) rename onnxruntime/core/codegen/{target => passes/op_ir_creator}/tvm_ir_builder.cc (96%) rename onnxruntime/core/codegen/{target => passes/op_ir_creator}/tvm_ir_builder.h (94%) rename onnxruntime/core/codegen/{target => passes/op_ir_creator}/tvm_op_creator.cc (90%) rename onnxruntime/core/codegen/{target => passes/op_ir_creator}/tvm_op_creator.h (100%) rename onnxruntime/core/codegen/{target/generic => passes}/scheduler/all_schedules.h (90%) rename onnxruntime/core/codegen/{target/generic => passes}/scheduler/ort_type_schedule.cc (80%) rename onnxruntime/core/codegen/{target/generic => passes}/scheduler/schedule_utils.cc (98%) rename onnxruntime/core/codegen/{target/generic => passes}/scheduler/schedule_utils.h (96%) rename onnxruntime/core/codegen/{target/generic => passes}/scheduler/tvm_rule_schedule.cc (88%) rename onnxruntime/core/codegen/{target => passes/scheduler}/tvm_schedule_builder.cc (97%) rename onnxruntime/core/codegen/{target => passes/scheduler}/tvm_schedule_builder.h (95%) rename onnxruntime/core/codegen/{target => passes/scheduler}/tvm_scheduler.cc (95%) rename onnxruntime/core/codegen/{target => passes/scheduler}/tvm_scheduler.h (98%) rename onnxruntime/core/codegen/{target => passes/utils}/codegen_context.cc (92%) rename onnxruntime/core/codegen/{target => passes/utils}/codegen_context.h (100%) rename onnxruntime/core/codegen/{target => passes/utils}/ort_tvm_utils.cc (98%) rename onnxruntime/core/codegen/{target => passes/utils}/ort_tvm_utils.h (100%) rename onnxruntime/core/codegen/{target/generic => passes}/weight_layout/transpose_2d.cc (93%) rename onnxruntime/core/codegen/{target/generic => passes}/weight_layout/transpose_2d.h (93%) rename onnxruntime/core/codegen/{target/generic => passes}/weight_layout/vertical_stripes_2d.cc (94%) rename onnxruntime/core/codegen/{target/generic => passes}/weight_layout/vertical_stripes_2d.h (95%) rename onnxruntime/core/codegen/{target => passes/weight_layout}/weight_layout.cc (96%) rename onnxruntime/core/codegen/{target => passes/weight_layout}/weight_layout.h (100%) diff --git a/cmake/onnxruntime_codegen.cmake b/cmake/onnxruntime_codegen.cmake index 56ce2f416db97..df90e36cbce6a 100644 --- a/cmake/onnxruntime_codegen.cmake +++ b/cmake/onnxruntime_codegen.cmake @@ -9,8 +9,8 @@ file(GLOB_RECURSE onnxruntime_codegen_common_srcs file(GLOB_RECURSE onnxruntime_codegen_tvm_srcs CONFIGURE_DEPENDS "${ONNXRUNTIME_ROOT}/core/codegen/mti/*.h" "${ONNXRUNTIME_ROOT}/core/codegen/mti/*.cc" - "${ONNXRUNTIME_ROOT}/core/codegen/target/*.h" - "${ONNXRUNTIME_ROOT}/core/codegen/target/*.cc" + "${ONNXRUNTIME_ROOT}/core/codegen/passes/*.h" + "${ONNXRUNTIME_ROOT}/core/codegen/passes/*.cc" ) source_group(TREE ${ONNXRUNTIME_ROOT}/core FILES ${onnxruntime_codegen_common_srcs} ${onnxruntime_codegen_tvm_srcs}) diff --git a/onnxruntime/core/codegen/mti/nn/conv_ops.cc b/onnxruntime/core/codegen/mti/nn/conv_ops.cc index ca2fd3ee897c6..e2d4acc8843ad 100644 --- a/onnxruntime/core/codegen/mti/nn/conv_ops.cc +++ b/onnxruntime/core/codegen/mti/nn/conv_ops.cc @@ -68,12 +68,6 @@ tvm::Tensor Conv2D(const tvm::Tensor& input, const tvm::Array& stride, const tvm::Array& padding, const std::string& name) { - // Gemm Convolution - const int64_t* batch_size = tvm::as_const_int(input->shape[0]); - if (batch_size != nullptr && *batch_size == 1) - return Conv2D_gemm(input, filter, output_shape, stride, padding); - - // Native Convolution return Conv2D_native(input, filter, output_shape, stride, padding); } diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/all_ops.h b/onnxruntime/core/codegen/passes/op_ir_creator/all_ops.h similarity index 91% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/all_ops.h rename to onnxruntime/core/codegen/passes/op_ir_creator/all_ops.h index eeba129a7d597..1463e50bd72fb 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/all_ops.h +++ b/onnxruntime/core/codegen/passes/op_ir_creator/all_ops.h @@ -2,9 +2,9 @@ // Licensed under the MIT License. #pragma once -#include "core/codegen/target/codegen_context.h" +#include "core/codegen/passes/utils/codegen_context.h" #include "core/codegen/common/op_macro.h" -#include "core/codegen/target/tvm_op_creator.h" +#include "core/codegen/passes/op_ir_creator/tvm_op_creator.h" namespace onnxruntime { namespace tvm_codegen { diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/binary_ops.cc b/onnxruntime/core/codegen/passes/op_ir_creator/math/binary_ops.cc similarity index 97% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/math/binary_ops.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/math/binary_ops.cc index 771b2f18d1dc8..9452146621ac7 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/binary_ops.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/math/binary_ops.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/common/op_macro.h" #include "core/codegen/mti/math/binary_ops.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/clip.cc b/onnxruntime/core/codegen/passes/op_ir_creator/math/clip.cc similarity index 93% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/math/clip.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/math/clip.cc index 3551af6682828..88383624f87b5 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/clip.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/math/clip.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/math/unary_ops.h" #include "core/framework/op_kernel_info.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/gemm.cc b/onnxruntime/core/codegen/passes/op_ir_creator/math/gemm.cc similarity index 94% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/math/gemm.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/math/gemm.cc index b8628511e6308..64f995076e1bb 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/gemm.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/math/gemm.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/math/gemm.h" #include "core/framework/op_kernel_info.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/logsoftmax.cc b/onnxruntime/core/codegen/passes/op_ir_creator/math/logsoftmax.cc similarity index 93% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/math/logsoftmax.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/math/logsoftmax.cc index 4543e7b6f6f20..cb09518bf63d1 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/logsoftmax.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/math/logsoftmax.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/math/logsoftmax.h" #include "core/framework/op_kernel_info.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/matmul.cc b/onnxruntime/core/codegen/passes/op_ir_creator/math/matmul.cc similarity index 90% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/math/matmul.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/math/matmul.cc index 6dd2bb840e067..ab1ac237bfa5d 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/matmul.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/math/matmul.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/math/matmul_ops.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/quantize/matmul_integer.cc b/onnxruntime/core/codegen/passes/op_ir_creator/math/quantize/matmul_integer.cc similarity index 95% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/math/quantize/matmul_integer.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/math/quantize/matmul_integer.cc index df67c3cb8ca7a..60841d049e734 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/quantize/matmul_integer.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/math/quantize/matmul_integer.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/math/binary_ops.h" #include "core/codegen/mti/math/matmul_ops.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/reduce_ops.cc b/onnxruntime/core/codegen/passes/op_ir_creator/math/reduce_ops.cc similarity index 98% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/math/reduce_ops.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/math/reduce_ops.cc index 1773dd11fbcc3..f29a3f3e7cdf7 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/reduce_ops.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/math/reduce_ops.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/common/op_macro.h" #include "core/codegen/mti/math/reduce_ops.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/softmax.cc b/onnxruntime/core/codegen/passes/op_ir_creator/math/softmax.cc similarity index 93% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/math/softmax.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/math/softmax.cc index b8ce32ca30f6e..7b13de5a94e48 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/softmax.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/math/softmax.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/math/softmax.h" #include "core/framework/op_kernel_info.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/unary_ops.cc b/onnxruntime/core/codegen/passes/op_ir_creator/math/unary_ops.cc similarity index 98% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/math/unary_ops.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/math/unary_ops.cc index fc31097fc5453..bd5b89c718435 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/unary_ops.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/math/unary_ops.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/common/op_macro.h" #include "core/codegen/mti/math/unary_ops.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/variadic_ops.cc b/onnxruntime/core/codegen/passes/op_ir_creator/math/variadic_ops.cc similarity index 95% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/math/variadic_ops.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/math/variadic_ops.cc index ff719b82f85e5..9559a713c2876 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/math/variadic_ops.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/math/variadic_ops.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/math/binary_ops.h" #include "core/codegen/mti/tensor/reshape_ops.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/conv.cc b/onnxruntime/core/codegen/passes/op_ir_creator/nn/conv.cc similarity index 97% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/nn/conv.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/nn/conv.cc index 31ca33df8f335..c3a9e5950acce 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/conv.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/nn/conv.cc @@ -1,13 +1,13 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/nn/conv_ops.h" #include "core/codegen/mti/mti_tvm_utils.h" #include "core/codegen/mti/tensor/concat_ops.h" #include "core/codegen/mti/tensor/split.h" -#include "core/codegen/target/ort_tvm_utils.h" +#include "core/codegen/passes/utils/ort_tvm_utils.h" #include "core/framework/op_kernel_info.h" namespace onnxruntime { diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/lstm.cc b/onnxruntime/core/codegen/passes/op_ir_creator/nn/lstm.cc similarity index 97% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/nn/lstm.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/nn/lstm.cc index 0f244a69a3205..5c2557142dd0e 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/lstm.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/nn/lstm.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/nn/lstm.h" #include "core/framework/op_kernel_info.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/pool_ops.cc b/onnxruntime/core/codegen/passes/op_ir_creator/nn/pool_ops.cc similarity index 98% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/nn/pool_ops.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/nn/pool_ops.cc index e14a3f1bfb9f8..556d175a96601 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/nn/pool_ops.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/nn/pool_ops.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/mti_tvm_utils.h" #include "core/codegen/mti/nn/pool_ops.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/cast.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/cast.cc similarity index 91% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/cast.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/tensor/cast.cc index a3b9f5305d984..bd324fd359edf 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/cast.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/cast.cc @@ -1,10 +1,10 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/tensor/cast_ops.h" -#include "core/codegen/target/ort_tvm_utils.h" +#include "core/codegen/passes/utils/ort_tvm_utils.h" #include "core/framework/op_kernel_info.h" namespace onnxruntime { diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/concat.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/concat.cc similarity index 92% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/concat.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/tensor/concat.cc index c04603266eaa2..418296889419e 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/concat.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/concat.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/tensor/concat_ops.h" #include "core/framework/op_kernel_info.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/crop.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/crop.cc similarity index 95% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/crop.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/tensor/crop.cc index a14bba5f9f8a8..46adb7e984f2d 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/crop.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/crop.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/mti_tvm_utils.h" #include "core/codegen/mti/tensor/crop.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/gather.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/gather.cc similarity index 92% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/gather.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/tensor/gather.cc index 7a314b572b49e..3a5d801b6839f 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/gather.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/gather.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/tensor/gather.h" #include "core/framework/op_kernel_info.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/pad.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/pad.cc similarity index 96% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/pad.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/tensor/pad.cc index 159825b0dc088..ecff2c7b73847 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/pad.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/pad.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/mti_tvm_utils.h" #include "core/codegen/mti/tensor/pad_ops.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/reshape_ops.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/reshape_ops.cc similarity index 96% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/reshape_ops.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/tensor/reshape_ops.cc index a45d82475f2a7..ec5862a8a688c 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/reshape_ops.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/reshape_ops.cc @@ -1,11 +1,11 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/mti_tvm_utils.h" #include "core/codegen/mti/tensor/reshape_ops.h" -#include "core/codegen/target/ort_tvm_utils.h" +#include "core/codegen/passes/utils/ort_tvm_utils.h" #include "core/framework/op_kernel_info.h" namespace onnxruntime { diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/slice.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/slice.cc similarity index 97% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/slice.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/tensor/slice.cc index 2915dcc701cd1..4c27da39db0c4 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/slice.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/slice.cc @@ -1,8 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" -#include "core/codegen/target/ort_tvm_utils.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" +#include "core/codegen/passes/utils/ort_tvm_utils.h" #include "core/codegen/mti/mti_tvm_utils.h" #include "core/codegen/mti/tensor/slice.h" #include "core/framework/op_kernel_info.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/split.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/split.cc similarity index 97% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/split.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/tensor/split.cc index 5c82ad405e17a..7a190b5617042 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/split.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/split.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/mti_tvm_utils.h" #include "core/codegen/mti/tensor/split.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/transpose.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/transpose.cc similarity index 95% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/transpose.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/tensor/transpose.cc index d47ea30084aec..f4d7bb1da5e97 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/transpose.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/transpose.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/mti_tvm_utils.h" #include "core/codegen/mti/tensor/transpose.h" diff --git a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/where.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/where.cc similarity index 92% rename from onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/where.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/tensor/where.cc index e20df5f3b0382..9d6df7c1c430d 100644 --- a/onnxruntime/core/codegen/target/generic/op_ir_creator/tensor/where.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/where.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/codegen/mti/mti_tvm_utils.h" #include "core/codegen/mti/tensor/where.h" diff --git a/onnxruntime/core/codegen/target/tvm_ir_builder.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tvm_ir_builder.cc similarity index 96% rename from onnxruntime/core/codegen/target/tvm_ir_builder.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/tvm_ir_builder.cc index c1531c8cfacc2..6933681dda6c0 100644 --- a/onnxruntime/core/codegen/target/tvm_ir_builder.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tvm_ir_builder.cc @@ -1,10 +1,10 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/tvm_ir_builder.h" +#include "core/codegen/passes/op_ir_creator/tvm_ir_builder.h" #include "core/codegen/common/op_macro.h" -#include "core/codegen/target/generic/op_ir_creator/all_ops.h" +#include "core/codegen/passes/op_ir_creator/all_ops.h" #include "core/common/common.h" namespace onnxruntime { diff --git a/onnxruntime/core/codegen/target/tvm_ir_builder.h b/onnxruntime/core/codegen/passes/op_ir_creator/tvm_ir_builder.h similarity index 94% rename from onnxruntime/core/codegen/target/tvm_ir_builder.h rename to onnxruntime/core/codegen/passes/op_ir_creator/tvm_ir_builder.h index f7f69647e573b..c80056e619d6d 100644 --- a/onnxruntime/core/codegen/target/tvm_ir_builder.h +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tvm_ir_builder.h @@ -2,8 +2,8 @@ // Licensed under the MIT License. #pragma once -#include "core/codegen/target/codegen_context.h" -#include "core/codegen/target/tvm_op_creator.h" +#include "core/codegen/passes/utils/codegen_context.h" +#include "core/codegen/passes/op_ir_creator/tvm_op_creator.h" #include "core/common/common.h" namespace onnxruntime { diff --git a/onnxruntime/core/codegen/target/tvm_op_creator.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tvm_op_creator.cc similarity index 90% rename from onnxruntime/core/codegen/target/tvm_op_creator.cc rename to onnxruntime/core/codegen/passes/op_ir_creator/tvm_op_creator.cc index ef2da19c6d8de..992272753f5a4 100644 --- a/onnxruntime/core/codegen/target/tvm_op_creator.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tvm_op_creator.cc @@ -1,11 +1,11 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/tvm_op_creator.h" +#include "core/codegen/passes/op_ir_creator/tvm_op_creator.h" #include "core/codegen/common/common.h" #include "core/codegen/common/dispatcher.h" -#include "core/codegen/target/codegen_context.h" +#include "core/codegen/passes/utils/codegen_context.h" namespace onnxruntime { namespace codegen { diff --git a/onnxruntime/core/codegen/target/tvm_op_creator.h b/onnxruntime/core/codegen/passes/op_ir_creator/tvm_op_creator.h similarity index 100% rename from onnxruntime/core/codegen/target/tvm_op_creator.h rename to onnxruntime/core/codegen/passes/op_ir_creator/tvm_op_creator.h diff --git a/onnxruntime/core/codegen/target/generic/scheduler/all_schedules.h b/onnxruntime/core/codegen/passes/scheduler/all_schedules.h similarity index 90% rename from onnxruntime/core/codegen/target/generic/scheduler/all_schedules.h rename to onnxruntime/core/codegen/passes/scheduler/all_schedules.h index 54177009c6b72..cb75c7fa639c0 100644 --- a/onnxruntime/core/codegen/target/generic/scheduler/all_schedules.h +++ b/onnxruntime/core/codegen/passes/scheduler/all_schedules.h @@ -2,7 +2,7 @@ // Licensed under the MIT License. #pragma once -#include "core/codegen/target/tvm_scheduler.h" +#include "core/codegen/passes/scheduler/tvm_scheduler.h" namespace onnxruntime { namespace tvm_codegen { diff --git a/onnxruntime/core/codegen/target/generic/scheduler/ort_type_schedule.cc b/onnxruntime/core/codegen/passes/scheduler/ort_type_schedule.cc similarity index 80% rename from onnxruntime/core/codegen/target/generic/scheduler/ort_type_schedule.cc rename to onnxruntime/core/codegen/passes/scheduler/ort_type_schedule.cc index 4cd9efb313fc1..59f492d164b14 100644 --- a/onnxruntime/core/codegen/target/generic/scheduler/ort_type_schedule.cc +++ b/onnxruntime/core/codegen/passes/scheduler/ort_type_schedule.cc @@ -1,9 +1,9 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/scheduler/all_schedules.h" +#include "core/codegen/passes/scheduler/all_schedules.h" -#include "core/codegen/target/generic/scheduler/schedule_utils.h" +#include "core/codegen/passes/scheduler/schedule_utils.h" namespace onnxruntime { namespace tvm_codegen { diff --git a/onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.cc b/onnxruntime/core/codegen/passes/scheduler/schedule_utils.cc similarity index 98% rename from onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.cc rename to onnxruntime/core/codegen/passes/scheduler/schedule_utils.cc index 319bdfed92ea4..8f1485235697c 100644 --- a/onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.cc +++ b/onnxruntime/core/codegen/passes/scheduler/schedule_utils.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/scheduler/schedule_utils.h" +#include "core/codegen/passes/scheduler/schedule_utils.h" namespace onnxruntime { namespace tvm_codegen { diff --git a/onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.h b/onnxruntime/core/codegen/passes/scheduler/schedule_utils.h similarity index 96% rename from onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.h rename to onnxruntime/core/codegen/passes/scheduler/schedule_utils.h index f6e1bdb5c4657..f928928f30c57 100644 --- a/onnxruntime/core/codegen/target/generic/scheduler/schedule_utils.h +++ b/onnxruntime/core/codegen/passes/scheduler/schedule_utils.h @@ -3,7 +3,7 @@ #pragma once #include -#include +#include namespace onnxruntime { namespace tvm_codegen { diff --git a/onnxruntime/core/codegen/target/generic/scheduler/tvm_rule_schedule.cc b/onnxruntime/core/codegen/passes/scheduler/tvm_rule_schedule.cc similarity index 88% rename from onnxruntime/core/codegen/target/generic/scheduler/tvm_rule_schedule.cc rename to onnxruntime/core/codegen/passes/scheduler/tvm_rule_schedule.cc index 9f882e6c0b5ad..33162deddc983 100644 --- a/onnxruntime/core/codegen/target/generic/scheduler/tvm_rule_schedule.cc +++ b/onnxruntime/core/codegen/passes/scheduler/tvm_rule_schedule.cc @@ -1,9 +1,9 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/scheduler/all_schedules.h" +#include "core/codegen/passes/scheduler/all_schedules.h" -#include "core/codegen/target/generic/scheduler/schedule_utils.h" +#include "core/codegen/passes/scheduler/schedule_utils.h" namespace onnxruntime { namespace tvm_codegen { diff --git a/onnxruntime/core/codegen/target/tvm_schedule_builder.cc b/onnxruntime/core/codegen/passes/scheduler/tvm_schedule_builder.cc similarity index 97% rename from onnxruntime/core/codegen/target/tvm_schedule_builder.cc rename to onnxruntime/core/codegen/passes/scheduler/tvm_schedule_builder.cc index be0d6ab44530f..6f0ffa14e8abb 100644 --- a/onnxruntime/core/codegen/target/tvm_schedule_builder.cc +++ b/onnxruntime/core/codegen/passes/scheduler/tvm_schedule_builder.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/tvm_schedule_builder.h" +#include "core/codegen/passes/scheduler/tvm_schedule_builder.h" #include "core/codegen/common/op_macro.h" #include "core/codegen/common/settings.h" diff --git a/onnxruntime/core/codegen/target/tvm_schedule_builder.h b/onnxruntime/core/codegen/passes/scheduler/tvm_schedule_builder.h similarity index 95% rename from onnxruntime/core/codegen/target/tvm_schedule_builder.h rename to onnxruntime/core/codegen/passes/scheduler/tvm_schedule_builder.h index e125b1c5db00a..9f0a1b3ef45c2 100644 --- a/onnxruntime/core/codegen/target/tvm_schedule_builder.h +++ b/onnxruntime/core/codegen/passes/scheduler/tvm_schedule_builder.h @@ -2,7 +2,7 @@ // Licensed under the MIT License. #pragma once -#include "core/codegen/target/tvm_scheduler.h" +#include "core/codegen/passes/scheduler/tvm_scheduler.h" #include "core/common/common.h" namespace onnxruntime { diff --git a/onnxruntime/core/codegen/target/tvm_scheduler.cc b/onnxruntime/core/codegen/passes/scheduler/tvm_scheduler.cc similarity index 95% rename from onnxruntime/core/codegen/target/tvm_scheduler.cc rename to onnxruntime/core/codegen/passes/scheduler/tvm_scheduler.cc index 697a334f469de..071200a234e33 100644 --- a/onnxruntime/core/codegen/target/tvm_scheduler.cc +++ b/onnxruntime/core/codegen/passes/scheduler/tvm_scheduler.cc @@ -1,11 +1,11 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/tvm_scheduler.h" +#include "core/codegen/passes/scheduler/tvm_scheduler.h" #include "core/codegen/common/common.h" #include "core/codegen/common/dispatcher.h" -#include "core/codegen/target/codegen_context.h" +#include "core/codegen/passes/utils/codegen_context.h" namespace onnxruntime { namespace codegen { diff --git a/onnxruntime/core/codegen/target/tvm_scheduler.h b/onnxruntime/core/codegen/passes/scheduler/tvm_scheduler.h similarity index 98% rename from onnxruntime/core/codegen/target/tvm_scheduler.h rename to onnxruntime/core/codegen/passes/scheduler/tvm_scheduler.h index 9c2dbfdd327c4..413e0fb504e89 100644 --- a/onnxruntime/core/codegen/target/tvm_scheduler.h +++ b/onnxruntime/core/codegen/passes/scheduler/tvm_scheduler.h @@ -5,7 +5,7 @@ #include "core/common/common.h" #include "core/codegen/common/creator.h" #include "core/codegen/common/registry.h" -#include "core/codegen/target/codegen_context.h" +#include "core/codegen/passes/utils/codegen_context.h" #include "core/graph/graph.h" #include diff --git a/onnxruntime/core/codegen/target/codegen_context.cc b/onnxruntime/core/codegen/passes/utils/codegen_context.cc similarity index 92% rename from onnxruntime/core/codegen/target/codegen_context.cc rename to onnxruntime/core/codegen/passes/utils/codegen_context.cc index 17bef98c0591c..2f1a59b4a92eb 100644 --- a/onnxruntime/core/codegen/target/codegen_context.cc +++ b/onnxruntime/core/codegen/passes/utils/codegen_context.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/codegen_context.h" +#include "core/codegen/passes/utils/codegen_context.h" #include "core/codegen/common/common.h" diff --git a/onnxruntime/core/codegen/target/codegen_context.h b/onnxruntime/core/codegen/passes/utils/codegen_context.h similarity index 100% rename from onnxruntime/core/codegen/target/codegen_context.h rename to onnxruntime/core/codegen/passes/utils/codegen_context.h diff --git a/onnxruntime/core/codegen/target/ort_tvm_utils.cc b/onnxruntime/core/codegen/passes/utils/ort_tvm_utils.cc similarity index 98% rename from onnxruntime/core/codegen/target/ort_tvm_utils.cc rename to onnxruntime/core/codegen/passes/utils/ort_tvm_utils.cc index 88aa987a371c2..f7906b71e1189 100644 --- a/onnxruntime/core/codegen/target/ort_tvm_utils.cc +++ b/onnxruntime/core/codegen/passes/utils/ort_tvm_utils.cc @@ -1,10 +1,10 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/ort_tvm_utils.h" +#include "core/codegen/passes/utils/ort_tvm_utils.h" #include "core/codegen/common/profile.h" -#include "core/codegen/target/codegen_context.h" +#include "core/codegen/passes/utils/codegen_context.h" #include "core/providers/common.h" #include "gsl/gsl_util" diff --git a/onnxruntime/core/codegen/target/ort_tvm_utils.h b/onnxruntime/core/codegen/passes/utils/ort_tvm_utils.h similarity index 100% rename from onnxruntime/core/codegen/target/ort_tvm_utils.h rename to onnxruntime/core/codegen/passes/utils/ort_tvm_utils.h diff --git a/onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.cc b/onnxruntime/core/codegen/passes/weight_layout/transpose_2d.cc similarity index 93% rename from onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.cc rename to onnxruntime/core/codegen/passes/weight_layout/transpose_2d.cc index 8508fa3e40aa8..df1767b81032a 100644 --- a/onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.cc +++ b/onnxruntime/core/codegen/passes/weight_layout/transpose_2d.cc @@ -1,9 +1,9 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/weight_layout/transpose_2d.h" +#include "core/codegen/passes/weight_layout/transpose_2d.h" -#include "core/codegen/target/codegen_context.h" +#include "core/codegen/passes/utils/codegen_context.h" namespace onnxruntime { namespace tvm_codegen { diff --git a/onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.h b/onnxruntime/core/codegen/passes/weight_layout/transpose_2d.h similarity index 93% rename from onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.h rename to onnxruntime/core/codegen/passes/weight_layout/transpose_2d.h index bfa192d1f432c..65babaaec8dac 100644 --- a/onnxruntime/core/codegen/target/generic/weight_layout/transpose_2d.h +++ b/onnxruntime/core/codegen/passes/weight_layout/transpose_2d.h @@ -3,7 +3,7 @@ #pragma once -#include "core/codegen/target/weight_layout.h" +#include "core/codegen/passes/weight_layout/weight_layout.h" #include namespace onnxruntime { diff --git a/onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.cc b/onnxruntime/core/codegen/passes/weight_layout/vertical_stripes_2d.cc similarity index 94% rename from onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.cc rename to onnxruntime/core/codegen/passes/weight_layout/vertical_stripes_2d.cc index 8792f753652aa..b1ddb791a3b3d 100644 --- a/onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.cc +++ b/onnxruntime/core/codegen/passes/weight_layout/vertical_stripes_2d.cc @@ -1,9 +1,9 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/generic/weight_layout/vertical_stripes_2d.h" +#include "core/codegen/passes/weight_layout/vertical_stripes_2d.h" -#include "core/codegen/target/codegen_context.h" +#include "core/codegen/passes/utils/codegen_context.h" namespace onnxruntime { namespace tvm_codegen { diff --git a/onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.h b/onnxruntime/core/codegen/passes/weight_layout/vertical_stripes_2d.h similarity index 95% rename from onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.h rename to onnxruntime/core/codegen/passes/weight_layout/vertical_stripes_2d.h index d9b6f3299aa52..b9b65025dc014 100644 --- a/onnxruntime/core/codegen/target/generic/weight_layout/vertical_stripes_2d.h +++ b/onnxruntime/core/codegen/passes/weight_layout/vertical_stripes_2d.h @@ -3,7 +3,7 @@ #pragma once #include "core/codegen/common/common.h" -#include "core/codegen/target/weight_layout.h" +#include "core/codegen/passes/weight_layout/weight_layout.h" #include namespace onnxruntime { diff --git a/onnxruntime/core/codegen/target/weight_layout.cc b/onnxruntime/core/codegen/passes/weight_layout/weight_layout.cc similarity index 96% rename from onnxruntime/core/codegen/target/weight_layout.cc rename to onnxruntime/core/codegen/passes/weight_layout/weight_layout.cc index f74c6ad0727e4..0b8ae71030779 100644 --- a/onnxruntime/core/codegen/target/weight_layout.cc +++ b/onnxruntime/core/codegen/passes/weight_layout/weight_layout.cc @@ -1,12 +1,12 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/codegen/target/weight_layout.h" +#include "core/codegen/passes/weight_layout/weight_layout.h" #include "core/codegen/common/common.h" #include "core/codegen/common/utils.h" #include "core/codegen/mti/mti_tvm_utils.h" -#include "core/codegen/target/ort_tvm_utils.h" +#include "core/codegen/passes/utils/ort_tvm_utils.h" namespace onnxruntime { namespace tvm_codegen { diff --git a/onnxruntime/core/codegen/target/weight_layout.h b/onnxruntime/core/codegen/passes/weight_layout/weight_layout.h similarity index 100% rename from onnxruntime/core/codegen/target/weight_layout.h rename to onnxruntime/core/codegen/passes/weight_layout/weight_layout.h diff --git a/onnxruntime/test/tvm/tvm_demo/demo_compiler.cc b/onnxruntime/test/tvm/tvm_demo/demo_compiler.cc index a7c2a47ef405d..c978675b06d8e 100644 --- a/onnxruntime/test/tvm/tvm_demo/demo_compiler.cc +++ b/onnxruntime/test/tvm/tvm_demo/demo_compiler.cc @@ -3,11 +3,11 @@ #include "test/tvm/tvm_demo/demo_compiler.h" -#include "core/codegen/target/generic/scheduler/schedule_utils.h" -#include "core/codegen/target/ort_tvm_utils.h" -#include "core/codegen/target/tvm_ir_builder.h" -#include "core/codegen/target/tvm_scheduler.h" -#include "core/codegen/target/tvm_schedule_builder.h" +#include "core/codegen/passes/scheduler/schedule_utils.h" +#include "core/codegen/passes/utils/ort_tvm_utils.h" +#include "core/codegen/passes/op_ir_creator/tvm_ir_builder.h" +#include "core/codegen/passes/scheduler/tvm_scheduler.h" +#include "core/codegen/passes/scheduler/tvm_schedule_builder.h" #include #include