Triton Codegen for ORTModule (#15831)

Fuse connected elementwise and reduce Ops to TritonOp and codegen triton
code to run the kernel.

This PR is co-edited by @wejoncy and @er3x3

cmake/CMakeLists.txt

@@ -217,6 +217,9 @@ option(onnxruntime_ENABLE_CPUINFO "Enable cpuinfo" ON)
 # ATen fallback support
 option(onnxruntime_ENABLE_ATEN "Enable ATen fallback" OFF)
 
+# Triton support
+option(onnxruntime_ENABLE_TRITON "Enable Triton" OFF)
+
 # composable kernel is managed automatically, unless user want to explicitly disable it, it should not be manually set
 option(onnxruntime_USE_COMPOSABLE_KERNEL "Enable composable kernel for ROCm EP" ON)
 option(onnxruntime_USE_ROCBLAS_EXTENSION_API "Enable rocblas tuning for ROCm EP" OFF)
@@ -247,6 +250,7 @@ if (onnxruntime_ENABLE_TRAINING)
   set(onnxruntime_ENABLE_TRAINING_APIS ON)
   set(onnxruntime_ENABLE_TRAINING_TORCH_INTEROP ON)
   set(onnxruntime_ENABLE_ATEN ON)
+  set(onnxruntime_ENABLE_TRITON ON)
 endif()
 
 if (onnxruntime_ENABLE_TRAINING_APIS)
@@ -409,6 +413,13 @@ if (NOT (UNIX AND onnxruntime_ENABLE_PYTHON AND onnxruntime_ENABLE_TRAINING AND
   set(onnxruntime_ENABLE_TRAINING_TORCH_INTEROP OFF)
 endif()
 
+if (NOT (UNIX AND onnxruntime_USE_CUDA AND onnxruntime_ENABLE_PYTHON AND onnxruntime_ENABLE_TRAINING AND (NOT onnxruntime_BUILD_SHARED_LIB)))
+  if (onnxruntime_ENABLE_TRITON)
+    message(WARNING "onnxruntime_ENABLE_TRITON is turned OFF because it's designed to support CUDA training on Linux only currently.")
+  endif()
+  set(onnxruntime_ENABLE_TRITON OFF)
+endif()
+
 set(onnxruntime_REQUIRE_PYTHON_EMBED_LIB OFF)
 if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
   add_compile_definitions(ENABLE_TRAINING_TORCH_INTEROP)
@@ -419,6 +430,17 @@ if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
   endif()
 endif()
 
+if (onnxruntime_ENABLE_TRITON)
+  # Need SetOutputMLValue.
+  set(onnxruntime_ENABLE_ATEN ON)
+  add_compile_definitions(ENABLE_TRITON)
+
+  # Python::Python is required for building unit test executables.
+  if (onnxruntime_BUILD_UNIT_TESTS)
+    set(onnxruntime_REQUIRE_PYTHON_EMBED_LIB ON)
+  endif()
+endif()
+
 # General C# properties
 if (onnxruntime_BUILD_CSHARP)
   check_language(CSharp)

cmake/onnxruntime_framework.cmake

@@ -8,12 +8,31 @@ file(GLOB_RECURSE onnxruntime_framework_srcs CONFIGURE_DEPENDS
 )
 
 if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
-file(GLOB_RECURSE onnxruntime_training_framework_torch_srcs CONFIGURE_DEPENDS
-    "${ORTTRAINING_SOURCE_DIR}/core/framework/torch/*.h"
-    "${ORTTRAINING_SOURCE_DIR}/core/framework/torch/*.cc"
-)
-
+  file(GLOB_RECURSE onnxruntime_training_framework_torch_srcs CONFIGURE_DEPENDS
+      "${ORTTRAINING_SOURCE_DIR}/core/framework/torch/*.h"
+      "${ORTTRAINING_SOURCE_DIR}/core/framework/torch/*.cc"
+  )
   list(APPEND onnxruntime_framework_srcs ${onnxruntime_training_framework_torch_srcs})
+  if (onnxruntime_ENABLE_TRITON)
+    file(GLOB_RECURSE onnxruntime_training_framework_triton_srcs CONFIGURE_DEPENDS
+      "${ORTTRAINING_SOURCE_DIR}/core/framework/triton/*.h"
+      "${ORTTRAINING_SOURCE_DIR}/core/framework/triton/*.cc"
+    )
+    list(APPEND onnxruntime_framework_srcs ${onnxruntime_training_framework_triton_srcs})
+  endif()
+elseif(onnxruntime_ENABLE_TRITON)
+  # Triton executor shares some code from torch_interop, such as python and dlpack related code files.
+  # When torch_interop is enabled, all these dependencies are already included.
+  # But if not, we need to include them explicitly.
+  file(GLOB_RECURSE onnxruntime_training_framework_triton_srcs CONFIGURE_DEPENDS
+    "${ORTTRAINING_SOURCE_DIR}/core/framework/torch/dlpack_python.h"
+    "${ORTTRAINING_SOURCE_DIR}/core/framework/torch/dlpack_python.cc"
+    "${ORTTRAINING_SOURCE_DIR}/core/framework/torch/gil.h"
+    "${ORTTRAINING_SOURCE_DIR}/core/framework/torch/python_common.h"
+    "${ORTTRAINING_SOURCE_DIR}/core/framework/triton/*.h"
+    "${ORTTRAINING_SOURCE_DIR}/core/framework/triton/*.cc"
+  )
+  list(APPEND onnxruntime_framework_srcs ${onnxruntime_training_framework_triton_srcs})
 endif()
 
 if (onnxruntime_MINIMAL_BUILD)
@@ -50,7 +69,7 @@ endif()
 # Needed for the provider interface, as it includes training headers when training is enabled
 if (onnxruntime_ENABLE_TRAINING_OPS)
   target_include_directories(onnxruntime_framework PRIVATE ${ORTTRAINING_ROOT})
-  if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
+  if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP OR onnxruntime_ENABLE_TRITON)
     onnxruntime_add_include_to_target(onnxruntime_framework Python::Module)
     target_include_directories(onnxruntime_framework PRIVATE ${dlpack_SOURCE_DIR}/include)
   endif()

cmake/onnxruntime_optimizer.cmake

@@ -110,6 +110,10 @@ target_include_directories(onnxruntime_optimizer PRIVATE ${ONNXRUNTIME_ROOT})
 if (onnxruntime_ENABLE_TRAINING)
   target_include_directories(onnxruntime_optimizer PRIVATE ${ORTTRAINING_ROOT})
 endif()
+if (onnxruntime_ENABLE_TRITON)
+  target_link_libraries(onnxruntime_optimizer PRIVATE nlohmann_json::nlohmann_json)
+  onnxruntime_add_include_to_target(onnxruntime_optimizer Python::Module)
+endif()
 add_dependencies(onnxruntime_optimizer ${onnxruntime_EXTERNAL_DEPENDENCIES})
 set_target_properties(onnxruntime_optimizer PROPERTIES FOLDER "ONNXRuntime")
 

cmake/onnxruntime_providers.cmake

@@ -204,6 +204,8 @@ if (onnxruntime_ENABLE_TRAINING_OPS AND NOT onnxruntime_ENABLE_TRAINING)
     "${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/tensorboard/*.h"
     "${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/torch/*.cc"
     "${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/torch/*.h"
+    "${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/triton/triton_op.cc"
+    "${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/triton/triton_op.h"
   )
 
   list(REMOVE_ITEM onnxruntime_providers_src ${onnxruntime_cpu_full_training_only_srcs})
@@ -233,6 +235,8 @@ if (onnxruntime_ENABLE_TRAINING)
   file(GLOB_RECURSE onnxruntime_training_framework_excude_srcs CONFIGURE_DEPENDS
       "${ORTTRAINING_SOURCE_DIR}/core/framework/torch/*.h"
       "${ORTTRAINING_SOURCE_DIR}/core/framework/torch/*.cc"
+      "${ORTTRAINING_SOURCE_DIR}/core/framework/triton/*.h"
+      "${ORTTRAINING_SOURCE_DIR}/core/framework/triton/*.cc"
   )
 
   list(REMOVE_ITEM onnxruntime_cpu_training_ops_srcs ${onnxruntime_training_framework_excude_srcs})
@@ -303,7 +307,7 @@ endif()
 if (onnxruntime_ENABLE_TRAINING)
   add_dependencies(onnxruntime_providers tensorboard)
   onnxruntime_add_include_to_target(onnxruntime_providers tensorboard)
-  if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
+  if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP OR onnxruntime_ENABLE_TRITON)
     onnxruntime_add_include_to_target(onnxruntime_providers Python::Module)
   endif()
 
@@ -427,12 +431,12 @@ if (onnxruntime_USE_CUDA)
         "${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/controlflow/wait.cc"
         "${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/controlflow/wait.h"
         "${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/controlflow/yield.cc"
-        "${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/controlflow/yield.h"
         "${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/gist/*.cc"
         "${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/gist/*.h"
         "${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/gist/*.cu"
         "${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/torch/*.cc"
         "${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/torch/*.h"
+        "${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/triton/triton_op.cc"
       )
 
       list(REMOVE_ITEM onnxruntime_providers_cuda_src ${onnxruntime_cuda_full_training_only_srcs})
@@ -496,7 +500,7 @@ if (onnxruntime_USE_CUDA)
       if (onnxruntime_ENABLE_TRAINING)
         target_link_libraries(${target} PRIVATE onnxruntime_training)
       endif()
-      if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
+      if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP OR onnxruntime_ENABLE_TRITON)
         onnxruntime_add_include_to_target(${target} Python::Module)
       endif()
     endif()

cmake/onnxruntime_python.cmake

@@ -387,8 +387,14 @@ if (onnxruntime_ENABLE_TRAINING)
   file(GLOB onnxruntime_python_ortmodule_torch_cpp_ext_fused_ops_srcs CONFIGURE_DEPENDS
     "${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/*"
   )
+  file(GLOB onnxruntime_python_ort_triton_srcs CONFIGURE_DEPENDS
+    "${ORTTRAINING_SOURCE_DIR}/python/training/ort_triton/*.py"
+  )
+  file(GLOB onnxruntime_python_ort_triton_kernel_srcs CONFIGURE_DEPENDS
+    "${ORTTRAINING_SOURCE_DIR}/python/training/ort_triton/kernel/*.py"
+  )
   file(GLOB onnxruntime_python_utils_data_srcs CONFIGURE_DEPENDS
-  "${ORTTRAINING_SOURCE_DIR}/python/training/utils/data/*"
+    "${ORTTRAINING_SOURCE_DIR}/python/training/utils/data/*"
   )
   file(GLOB onnxruntime_python_utils_hooks_srcs CONFIGURE_DEPENDS
   "${ORTTRAINING_SOURCE_DIR}/python/training/utils/hooks/*"
@@ -725,6 +731,8 @@ if (onnxruntime_ENABLE_TRAINING)
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/fused_ops
+    COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ort_triton
+    COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ort_triton/kernel
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/utils/data/
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/utils/hooks/
     COMMAND ${CMAKE_COMMAND} -E copy
@@ -775,6 +783,12 @@ if (onnxruntime_ENABLE_TRAINING)
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_ortmodule_torch_cpp_ext_fused_ops_srcs}
         $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/
+    COMMAND ${CMAKE_COMMAND} -E copy
+        ${onnxruntime_python_ort_triton_srcs}
+        $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ort_triton/
+    COMMAND ${CMAKE_COMMAND} -E copy
+        ${onnxruntime_python_ort_triton_kernel_srcs}
+        $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ort_triton/kernel/
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_utils_data_srcs}
         $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/utils/data/
@@ -801,7 +815,6 @@ if (onnxruntime_ENABLE_TRAINING)
       COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_api_srcs}
         $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/api/
-
     )
   endif()
 endif()

cmake/onnxruntime_training.cmake

@@ -22,6 +22,8 @@ file(GLOB_RECURSE onnxruntime_training_srcs
 file(GLOB_RECURSE onnxruntime_training_framework_excluded_srcs CONFIGURE_DEPENDS
     "${ORTTRAINING_SOURCE_DIR}/core/framework/torch/*.h"
     "${ORTTRAINING_SOURCE_DIR}/core/framework/torch/*.cc"
+    "${ORTTRAINING_SOURCE_DIR}/core/framework/triton/*.h"
+    "${ORTTRAINING_SOURCE_DIR}/core/framework/triton/*.cc"
 )
 
 list(REMOVE_ITEM onnxruntime_training_srcs ${onnxruntime_training_framework_excluded_srcs})

docs/ORTModule_Training_Guidelines.md

@@ -268,7 +268,60 @@ Check [FP16_Optimizer implementation](../orttraining/orttraining/python/training
 
 ```
 
-## 6. One More Thing - `LoadBalancingDistributedBatchSampler`
+
+## 6. Use OpenAI Triton to Compute ONNX Sub-graph
+
+`ORTModule` provides a way to switch to OpenAI Triton for executing some Ops to further accelerate training.
+
+### 6.1 Environment Variables
+
+#### ORTMODULE_USE_TRITON
+
+- **Feature Area**: *ORTMODULE/TritonOp*
+- **Description**: By default, this is disabled. This env var can be used for enabling Triton optimization.
+
+    ```bash
+    export ORTMODULE_USE_TRITON=1
+    ```
+
+#### ORTMODULE_ENABLE_TUNING
+
+- **Feature Area**: *ORTMODULE/TritonOp*
+- **Description**: By default, this is disabled. This env var can be used for enabling online Op tuning for those Ops that have multiple implementations on target EP.
+
+    ```bash
+    export ORTMODULE_ENABLE_TUNING=1
+    ```
+
+#### ORTMODULE_MAX_TUNING_DURATION_MS
+
+- **Feature Area**: *ORTMODULE/TritonOp*
+- **Description**: When `ORTMODULE_ENABLE_TUNING` is enabled, this env var can be used to set max tuning duration in ms to avoid long tuning time.
+
+    ```bash
+    export ORTMODULE_MAX_TUNING_DURATION_MS=9999
+    ```
+
+#### ORTMODULE_TUNING_RESULTS_PATH
+
+- **Feature Area**: *ORTMODULE/TritonOp*
+- **Description**: When `ORTMODULE_ENABLE_TUNING` is enabled, this env var can be used to specify where the online Op tuning results be saved for later use. By default the results will not be saved. When `ORTMODULE_ENABLE_TUNING` is NOT enabled, this env var can be used to specify where Op tuning results can be fetched as offline tuning results.
+
+    ```bash
+    export ORTMODULE_TUNING_RESULTS_PATH=/tmp/tuning_results
+    ```
+
+#### ORTMODULE_TRITON_DEBUG
+
+- **Feature Area**: *ORTMODULE/TritonOp*
+- **Description**: By default, this is disabled. This env var can be used for enabling Triton debug mode. All original and processed sub-graphs and corresponding generated Triton codes will be saved into a triton_debug folder under working directory.
+
+    ```bash
+    export ORTMODULE_TRITON_DEBUG=1
+    ```
+
+
+## 7. One More Thing - `LoadBalancingDistributedBatchSampler`
 
 `LoadBalancingDistributedBatchSampler` balances the data load across workers based on the sample's complexity.
 This is useful in scenarios like speech and NLP, where each batch has variable length and distributed training suffers from **straggler problem**. In such scenarios, the complexity function could be defined to return the length of the input sample sequence. The usage is similar to `torch.utils.data.DistributedSampler`, where each process loads a subset of the original dataset that is exclusive to it.

onnxruntime/core/graph/graph_utils.cc

@@ -214,6 +214,10 @@ bool MatchesOpSinceVersion(const Node& node, std::initializer_list<ONNX_NAMESPAC
   return std::find(versions.begin(), versions.end(), node.SinceVersion()) != versions.end();
 }
 
+bool MatchesOpSinceVersion(const Node& node, gsl::span<const ONNX_NAMESPACE::OperatorSetVersion> versions) {
+  return std::find(versions.begin(), versions.end(), node.SinceVersion()) != versions.end();
+}
+
 bool MatchesOpSetDomain(const Node& node, std::string_view domain) {
   const auto& node_domain = node.Domain();
   return node_domain == domain;
@@ -223,6 +227,13 @@ bool IsSupportedOptypeVersionAndDomain(const Node& node,
                                        std::string_view op_type,
                                        std::initializer_list<ONNX_NAMESPACE::OperatorSetVersion> versions,
                                        std::string_view domain) {
+  std::vector<ONNX_NAMESPACE::OperatorSetVersion> versions_vec(versions);
+  return IsSupportedOptypeVersionAndDomain(node, op_type, versions_vec, domain);
+}
+
+bool IsSupportedOptypeVersionAndDomain(const Node& node, std::string_view op_type,
+                                       gsl::span<const ONNX_NAMESPACE::OperatorSetVersion> versions,
+                                       std::string_view domain) {
   return (node.OpType() == op_type &&
   // we don't have op schemas in the minimal build so there's no way to check the deprecated flag
 #if !defined(ORT_MINIMAL_BUILD)
@@ -412,10 +423,6 @@ void GraphEdge::RemoveGraphEdges(Graph& graph, const std::vector<GraphEdge>& edg
 
 #if !defined(ORT_MINIMAL_BUILD)
 
-bool MatchesOpSinceVersion(const Node& node, gsl::span<const ONNX_NAMESPACE::OperatorSetVersion> versions) {
-  return std::find(versions.begin(), versions.end(), node.SinceVersion()) != versions.end();
-}
-
 int GetNodeInputIndexFromInputName(const Node& node, const std::string& input_name) {
   return GetIndexFromName(node, input_name, true);
 }

onnxruntime/core/graph/graph_utils.h

@@ -25,6 +25,10 @@ bool IsSupportedOptypeVersionAndDomain(const Node& node,
                                        std::string_view op_type,
                                        std::initializer_list<ONNX_NAMESPACE::OperatorSetVersion> versions,
                                        std::string_view domain = kOnnxDomain);
+bool IsSupportedOptypeVersionAndDomain(const Node& node,
+                                       std::string_view op_type,
+                                       gsl::span<const ONNX_NAMESPACE::OperatorSetVersion> versions,
+                                       std::string_view domain = kOnnxDomain);
 
 #if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
 
@@ -104,7 +108,6 @@ struct GraphEdge {
 
 #endif  // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
 
-#if !defined(ORT_MINIMAL_BUILD)
 /** Checks if the node has the same operator since version as the given one. */
 bool MatchesOpSinceVersion(const Node& node, std::initializer_list<ONNX_NAMESPACE::OperatorSetVersion> versions);
 bool MatchesOpSinceVersion(const Node& node, gsl::span<const ONNX_NAMESPACE::OperatorSetVersion> versions);
@@ -112,6 +115,7 @@ bool MatchesOpSinceVersion(const Node& node, gsl::span<const ONNX_NAMESPACE::Ope
 /** Checks if the node has the same op set domain as the given one. */
 bool MatchesOpSetDomain(const Node& node, std::string_view domain);
 
+#if !defined(ORT_MINIMAL_BUILD)
 /** Returns true if the execution provider assigned to current node is present in the compatible providers list
     or if the compatible_providers list is empty. */
 bool IsSupportedProvider(const Node& node,

onnxruntime/core/optimizer/graph_transformer_utils.cc

@@ -77,6 +77,10 @@
 #include "orttraining/core/optimizer/sce_loss_grad_bias_fusion.h"
 #include "orttraining/core/optimizer/memory_optimizer.h"
 #endif
+#ifdef ENABLE_TRITON
+#include "orttraining/core/optimizer/triton_fusion.h"
+#include "orttraining/core/framework/triton/triton_op_executor.h"
+#endif  // ENABLE_TRITON
 
 #endif  // !defined(ORT_MINIMAL_BUILD)
 
@@ -299,6 +303,27 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
 
       transformers.emplace_back(std::make_unique<MatmulTransposeFusion>(cpu_cuda_dml_rocm_eps));
       transformers.emplace_back(std::make_unique<BiasGeluFusion>(cpu_cuda_dml_rocm_eps));
+
+      transformers.emplace_back(std::make_unique<SkipLayerNormFusion>(cpu_cuda_dml_rocm_eps));
+
+      transformers.emplace_back(std::make_unique<FastGeluFusion>(cpu_cuda_rocm_eps));
+      transformers.emplace_back(std::make_unique<QuickGeluFusion>(cpu_cuda_dml_rocm_eps));
+
+      // GeluApproximation has side effects which may change results. It needs to be manually enabled,
+      // or alternatively the model can be updated offline using a model conversion script
+      //   e.g. fusion_gelu_approximation function used by onnxruntime/python/tools/transformers/onnx_model_bert.py
+      if (enable_gelu_approximation) {
+        transformers.emplace_back(std::make_unique<GeluApproximation>(cpu_cuda_rocm_eps));
+      }
+
+#ifdef ENABLE_TRITON
+      if (training::framework::triton::TritonOpExecutor::Instance().IsInitialized()) {
+        transformers.emplace_back(
+            std::make_unique<TritonFusion>(training::framework::triton::TritonOpExecutor::Instance().GetConfigJson(),
+                                           InlinedHashSet<std::string_view>{onnxruntime::kCudaExecutionProvider}));
+      }
+#endif  // ENABLE_TRITON
+
       transformers.emplace_back(std::make_unique<BiasSoftmaxFusion>(cpu_cuda_rocm_eps));
       transformers.emplace_back(std::make_unique<BiasDropoutFusion>(cuda_rocm_eps));
 #ifdef ENABLE_TRAINING
@@ -307,21 +332,9 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
       transformers.emplace_back(std::make_unique<SceLossGradBiasFusion>(cpu_cuda_rocm_eps));
 #endif
 
-      transformers.emplace_back(std::make_unique<SkipLayerNormFusion>(cpu_cuda_dml_rocm_eps));
-
-      transformers.emplace_back(std::make_unique<FastGeluFusion>(cpu_cuda_rocm_eps));
-      transformers.emplace_back(std::make_unique<QuickGeluFusion>(cpu_cuda_dml_rocm_eps));
-
       transformers.emplace_back(std::make_unique<MatMulScaleFusion>(cpu_cuda_dml_rocm_eps));
       transformers.emplace_back(std::make_unique<MatMulActivationFusion>(dml_ep));
 
-      // GeluApproximation has side effects which may change results. It needs to be manually enabled,
-      // or alternatively the model can be updated offline using a model conversion script
-      //   e.g. fusion_gelu_approximation function used by onnxruntime/python/tools/transformers/onnx_model_bert.py
-      if (enable_gelu_approximation) {
-        transformers.emplace_back(std::make_unique<GeluApproximation>(cpu_cuda_rocm_eps));
-      }
-
 #ifdef MLAS_TARGET_AMD64_IX86
       if (avx2_precision_mode) {
         transformers.emplace_back(std::make_unique<Avx2WeightS8ToU8Transformer>(cpu_ep));

onnxruntime/core/providers/cpu/cpu_provider_shared.cc

@@ -57,6 +57,10 @@
 #include "orttraining/training_ops/cpu/controlflow/yield.h"
 #endif
 
+#ifdef ENABLE_TRITON
+#include "orttraining/training_ops/cpu/triton/triton_op.h"
+#endif
+
 #include "cpu_provider_shared.h"
 
 namespace onnxruntime {
@@ -300,6 +304,19 @@ struct ProviderHostCPUImpl : ProviderHostCPU {
     return contrib::ExecuteReduceSumATen(p_ctx, axes, keepdims);
   }
 #endif
+
+#ifdef ENABLE_TRITON
+  Status contrib__TritonOp__Compute(const contrib::TritonOp* p, OpKernelContext* context) override {
+    return p->TritonOp::Compute(context);
+  }
+  bool contrib__IsTritonOpExecutorInitialized() override { return contrib::IsTritonOpExecutorInitialized(); }
+  Status contrib__ExecuteTritonOpByFuncName(
+      OpKernelContext* p_ctx, const std::string& func_name, size_t input_count, size_t output_count,
+      const InlinedHashMap<std::string, std::pair<std::string, int>>& kwargs) override {
+    return contrib::ExecuteTritonOpByFuncName(p_ctx, func_name, input_count, output_count, kwargs);
+  }
+#endif
+
 #endif
 };
 #if defined(_MSC_VER) && !defined(__clang__)

onnxruntime/core/providers/cpu/cpu_provider_shared.h

@@ -211,6 +211,15 @@ struct ProviderHostCPU {
   virtual bool contrib__IsATenOperatorExecutorInitialized() = 0;
   virtual Status contrib__ExecuteReduceSumATen(OpKernelContext* p_ctx, const gsl::span<const int64_t>& axes, bool keepdims) = 0;
 #endif
+
+#ifdef ENABLE_TRITON
+  virtual Status contrib__TritonOp__Compute(const contrib::TritonOp* p, OpKernelContext* context) = 0;
+  virtual bool contrib__IsTritonOpExecutorInitialized() = 0;
+  virtual Status contrib__ExecuteTritonOpByFuncName(
+      OpKernelContext* p_ctx, const std::string& func_name, size_t input_count, size_t output_count,
+      const InlinedHashMap<std::string, std::pair<std::string, int>>& kwargs) = 0;
+#endif
+
 #endif
 };
 
@@ -288,6 +297,18 @@ inline Status ExecuteReduceSumATen(OpKernelContext* p_ctx, const gsl::span<const
 }
 }  // namespace contrib
 #endif  // ENABLE_TRAINING
+
+#ifdef ENABLE_TRITON
+namespace contrib {
+inline bool IsTritonOpExecutorInitialized() { return g_host_cpu.contrib__IsTritonOpExecutorInitialized(); }
+inline Status ExecuteTritonOpByFuncName(OpKernelContext* p_ctx, const std::string& func_name, size_t input_count,
+                                        size_t output_count,
+                                        const InlinedHashMap<std::string, std::pair<std::string, int>>& kwargs) {
+  return g_host_cpu.contrib__ExecuteTritonOpByFuncName(p_ctx, func_name, input_count, output_count, kwargs);
+}
+}  // namespace contrib
+#endif  // ENABLE_TRITON
+
 #endif  // USE_CUDA || USE_ROCM
 #endif
 

onnxruntime/core/providers/cuda/math/gemm.cc

@@ -2,9 +2,10 @@
 // Licensed under the MIT License.
 
 #include "core/providers/cuda/math/gemm.h"
+
 #include "core/providers/cpu/math/gemm_helper.h"
-#include "core/providers/cuda/cuda_common.h"
 #include "core/providers/cuda/shared_inc/fpgeneric.h"
+#include "core/providers/cuda/tunable/math/gemm.h"
 
 namespace onnxruntime {
 namespace cuda {
@@ -57,20 +58,34 @@ REGISTER_KERNEL_TYPED(BFloat16)
 
 template <typename T>
 Status Gemm<T>::ComputeInternal(OpKernelContext* ctx) const {
-  typedef typename ToCudaType<T>::MappedType CudaT;
-
   const auto* X = ctx->Input<Tensor>(0);
   const auto* W = ctx->Input<Tensor>(1);
   const auto* B = ctx->Input<Tensor>(2);
   // Bias could be missing. Treat as scalar 0 if that is the case.
   GemmHelper helper(X->Shape(), trans_A_, W->Shape(), trans_B_, B != nullptr ? B->Shape() : TensorShape({}));
-
-  if (!helper.State().IsOK())
-    return helper.State();
-
+  if (!helper.State().IsOK()) return helper.State();
   int M = gsl::narrow_cast<int>(helper.M());
   int N = gsl::narrow_cast<int>(helper.N());
   int K = gsl::narrow_cast<int>(helper.K());
+
+  auto* Y = ctx->Output(0, {M, N});
+  // Bail out early if the output is going to be empty
+  if (Y->Shape().Size() == 0) return Status::OK();
+
+  if (GetTuningContext()->IsTunableOpEnabled()) {
+    return tunable::TunableGemm<T>(M, N, K, trans_A_, trans_B_, alpha_, B ? beta_ : 0.0f, this, ctx);
+  }
+
+  return ComputeDefault(ctx, M, N, K);
+}
+
+template <typename T>
+Status Gemm<T>::ComputeDefault(OpKernelContext* ctx, int M, int N, int K) const {
+  typedef typename ToCudaType<T>::MappedType CudaT;
+
+  const auto* X = ctx->Input<Tensor>(0);
+  const auto* W = ctx->Input<Tensor>(1);
+  const auto* B = ctx->Input<Tensor>(2);
   auto* Y = ctx->Output(0, {M, N});
   CudaT* out_data = reinterpret_cast<CudaT*>(Y->MutableData<T>());
 

onnxruntime/core/providers/cuda/math/gemm.h

@@ -25,6 +25,7 @@ class Gemm final : public CudaKernel {
   }
 
   Status ComputeInternal(OpKernelContext* context) const override;
+  Status ComputeDefault(OpKernelContext* context, int M, int N, int K) const;
 
  private:
   bool trans_A_;

onnxruntime/core/providers/cuda/math/matmul.cc

@@ -2,9 +2,10 @@
 // Licensed under the MIT License.
 
 #include "core/providers/cuda/math/matmul.h"
-#include "core/providers/cpu/math/matmul_helper.h"
+
 #include "core/providers/cuda/shared_inc/fpgeneric.h"
 #include "core/providers/cuda/cuda_allocator.h"
+#include "core/providers/cuda/tunable/math/matmul.h"
 
 namespace onnxruntime {
 namespace cuda {
@@ -89,6 +90,37 @@ static bool CanUseStridedBatchedGemm(const TensorShape& left_shape, const Tensor
 
 template <typename T>
 Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
+  const Tensor* left_X = ctx->Input<Tensor>(0);
+  const Tensor* right_X = ctx->Input<Tensor>(1);
+
+  // Ignore the transpose flag if rank of input being 1.
+  // Be noted: numpy.transpose on vector does not change anything.
+  bool trans_a = trans_A_;
+  bool trans_b = trans_B_;
+  if (left_X->Shape().NumDimensions() == 1) {
+    trans_a = false;
+  }
+  if (right_X->Shape().NumDimensions() == 1) {
+    trans_b = false;
+  }
+
+  MatMulComputeHelper helper;
+  ORT_RETURN_IF_ERROR(
+      helper.Compute(left_X->Shape(), right_X->Shape(), trans_a, trans_b, trans_batch_a_, trans_batch_b_, false));
+
+  Tensor* Y = ctx->Output(0, helper.OutputShape());
+  // Bail out early if the output is going to be empty
+  if (Y->Shape().Size() == 0) return Status::OK();
+
+  if (GetTuningContext()->IsTunableOpEnabled()) {
+    return tunable::TunableMatMul<T>(alpha_, trans_a, trans_b, trans_batch_a_, trans_batch_b_, helper, this, ctx);
+  }
+
+  return ComputeDefault(ctx, helper);
+}
+
+template <typename T>
+Status MatMul<T>::ComputeDefault(OpKernelContext* ctx, MatMulComputeHelper& helper) const {
   typedef typename ToCudaType<T>::MappedType CudaT;
 
   const Tensor* left_X = ctx->Input<Tensor>(0);
@@ -105,15 +137,8 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
     transb = false;
   }
 
-  MatMulComputeHelper helper;
-  ORT_RETURN_IF_ERROR(helper.Compute(left_X->Shape(), right_X->Shape(), transa, transb, trans_batch_a_, trans_batch_b_, false));
-
   Tensor* Y = ctx->Output(0, helper.OutputShape());
 
-  // Bail out early if the output is going to be empty
-  if (Y->Shape().Size() == 0)
-    return Status::OK();
-
   const CudaT alpha = ToCudaType<T>::FromFloat(alpha_);
   const CudaT zero = ToCudaType<T>::FromFloat(0.0f);
 

onnxruntime/core/providers/cuda/math/matmul.h

@@ -4,6 +4,7 @@
 #pragma once
 
 #include "core/providers/cuda/cuda_kernel.h"
+#include "core/providers/cpu/math/matmul_helper.h"
 
 namespace onnxruntime {
 namespace cuda {
@@ -21,6 +22,7 @@ class MatMul final : public CudaKernel {
         trans_batch_b_{info.GetAttrOrDefault<int64_t>("transBatchB", 0) != 0} {}
 
   Status ComputeInternal(OpKernelContext* context) const override;
+  Status ComputeDefault(OpKernelContext* context, MatMulComputeHelper& helper) const;
 
  private:
   const float alpha_;

onnxruntime/core/providers/cuda/tunable/cuda_tuning_context.cc

@@ -27,6 +27,17 @@ static Status ValidateCudaVersion(const std::string& value) {
   return Status::OK();
 }
 
+std::string CudaTuningResultsValidator::GetOrtBuildConfig() const {
+  std::ostringstream oss;
+#ifdef ENABLE_TRITON
+  constexpr int kTriton = 1;
+#else
+  constexpr int kTriton = 0;
+#endif
+  oss << "ENABLE_TRITON=" << kTriton << "|";
+  return oss.str();
+}
+
 std::string CudaTuningResultsValidator::GetDeviceModel() const {
   return ep_->GetDeviceProp().name;
 }

onnxruntime/core/providers/cuda/tunable/cuda_tuning_context.h

@@ -20,6 +20,8 @@ class CudaTuningResultsValidator : public TuningResultsValidator {
   CudaTuningResultsValidator(CUDAExecutionProvider* ep);
 
  protected:
+  std::string GetOrtBuildConfig() const override;
+
   std::string GetDeviceModel() const;
   Status ValidateDeviceModel(const std::string& value) const;
 

onnxruntime/core/providers/cuda/tunable/math/gemm.cc

@@ -0,0 +1,105 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/cuda/tunable/math/gemm.h"
+
+#include "core/common/inlined_containers.h"
+
+namespace onnxruntime {
+namespace cuda {
+namespace tunable {
+
+template <typename T>
+GemmParams<T>::GemmParams(int m, int n, int k, bool trans_a, bool trans_b, float alpha, float beta,
+                          const Gemm<T>* gemm_kernel, OpKernelContext* ctx)
+    : OpParams(gemm_kernel->GetTuningContext(), gemm_kernel->Stream(ctx)),
+      trans_a_(trans_a),
+      trans_b_(trans_b),
+      alpha_(alpha),
+      beta_(beta),
+      m_(m),
+      n_(n),
+      k_(k),
+      gemm_kernel_(gemm_kernel),
+      ctx_(ctx) {
+  const auto* b = ctx->Input<Tensor>(2);
+  bm_ = gsl::narrow_cast<int>(beta_ == 0.0f ? 0 : (b->Shape().NumDimensions() > 1 ? b->Shape()[0] : 1));
+  bn_ = gsl::narrow_cast<int>(
+      beta_ == 0.0f
+          ? 0
+          : (b->Shape().NumDimensions() > 1 ? b->Shape()[1] : (b->Shape().NumDimensions() > 0 ? b->Shape()[0] : 1)));
+
+#ifdef ENABLE_TRITON
+  const auto* x = ctx->Input<Tensor>(0);
+  has_triton_support_ = contrib::IsTritonOpExecutorInitialized() &&
+                        (std::is_same<T, MLFloat16>::value || std::is_same<T, float>::value) &&
+                        x->Shape().NumDimensions() > 1;
+#endif
+}
+
+namespace {
+
+template <typename T>
+common::Status DefaultGemmOp(const GemmParams<T>* params) {
+  return params->gemm_kernel_->ComputeDefault(params->ctx_, params->m_, params->n_, params->k_);
+}
+
+#ifdef ENABLE_TRITON
+template <typename T>
+common::Status TritonGemmOp(const GemmParams<T>* params) {
+  TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(!params->has_triton_support_);
+  size_t input_count = params->beta_ == 0.0f ? 2 : 3;
+  size_t output_count = 1;
+  std::string func_name = params->beta_ == 0.0f ? "triton_matmul_out" : "triton_gemm_out";
+  InlinedHashMap<std::string, std::pair<std::string, int>> kwargs;
+  if (params->trans_a_) kwargs.insert({"trans_a", {"true", ONNX_NAMESPACE::TensorProto_DataType_BOOL}});
+  if (params->trans_b_) kwargs.insert({"trans_b", {"true", ONNX_NAMESPACE::TensorProto_DataType_BOOL}});
+  if (params->alpha_ != 1.0f) {
+    kwargs.insert({"alpha", {std::to_string(params->alpha_), ONNX_NAMESPACE::TensorProto_DataType_FLOAT}});
+  }
+  if (params->beta_ != 0.0f && params->beta_ != 1.0f) {
+    kwargs.insert({"beta", {std::to_string(params->beta_), ONNX_NAMESPACE::TensorProto_DataType_FLOAT}});
+  }
+  return contrib::ExecuteTritonOpByFuncName(params->ctx_, func_name, input_count, output_count, kwargs);
+}
+#endif
+
+template <typename T>
+class GemmTunableOp : public TunableOp<GemmParams<T>> {
+ public:
+  GemmTunableOp() {
+    this->RegisterOp(DefaultGemmOp<T>);
+#ifdef ENABLE_TRITON
+    this->RegisterOp(TritonGemmOp<T>);
+#endif
+  }
+};
+
+}  // namespace
+
+template <typename T>
+inline common::Status TunableGemm(int m, int n, int k, bool trans_a, bool trans_b, float alpha, float beta,
+                                  const Gemm<T>* gemm_kernel, OpKernelContext* ctx) {
+  GemmParams<T> params(m, n, k, trans_a, trans_b, alpha, beta, gemm_kernel, ctx);
+  if (params.tuning_ctx->IsTunableOpEnabled()) {
+    static GemmTunableOp<T> gemm{};
+    return gemm(&params);
+  }
+
+  return DefaultGemmOp(&params);
+}
+
+#define SPECIALIZE_TUNABLE_GEMM(T)                                                                                 \
+  template common::Status TunableGemm<T>(int m, int n, int k, bool trans_a, bool trans_b, float alpha, float beta, \
+                                         const Gemm<T>* gemm_kernel, OpKernelContext* ctx);
+
+SPECIALIZE_TUNABLE_GEMM(float)
+SPECIALIZE_TUNABLE_GEMM(double)
+SPECIALIZE_TUNABLE_GEMM(MLFloat16)
+SPECIALIZE_TUNABLE_GEMM(BFloat16)
+
+#undef SPECIALIZE_TUNABLE_GEMM
+
+}  // namespace tunable
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/core/providers/cuda/tunable/math/gemm.h

@@ -0,0 +1,46 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/common/status.h"
+#include "core/providers/cuda/cuda_common.h"
+#include "core/providers/cuda/tunable/cuda_tunable.h"
+#include "core/providers/cuda/math/gemm.h"
+
+namespace onnxruntime {
+namespace cuda {
+namespace tunable {
+
+template <typename T>
+struct GemmParams : OpParams {
+  GemmParams(int m, int n, int k, bool trans_a, bool trans_b, float alpha, float beta, const Gemm<T>* gemm_kernel,
+             OpKernelContext* ctx);
+
+  std::string Signature() const override {
+    return MakeString((trans_a_ ? "T" : "N"), (trans_b_ ? "T" : "N"), "_", m_, "_", n_, "_", k_, "_", bm_, "_", bn_);
+  }
+
+  bool trans_a_;
+  bool trans_b_;
+  float alpha_;
+  float beta_;
+  int m_;
+  int n_;
+  int k_;
+  int bm_;
+  int bn_;
+  const Gemm<T>* gemm_kernel_;
+  OpKernelContext* ctx_;
+#ifdef ENABLE_TRITON
+  bool has_triton_support_ = false;
+#endif
+};
+
+template <typename T>
+common::Status TunableGemm(int m, int n, int k, bool trans_a, bool trans_b, float alpha, float beta,
+                           const Gemm<T>* gemm_kernel, OpKernelContext* ctx);
+
+}  // namespace tunable
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/core/providers/cuda/tunable/math/matmul.cc

@@ -0,0 +1,114 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/cuda/tunable/math/matmul.h"
+
+#include "core/common/inlined_containers.h"
+
+namespace onnxruntime {
+namespace cuda {
+namespace tunable {
+
+template <typename T>
+MatMulParams<T>::MatMulParams(float alpha, bool trans_a, bool trans_b, bool trans_batch_a, bool trans_batch_b,
+                              MatMulComputeHelper& helper, const MatMul<T>* matmul_kernel, OpKernelContext* ctx)
+    : OpParams(matmul_kernel->GetTuningContext(), matmul_kernel->Stream(ctx)),
+      alpha_(alpha),
+      trans_a_(trans_a),
+      trans_b_(trans_b),
+      trans_batch_a_(trans_batch_a),
+      trans_batch_b_(trans_batch_b),
+      helper_(helper),
+      matmul_kernel_(matmul_kernel),
+      ctx_(ctx) {
+#ifdef ENABLE_TRITON
+  const TensorShape& shape_x = ctx->Input<Tensor>(0)->Shape();
+  const TensorShape& shape_y = ctx->Input<Tensor>(1)->Shape();
+  size_t rank_x = shape_x.NumDimensions();
+  size_t rank_y = shape_y.NumDimensions();
+  has_triton_support_ = contrib::IsTritonOpExecutorInitialized() &&
+                        (std::is_same<T, MLFloat16>::value || std::is_same<T, float>::value) && !trans_batch_a &&
+                        !trans_batch_b && rank_x > 1 && rank_y > 1;
+  if (has_triton_support_ && rank_x > 2 && rank_y > 2) {
+    has_triton_support_ = rank_x == rank_y;
+    if (has_triton_support_) {
+      for (size_t i = 0; i < rank_x - 2; ++i) {
+        if (shape_x[i] != shape_y[i]) {
+          has_triton_support_ = false;
+          break;
+        }
+      }
+    }
+  }
+#endif
+}
+
+template <typename T>
+std::string MatMulParams<T>::Signature() const {
+  const TensorShape& shape_x = ctx_->Input<Tensor>(0)->Shape();
+  const TensorShape& shape_y = ctx_->Input<Tensor>(1)->Shape();
+  return MakeString((trans_a_ ? "T" : "N"), (trans_b_ ? "T" : "N"), (trans_batch_a_ ? "T" : "N"),
+                    (trans_batch_a_ ? "T" : "N"), "_", shape_x.ToString(), "_", shape_y.ToString());
+}
+
+namespace {
+
+template <typename T>
+common::Status DefaultMatMulOp(const MatMulParams<T>* params) {
+  return params->matmul_kernel_->ComputeDefault(params->ctx_, params->helper_);
+}
+
+#ifdef ENABLE_TRITON
+template <typename T>
+common::Status TritonMatMulOp(const MatMulParams<T>* params) {
+  TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(!params->has_triton_support_);
+  InlinedHashMap<std::string, std::pair<std::string, int>> kwargs;
+  if (params->trans_a_) kwargs.insert({"trans_a", {"true", ONNX_NAMESPACE::TensorProto_DataType_BOOL}});
+  if (params->trans_b_) kwargs.insert({"trans_b", {"true", ONNX_NAMESPACE::TensorProto_DataType_BOOL}});
+  if (params->alpha_ != 1.0f) {
+    kwargs.insert({"alpha", {std::to_string(params->alpha_), ONNX_NAMESPACE::TensorProto_DataType_FLOAT}});
+  }
+  return contrib::ExecuteTritonOpByFuncName(params->ctx_, "triton_matmul_out", 2, 1, kwargs);
+}
+#endif
+
+template <typename T>
+class MatMulTunableOp : public TunableOp<MatMulParams<T>> {
+ public:
+  MatMulTunableOp() {
+    this->RegisterOp(DefaultMatMulOp<T>);
+#ifdef ENABLE_TRITON
+    this->RegisterOp(TritonMatMulOp<T>);
+#endif
+  }
+};
+
+}  // namespace
+
+template <typename T>
+inline common::Status TunableMatMul(float alpha, bool trans_a, bool trans_b, bool trans_batch_a, bool trans_batch_b,
+                                    MatMulComputeHelper& helper, const MatMul<T>* matmul_kernel, OpKernelContext* ctx) {
+  MatMulParams<T> params(alpha, trans_a, trans_b, trans_batch_a, trans_batch_b, helper, matmul_kernel, ctx);
+  if (params.tuning_ctx->IsTunableOpEnabled()) {
+    static MatMulTunableOp<T> matmul{};
+    return matmul(&params);
+  }
+
+  return DefaultMatMulOp(&params);
+}
+
+#define SPECIALIZE_TUNABLE_MATMUL(T)                                                                    \
+  template common::Status TunableMatMul<T>(float alpha, bool trans_a, bool trans_b, bool trans_batch_a, \
+                                           bool trans_batch_b, MatMulComputeHelper& helper,             \
+                                           const MatMul<T>* matmul_kernel, OpKernelContext* ctx);
+
+SPECIALIZE_TUNABLE_MATMUL(float)
+SPECIALIZE_TUNABLE_MATMUL(double)
+SPECIALIZE_TUNABLE_MATMUL(MLFloat16)
+SPECIALIZE_TUNABLE_MATMUL(BFloat16)
+
+#undef SPECIALIZE_TUNABLE_MATMUL
+
+}  // namespace tunable
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/core/providers/cuda/tunable/math/matmul.h

@@ -0,0 +1,41 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/common/status.h"
+#include "core/providers/cuda/cuda_common.h"
+#include "core/providers/cuda/tunable/cuda_tunable.h"
+#include "core/providers/cuda/math/matmul.h"
+
+namespace onnxruntime {
+namespace cuda {
+namespace tunable {
+
+template <typename T>
+struct MatMulParams : OpParams {
+  MatMulParams(float alpha, bool trans_a, bool trans_b, bool trans_batch_a, bool trans_batch_b,
+               MatMulComputeHelper& helper, const MatMul<T>* matmul_kernel, OpKernelContext* ctx);
+
+  std::string Signature() const override;
+
+  float alpha_;
+  bool trans_a_;
+  bool trans_b_;
+  bool trans_batch_a_;
+  bool trans_batch_b_;
+  MatMulComputeHelper& helper_;
+  const MatMul<T>* matmul_kernel_;
+  OpKernelContext* ctx_;
+#ifdef ENABLE_TRITON
+  bool has_triton_support_ = false;
+#endif
+};
+
+template <typename T>
+common::Status TunableMatMul(float alpha, bool trans_a, bool trans_b, bool trans_batch_a, bool trans_batch_b,
+                             MatMulComputeHelper& helper, const MatMul<T>* matmul_kernel, OpKernelContext* ctx);
+
+}  // namespace tunable
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/core/providers/shared_library/provider_api.h

@@ -179,6 +179,7 @@ class ATen;
 class Group;
 class PassThrough;
 class YieldOp;
+class TritonOp;
 class AdamWOptimizerBase;
 class SGDOptimizerV2Base;
 class ShrunkenGatherCommon;

onnxruntime/core/providers/shared_library/provider_bridge_provider.cc

@@ -59,6 +59,10 @@
 
 #endif
 
+#ifdef ENABLE_TRITON
+#include "orttraining/training_ops/cpu/triton/triton_op.h"
+#endif
+
 #ifndef _Ret_notnull_
 #define _Ret_notnull_
 #endif
@@ -708,6 +712,15 @@ void RefCountTracker::DumpDetails(const std::string& phase_name) const {
 #endif
 
 #endif
+
+#ifdef ENABLE_TRITON
+namespace contrib {
+Status TritonOp::Compute(OpKernelContext* context) const {
+  return g_host_cpu.contrib__TritonOp__Compute(this, context);
+}
+}  // namespace contrib
+#endif
+
 #endif
 
 #if defined(USE_CANN)

onnxruntime/python/tools/symbolic_shape_infer.py

@@ -2273,6 +2273,23 @@ class SymbolicShapeInference:
 
     def _infer_LayerNormalization(self, node):  # noqa: N802
         self._propagate_shape_and_type(node)
+        if len(node.output) > 1:
+            axis = get_attribute(node, "axis")
+            if axis is None:
+                axis = -1
+            x_shape = self._get_shape(node, 0)
+            if x_shape is not None:
+                rank = len(x_shape)
+                axis = handle_negative_axis(axis, rank)
+                mean_shape = x_shape[:axis] + [1 for _ in range(rank - axis)]
+                mean_dtype = self.known_vi_[node.input[0]].type.tensor_type.elem_type
+                if mean_dtype == onnx.TensorProto.FLOAT16 or mean_dtype == onnx.TensorProto.BFLOAT16:
+                    mean_dtype = onnx.TensorProto.FLOAT
+                vi = self.known_vi_[node.output[1]]
+                vi.CopyFrom(helper.make_tensor_value_info(node.output[1], mean_dtype, mean_shape))
+                if len(node.output) > 2:
+                    vi = self.known_vi_[node.output[2]]
+                    vi.CopyFrom(helper.make_tensor_value_info(node.output[2], mean_dtype, mean_shape))
 
     def _infer_LongformerAttention(self, node):  # noqa: N802
         self._propagate_shape_and_type(node)

orttraining/orttraining/core/framework/triton/triton_op_executor.cc

@@ -0,0 +1,132 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "orttraining/core/framework/triton/triton_op_executor.h"
+
+#include "orttraining/core/framework/torch/dlpack_python.h"
+#include "orttraining/core/framework/torch/gil.h"
+
+namespace onnxruntime {
+namespace training {
+namespace framework {
+namespace triton {
+
+namespace {
+void PythonObjectDeleter(PyObject* ptr) { Py_XDECREF(ptr); }
+}  // namespace
+
+void TritonOpExecutor::Initialize(PyObject* config_getter, PyObject* executor_by_name, PyObject* executor_by_onnx) {
+  ORT_ENFORCE(config_getter_.get() == nullptr && config_getter != nullptr && executor_by_name_.get() == nullptr &&
+              executor_by_name != nullptr && executor_by_onnx_.get() == nullptr && executor_by_onnx != nullptr);
+  Py_INCREF(config_getter);
+  Py_INCREF(executor_by_name);
+  Py_INCREF(executor_by_onnx);
+  PythonObjectPtr config_getter_ptr(config_getter, PythonObjectDeleter);
+  config_getter_ = std::move(config_getter_ptr);
+  PythonObjectPtr executor_by_name_ptr(executor_by_name, PythonObjectDeleter);
+  executor_by_name_ = std::move(executor_by_name_ptr);
+  PythonObjectPtr executor_by_onnx_ptr(executor_by_onnx, PythonObjectDeleter);
+  executor_by_onnx_ = std::move(executor_by_onnx_ptr);
+}
+
+std::string TritonOpExecutor::GetConfigJson() {
+  ORT_ENFORCE(IsInitialized());
+  // Python-related calls should happen only if guard is alive.
+  GilGuard guard;
+  PythonObjectPtr ret(PyObject_CallObject(config_getter_.get(), nullptr), PythonObjectDeleter);
+  char* buffer = nullptr;
+  Py_ssize_t length;
+  buffer = const_cast<char*>(PyUnicode_AsUTF8AndSize(ret.get(), &length));
+  return std::string(buffer, length);
+}
+
+void TritonOpExecutor::ExecuteByOnnx(int64_t onnx_key, const std::string& onnx_string,
+                                     const InlinedVector<const OrtValue*>& inputs, InlinedVector<OrtValue>& outputs,
+                                     const InlinedHashSet<size_t>& bool_outputs) {
+  ORT_ENFORCE(IsInitialized());
+  // Python-related calls should happen only if guard is alive.
+  GilGuard guard;
+  PythonObjectPtr args(PyTuple_New(static_cast<Py_ssize_t>(2 + inputs.size())), PythonObjectDeleter);
+  ORT_ENFORCE(args, "Failed to create args.");
+  PyTuple_SetItem(args.get(), 0, PyLong_FromLongLong(static_cast<long long>(onnx_key)));
+  PyTuple_SetItem(args.get(), 1, PyBytes_FromStringAndSize(onnx_string.c_str(), onnx_string.size()));
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    if (!inputs[i]) {
+      PyTuple_SetItem(args.get(), static_cast<Py_ssize_t>(2 + i), Py_None);
+      Py_INCREF(Py_None);
+    } else {
+      PyTuple_SetItem(args.get(), static_cast<Py_ssize_t>(2 + i), torch::ToDlpack(*inputs[i]));
+    }
+  }
+
+  PythonObjectPtr ret(PyObject_CallObject(executor_by_onnx_.get(), args.get()), PythonObjectDeleter);
+  if (ret == nullptr) {
+    PyErr_Print();
+    ORT_THROW("Python function execution fails with the above information.");
+  }
+  ORT_ENFORCE(ret.get() != Py_None);
+  if (PyTuple_Check(ret.get())) {
+    for (size_t i = 0; i < static_cast<size_t>(PyTuple_Size(ret.get())); ++i) {
+      outputs.emplace_back(torch::FromDlpack(PyTuple_GetItem(ret.get(), static_cast<Py_ssize_t>(i)),
+                                             bool_outputs.find(i) != bool_outputs.end()));
+    }
+  } else {
+    outputs.emplace_back(torch::FromDlpack(ret.get(), bool_outputs.find(0) != bool_outputs.end()));
+  }
+}
+
+void TritonOpExecutor::ExecuteByFuncName(const std::string& func_name, const InlinedVector<const OrtValue*>& inputs,
+                                         InlinedVector<OrtValue>& outputs, const InlinedHashSet<size_t>& bool_outputs,
+                                         const InlinedHashMap<std::string, std::pair<std::string, int>>& kwargs) {
+  ORT_ENFORCE(IsInitialized());
+  // Python-related calls should happen only if guard is alive.
+  GilGuard guard;
+  PythonObjectPtr args(PyTuple_New(static_cast<Py_ssize_t>(1 + inputs.size())), PythonObjectDeleter);
+  ORT_ENFORCE(args, "Failed to create args.");
+  PyTuple_SetItem(args.get(), 0, PyUnicode_FromString(func_name.c_str()));
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    if (!inputs[i]) {
+      PyTuple_SetItem(args.get(), static_cast<Py_ssize_t>(1 + i), Py_None);
+      Py_INCREF(Py_None);
+    } else {
+      PyTuple_SetItem(args.get(), static_cast<Py_ssize_t>(1 + i), torch::ToDlpack(*inputs[i]));
+    }
+  }
+
+  PythonObjectPtr python_kwargs(PyDict_New(), PythonObjectDeleter);
+  ORT_ENFORCE(python_kwargs, "Failed to create kwargs.");
+  for (const auto& kv : kwargs) {
+    if (kv.second.second == ONNX_NAMESPACE::TensorProto_DataType_BOOL) {
+      std::string bool_str = kv.second.first;
+      std::transform(bool_str.begin(), bool_str.end(), bool_str.begin(),
+                     [](unsigned char c) { return std::tolower(c); });
+      int bool_value = bool_str == "" || bool_str == "false" || bool_str == "0" ? 0 : 1;
+      PyDict_SetItemString(python_kwargs.get(), kv.first.c_str(), PyBool_FromLong(bool_value));
+    } else if (kv.second.second == ONNX_NAMESPACE::TensorProto_DataType_INT64) {
+      PyDict_SetItemString(python_kwargs.get(), kv.first.c_str(), PyLong_FromLongLong(std::stoll(kv.second.first)));
+    } else if (kv.second.second == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) {
+      PyDict_SetItemString(python_kwargs.get(), kv.first.c_str(), PyFloat_FromDouble(std::stod(kv.second.first)));
+    } else {
+      ORT_THROW("Unsupported kwargs data type: ", kv.second.second);
+    }
+  }
+  PythonObjectPtr ret(PyObject_Call(executor_by_name_.get(), args.get(), python_kwargs.get()), PythonObjectDeleter);
+  if (ret == nullptr) {
+    PyErr_Print();
+    ORT_THROW("Python function execution fails with the above information.");
+  }
+  if (ret.get() == Py_None) return;
+  if (PyTuple_Check(ret.get())) {
+    for (size_t i = 0; i < static_cast<size_t>(PyTuple_Size(ret.get())); ++i) {
+      outputs.emplace_back(torch::FromDlpack(PyTuple_GetItem(ret.get(), static_cast<Py_ssize_t>(i)),
+                                             bool_outputs.find(i) != bool_outputs.end()));
+    }
+  } else {
+    outputs.emplace_back(torch::FromDlpack(ret.get(), bool_outputs.find(0) != bool_outputs.end()));
+  }
+}
+
+}  // namespace triton
+}  // namespace framework
+}  // namespace training
+}  // namespace onnxruntime

orttraining/orttraining/core/framework/triton/triton_op_executor.h

@@ -0,0 +1,53 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/common/inlined_containers.h"
+#include "orttraining/core/framework/torch/python_common.h"
+
+#ifndef SHARED_PROVIDER
+#include "core/framework/ort_value.h"
+#endif
+
+namespace onnxruntime {
+namespace training {
+namespace framework {
+namespace triton {
+
+using PythonObjectPtr = std::unique_ptr<PyObject, std::function<void(PyObject*)>>;
+
+class TritonOpExecutor final {
+ public:
+  static TritonOpExecutor& Instance() {
+    static TritonOpExecutor instance;
+    return instance;
+  }
+
+  void Initialize(PyObject* config_getter, PyObject* executor_by_name, PyObject* executor_by_onnx);
+
+  bool IsInitialized() {
+    return config_getter_.get() != nullptr && executor_by_name_.get() != nullptr && executor_by_onnx_.get() != nullptr;
+  }
+
+  std::string GetConfigJson();
+
+  // Execute ONNX graph by codegening, compiling and executing Triton kernels.
+  void ExecuteByOnnx(int64_t onnx_key, const std::string& onnx_string, const InlinedVector<const OrtValue*>& inputs,
+                     InlinedVector<OrtValue>& outputs, const InlinedHashSet<size_t>& bool_outputs = {});
+
+  // Execute existing Triton kernel by Python function name.
+  void ExecuteByFuncName(const std::string& func_name, const InlinedVector<const OrtValue*>& inputs,
+                         InlinedVector<OrtValue>& outputs, const InlinedHashSet<size_t>& bool_outputs = {},
+                         const InlinedHashMap<std::string, std::pair<std::string, int>>& kwargs = {});
+
+ private:
+  PythonObjectPtr config_getter_;
+  PythonObjectPtr executor_by_name_;
+  PythonObjectPtr executor_by_onnx_;
+};
+
+}  // namespace triton
+}  // namespace framework
+}  // namespace training
+}  // namespace onnxruntime

orttraining/orttraining/core/graph/training_op_defs.cc

@@ -4094,6 +4094,33 @@ Return true if all elements are true and false otherwise.
         }
       });
 
+#ifdef ENABLE_TRITON
+  ONNX_CONTRIB_OPERATOR_SCHEMA(TritonOp)
+      .SetDomain(kMSDomain)
+      .SinceVersion(1)
+      .SetDoc(
+          "Calling an existing Python Triton kernel by function name, "
+          "or compute an ONNX graph through Python code to codegen, compile and execute Triton kernels.")
+      .Attr("func_name", "Function name of the Python Triton kernel.", AttributeProto::STRING, std::string(""))
+      .Attr("onnx_key", "The hash key for the ONNX graph.", AttributeProto::INT, static_cast<int64_t>(0))
+      .Attr("onnx_string", "The onnx string of the triton kernel.", AttributeProto::STRING, std::string(""))
+      .Input(0, "inputs",
+             "Input tensors. If to call an existing Python Triton kernel, "
+             "the input count and order should match the arguments of the function. If to compute an ONNX graph, "
+             "the input count and order should match the input count and order of the ONNX graph.",
+             "T", OpSchema::Variadic,
+             /*is_homogeneous*/ false,
+             /*min_arity*/ 0)
+      .Output(0, "outputs",
+              "Output tensors. If to compute an ONNX graph, "
+              "the output count and order should match the output count and order of the ONNX graph.",
+              "T", OpSchema::Variadic,
+              /*is_homogeneous*/ false,
+              /*min_arity*/ 1)
+      .TypeConstraint("T", OpSchema::all_tensor_types_with_bfloat(),
+                      "Allow inputs and outputs to be any kind of tensor.");
+#endif  // ENABLE_TRITON
+
   ONNX_CONTRIB_OPERATOR_SCHEMA(SoftmaxCrossEntropyLossInternal)
       .SetDomain(kMSDomain)
       .SinceVersion(1)

orttraining/orttraining/core/optimizer/triton_fusion.cc

@@ -0,0 +1,436 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#ifdef ENABLE_TRITON
+
+#include <nlohmann/json.hpp>
+#include <string_view>
+
+#include "orttraining/core/optimizer/triton_fusion.h"
+
+#include "core/framework/compute_capability.h"
+#include "core/graph/model.h"
+#include "core/graph/graph_utils.h"
+#include "core/optimizer/initializer.h"
+#include "core/optimizer/utils.h"
+#include "core/providers/partitioning_utils.h"
+
+using namespace ONNX_NAMESPACE;
+using json = nlohmann::json;
+
+namespace onnxruntime {
+
+namespace {
+
+using SizeTypeVec = InlinedVector<size_t>;
+using NodeVec = InlinedVector<Node*>;
+using NodeArgVec = InlinedVector<NodeArg*>;
+using ConstNodeArgVec = InlinedVector<const NodeArg*>;
+using NodeArgSet = InlinedHashSet<NodeArg*>;
+using IsSupportedFunc = std::function<bool(const Graph&, const Node&)>;
+
+int64_t Hash(const std::string& str) {
+  uint32_t hash = 0;
+  for (char const& c : str) {
+    hash = hash * 101 + c;
+  }
+
+  return static_cast<int64_t>(hash);
+}
+
+bool CheckAxis(const Node& node, int64_t expected_axis) {
+  const auto& attributes = node.GetAttributes();
+  if (attributes.find("axis") == attributes.end() || !attributes.at("axis").has_i()) return false;
+  int64_t axis = attributes.at("axis").i();
+  if (axis == expected_axis) return true;
+  if (axis < 0 || expected_axis < 0) {
+    const auto& input_shape = node.InputDefs()[0]->Shape();
+    if (input_shape == nullptr) return false;
+    int64_t rank = static_cast<int64_t>(input_shape->dim_size());
+    if (axis < 0) axis += rank;
+    int64_t non_neg_expected_axis = expected_axis < 0 ? expected_axis + rank : expected_axis;
+    return axis == non_neg_expected_axis;
+  }
+  return false;
+}
+
+bool CheckAxes(const Graph& graph, const Node& node, bool single_axis, const std::vector<int>& expected_axes) {
+  std::vector<int64_t> axes_values;
+  const auto& attributes = node.GetAttributes();
+  if (attributes.find("axes") != attributes.end()) {
+    axes_values = RetrieveValues<int64_t>(attributes.at("axes"));
+  } else if (node.InputDefs().size() == 2) {
+    auto axes_const = graph.GetConstantInitializer(node.InputDefs()[1]->Name(), true);
+    if (!axes_const) {
+      return false;
+    }
+    Initializer initializer{*axes_const, graph.ModelPath()};
+    axes_values.insert(axes_values.end(), initializer.DataAsSpan<int64_t>().begin(),
+                       initializer.DataAsSpan<int64_t>().end());
+  } else {
+    return false;
+  }
+
+  if (expected_axes.empty()) {
+    return !single_axis || axes_values.size() == 1;
+  }
+
+  std::vector<int64_t> expected_axes_values(expected_axes.begin(), expected_axes.end());
+  bool has_negative_axis = false;
+  for (auto axis : axes_values) {
+    if (axis < 0) {
+      has_negative_axis = true;
+      break;
+    }
+  }
+  if (!has_negative_axis) {
+    for (auto axis : expected_axes_values) {
+      if (axis < 0) {
+        has_negative_axis = true;
+        break;
+      }
+    }
+  }
+  if (has_negative_axis) {
+    const auto& input_shape = node.InputDefs()[0]->Shape();
+    if (input_shape == nullptr) return false;
+    int64_t rank = static_cast<int64_t>(input_shape->dim_size());
+    for (auto& axis : axes_values) {
+      if (axis < 0) axis += rank;
+    }
+    for (auto& axis : expected_axes_values) {
+      if (axis < 0) axis += rank;
+    }
+  }
+  std::sort(axes_values.begin(), axes_values.end());
+  std::sort(expected_axes_values.begin(), expected_axes_values.end());
+  return axes_values == expected_axes_values;
+}
+
+// A TritonOpPartition contains all connected nodes in ONNX graph which are supported by ORTModule's Triton module.
+// When building the TritonOpPartition, we keep all the dependencies and output ref counts so that to make sure
+// there is no dependency between two nodes in the partition through any node outside the partition.
+struct TritonOpPartition {
+  NodeVec nodes;
+  NodeArgSet outputs;
+  NodeArgSet dependencies;
+  size_t output_ref_count;
+
+  void MergeFrom(const TritonOpPartition& other) {
+    nodes.insert(nodes.end(), other.nodes.begin(), other.nodes.end());
+    outputs.insert(other.outputs.begin(), other.outputs.end());
+    dependencies.insert(other.dependencies.begin(), other.dependencies.end());
+    output_ref_count += other.output_ref_count;
+  }
+
+  bool IsValid(const TritonFusionConfig& config) const {
+    size_t count = 0;
+    bool all_ignore_min_nodes = true;
+    for (const auto& node : nodes) {
+      if (!config.IsNoOp(*node)) {
+        ++count;
+        if (count >= config.min_nodes) return true;
+      }
+      if (!config.IgnoreMinNodes(*node)) all_ignore_min_nodes = false;
+    }
+    return all_ignore_min_nodes;
+  }
+};
+
+}  // namespace
+
+void from_json(const json& j, TritonFusionConfig::OpInfo& op_info) {
+  if (j.contains("domain")) j.at("domain").get_to(op_info.domain);
+  j.at("versions").get_to(op_info.versions);
+  if (j.contains("is_no_op")) j.at("is_no_op").get_to(op_info.is_no_op);
+  if (j.contains("conditions")) j.at("conditions").get_to(op_info.conditions);
+  if (j.contains("ignore_min_nodes")) j.at("ignore_min_nodes").get_to(op_info.ignore_min_nodes);
+}
+
+TritonFusionConfig::TritonFusionConfig(std::string_view config_json) {
+  const auto& config = json::parse(config_json);
+  if (config.contains("ops")) {
+    ops = config.at("ops").get<std::unordered_map<std::string, OpInfo>>();
+  }
+  if (config.contains("initializer")) {
+    initializer = config.at("initializer").get<std::string>();
+  }
+  if (config.contains("min_nodes")) {
+    min_nodes = static_cast<size_t>(config.at("min_nodes").get<int>());
+  }
+}
+
+bool TritonFusionConfig::IsSupported(const Graph& graph, const Node& node) const {
+  const auto& op_type = node.OpType();
+  auto it = ops.find(op_type);
+  if (it == ops.end()) {
+    return false;
+  }
+
+  const auto& op_info = it->second;
+  if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, op_type, op_info.versions, op_info.domain)) {
+    return false;
+  }
+
+  for (const auto& pair : op_info.conditions) {
+    if (pair.first == "axis") {
+      if (!CheckAxis(node, static_cast<int64_t>(std::stoi(pair.second)))) {
+        return false;
+      }
+    } else if (pair.first == "axes") {
+      if (pair.second == "constant") {
+        if (!CheckAxes(graph, node, false, {})) {
+          return false;
+        }
+      } else if (pair.second == "single") {
+        if (!CheckAxes(graph, node, true, {})) {
+          return false;
+        }
+      } else {
+        const auto& axes = json::parse(pair.second);
+        std::vector<int> axes_values = axes.get<std::vector<int>>();
+        if (!CheckAxes(graph, node, false, axes_values)) {
+          return false;
+        }
+      }
+    } else {
+      return false;
+    }
+  }
+
+  return true;
+}
+
+bool TritonFusionConfig::IsNoOp(const Node& node) const {
+  const auto& op_type = node.OpType();
+  return ops.find(op_type) != ops.end() && ops.at(op_type).is_no_op;
+}
+
+bool TritonFusionConfig::IgnoreMinNodes(const Node& node) const {
+  const auto& op_type = node.OpType();
+  return ops.find(op_type) != ops.end() && ops.at(op_type).ignore_min_nodes;
+}
+
+const ONNX_NAMESPACE::TensorProto* TritonFusionConfig::TryGetInitializer(const Graph& graph, const Node& node,
+                                                                         NodeArg* node_arg) const {
+  if (initializer == "none" || !graph_utils::IsInitializer(graph, node_arg->Name(), true)) {
+    return nullptr;
+  }
+
+  const ONNX_NAMESPACE::TensorProto* tensor = nullptr;
+  if (!graph.GetInitializedTensor(node_arg->Name(), tensor) || !tensor) {
+    return nullptr;
+  }
+
+  if (initializer == "all") {
+    return tensor;
+  }
+
+  if ((node.OpType() == "ReduceSum" || node.OpType() == "ReduceMean" || node.OpType() == "ReduceMax" ||
+       node.OpType() == "ReduceMin") &&
+      node.InputDefs().size() >= 2 and node.InputDefs()[1]->Name() == node_arg->Name()) {
+    return tensor;
+  }
+
+  return optimizer_utils::IsScalar(*node_arg) ? tensor : nullptr;
+}
+
+Status TritonFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const {
+  GraphViewer graph_viewer(graph);
+  const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder();
+  size_t global_id = 0;
+  InlinedHashMap<size_t, TritonOpPartition> partitions;
+  InlinedHashMap<size_t, TritonOpPartition> partitions_to_fuse;
+  InlinedHashMap<NodeArg*, size_t> active_outputs;
+  for (auto node_index : node_topology_list) {
+    auto* p_node = graph.GetNode(node_index);
+    if (!p_node) continue;
+    auto& node = *p_node;
+    ORT_RETURN_IF_ERROR(Recurse(node, modified, graph_level, logger));
+    bool is_supported =
+        graph_utils::IsSupportedProvider(node, GetCompatibleExecutionProviders()) && config_.IsSupported(graph, node);
+    SizeTypeVec partitions_to_merge;
+    for (auto& pair : partitions) {
+      auto& partition = pair.second;
+      bool connect_to_output = false;
+      bool connect_to_dependency = false;
+      for (auto& input : node.MutableInputDefs()) {
+        if (partition.outputs.find(input) != partition.outputs.end()) {
+          partition.output_ref_count--;
+          connect_to_output = true;
+        }
+        if (partition.dependencies.find(input) != partition.dependencies.end()) {
+          connect_to_dependency = true;
+        }
+      }
+      if (is_supported && connect_to_output && !connect_to_dependency) {
+        partitions_to_merge.emplace_back(pair.first);
+      } else if (connect_to_output || connect_to_dependency) {
+        for (auto& output : node.MutableOutputDefs()) {
+          partition.dependencies.emplace(output);
+        }
+      }
+    }
+
+    if (!partitions_to_merge.empty()) {
+      std::sort(partitions_to_merge.begin(), partitions_to_merge.end());
+      TritonOpPartition& dst = partitions.at(partitions_to_merge[0]);
+      for (size_t i = partitions_to_merge.size() - 1; i > 0; --i) {
+        dst.MergeFrom(partitions.at(partitions_to_merge[i]));
+        partitions.erase(partitions_to_merge[i]);
+      }
+
+      dst.nodes.emplace_back(&node);
+      for (auto& output : node.MutableOutputDefs()) {
+        dst.outputs.emplace(output);
+      }
+      dst.output_ref_count += node.GetOutputEdgesCount();
+    } else if (is_supported) {
+      TritonOpPartition partition;
+      partition.nodes.emplace_back(&node);
+      for (auto& node_def : node.MutableOutputDefs()) {
+        partition.outputs.emplace(node_def);
+      }
+      partition.output_ref_count = node.GetOutputEdgesCount();
+      partitions.emplace(global_id++, partition);
+    }
+
+    SizeTypeVec partitions_to_erase;
+    for (auto& pair : partitions) {
+      if (pair.second.output_ref_count == 0) {
+        if (pair.second.IsValid(config_)) {
+          pair.second.outputs.clear();
+          pair.second.dependencies.clear();
+          partitions_to_fuse.emplace(pair);
+        }
+        partitions_to_erase.emplace_back(pair.first);
+      }
+    }
+
+    for (auto& id : partitions_to_erase) {
+      partitions.erase(id);
+    }
+
+    for (auto& input : node.MutableInputDefs()) {
+      if (active_outputs.find(input) != active_outputs.end()) {
+        active_outputs.at(input)--;
+        if (active_outputs.at(input) == 0) {
+          active_outputs.erase(input);
+          for (auto& pair : partitions) {
+            pair.second.dependencies.erase(input);
+          }
+        }
+      }
+    }
+
+    for (auto it = node.OutputEdgesBegin(), end = node.OutputEdgesEnd(); it != end; ++it) {
+      NodeArg* output = node.MutableOutputDefs()[it->GetSrcArgIndex()];
+      if (active_outputs.find(output) == active_outputs.end()) {
+        active_outputs.emplace(output, 1);
+      } else {
+        active_outputs.at(output)++;
+      }
+    }
+  }
+
+  SizeTypeVec partition_ids;
+  for (auto& pair : partitions_to_fuse) {
+    partition_ids.emplace_back(pair.first);
+  }
+  std::sort(partition_ids.begin(), partition_ids.end());
+
+  for (auto& id : partition_ids) {
+    auto& partition = partitions_to_fuse.at(id);
+
+    Model sub_model("test", false, logger);
+    Graph& sub_graph = sub_model.MainGraph();
+
+    NodeArgVec graph_inputs;
+    NodeArgVec node_outputs;
+    NodeArgSet graph_input_set;
+    NodeArgSet initializers;
+    ConstNodeArgVec graph_const_inputs;
+    InlinedHashMap<NodeArg*, size_t> output_consumer_counts;
+    NodeArgSet no_consumer_outputs;
+    for (auto& p_node : partition.nodes) {
+      auto& node = *p_node;
+      sub_graph.AddNode(node);
+      for (auto& input : node.MutableInputDefs()) {
+        if (initializers.find(input) != initializers.end()) {
+          continue;
+        }
+        const ONNX_NAMESPACE::TensorProto* tensor = config_.TryGetInitializer(graph, node, input);
+        if (tensor) {
+          initializers.emplace(input);
+          sub_graph.AddInitializedTensor(*tensor);
+          continue;
+        }
+
+        if (output_consumer_counts.find(input) != output_consumer_counts.end()) {
+          output_consumer_counts.at(input)--;
+          if (output_consumer_counts.at(input) == 0) {
+            output_consumer_counts.erase(input);
+          }
+        } else if (graph_input_set.find(input) == graph_input_set.end()) {
+          graph_inputs.emplace_back(input);
+          graph_input_set.insert(input);
+          graph_const_inputs.emplace_back(input);
+        }
+      }
+
+      for (auto it = p_node->OutputEdgesBegin(), end = p_node->OutputEdgesEnd(); it != end; ++it) {
+        NodeArg* output = p_node->MutableOutputDefs()[it->GetSrcArgIndex()];
+        if (output_consumer_counts.find(output) == output_consumer_counts.end()) {
+          output_consumer_counts.emplace(output, 1);
+        } else {
+          output_consumer_counts.at(output)++;
+        }
+      }
+
+      for (auto& output : node.MutableOutputDefs()) {
+        if (output->Name() != "") {
+          node_outputs.emplace_back(output);
+          if (output_consumer_counts.find(output) == output_consumer_counts.end()) {
+            no_consumer_outputs.emplace(output);
+          }
+        }
+      }
+    }
+
+    NodeArgVec graph_outputs;
+    ConstNodeArgVec graph_const_outputs;
+    for (auto& output : node_outputs) {
+      if (no_consumer_outputs.find(output) != no_consumer_outputs.end() ||
+          output_consumer_counts.find(output) != output_consumer_counts.end()) {
+        graph_outputs.emplace_back(output);
+        graph_const_outputs.emplace_back(output);
+      }
+    }
+
+    sub_graph.SetInputs(graph_const_inputs);
+    sub_graph.SetOutputs(graph_const_outputs);
+
+    auto model_proto = sub_model.ToProto();
+    std::string model_str;
+    model_proto.SerializeToString(&model_str);
+
+    Node& fused_node = graph.AddNode(graph.GenerateNodeName("TritonOp"), "TritonOp", "Fused nodes for TritonOp",
+                                     graph_inputs, graph_outputs, {}, kMSDomain);
+    fused_node.AddAttribute("onnx_key", Hash(model_str));
+    fused_node.AddAttribute("onnx_string", model_str);
+    fused_node.SetExecutionProviderType(partition.nodes[0]->GetExecutionProviderType());
+
+    for (auto& p_node : partition.nodes) {
+      graph_utils::RemoveNodeOutputEdges(graph, *p_node);
+      graph.RemoveNode(p_node->Index());
+    }
+
+    modified = true;
+  }
+
+  return Status::OK();
+}
+
+}  // namespace onnxruntime
+
+#endif  // ENABLE_TRITON

orttraining/orttraining/core/optimizer/triton_fusion.h

@@ -0,0 +1,49 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#ifdef ENABLE_TRITON
+
+#pragma once
+
+#include "core/optimizer/graph_transformer.h"
+
+namespace onnxruntime {
+
+struct TritonFusionConfig {
+  struct OpInfo {
+    std::string domain = "";
+    std::vector<ONNX_NAMESPACE::OperatorSetVersion> versions;
+    bool is_no_op = false;
+    std::unordered_map<std::string, std::string> conditions = {};
+    bool ignore_min_nodes = false;
+  };
+
+  TritonFusionConfig(std::string_view config_json = "{}");
+
+  bool IsSupported(const Graph& graph, const Node& node) const;
+  bool IsNoOp(const Node& node) const;
+  bool IgnoreMinNodes(const Node& node) const;
+  const ONNX_NAMESPACE::TensorProto* TryGetInitializer(const Graph& graph, const Node& node, NodeArg* node_arg) const;
+
+  std::unordered_map<std::string, OpInfo> ops;
+  std::string initializer = "none";
+  size_t min_nodes = 2;
+};
+
+class TritonFusion : public GraphTransformer {
+ public:
+  TritonFusion(std::string_view config_json = "{}",
+               const InlinedHashSet<std::string_view>& compatible_execution_providers = {}) noexcept
+      : GraphTransformer("TritonFusion", compatible_execution_providers), config_(config_json) {}
+
+  Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override;
+
+ private:
+  bool IsSupportedNode(const Graph& graph, const Node& node) const;
+
+  TritonFusionConfig config_;
+};
+
+}  // namespace onnxruntime
+
+#endif  // ENABLE_TRITON

orttraining/orttraining/python/orttraining_pybind_state.cc

@@ -34,6 +34,10 @@
 #include "orttraining/core/framework/torch/custom_function_register.h"
 #endif
 
+#ifdef ENABLE_TRITON
+#include "orttraining/core/framework/triton/triton_op_executor.h"
+#endif
+
 #ifdef ENABLE_TRAINING_APIS
 #include "orttraining/training_api/checkpoint.h"
 #include "orttraining/training_api/lr_scheduler.h"
@@ -546,6 +550,20 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
         return false;
 #endif
   });
+  m.def("is_triton_enabled", []() -> bool {
+#ifdef ENABLE_TRITON
+    return true;
+#else
+        return false;
+#endif
+  });
+#ifdef ENABLE_TRITON
+  m.def("register_triton_op_executor",
+        [](py::object config_getter, py::object executor_by_name, py::object executor_by_onnx) -> void {
+          training::framework::triton::TritonOpExecutor::Instance().Initialize(
+              config_getter.ptr(), executor_by_name.ptr(), executor_by_onnx.ptr());
+        });
+#endif
 
   py::class_<TrainingConfigurationResult> config_result(m, "TrainingConfigurationResult", "pbdoc(Configuration result for training.)pbdoc");
   config_result.def(py::init())

orttraining/orttraining/python/training/ort_triton/__init__.py

@@ -0,0 +1,37 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import threading
+from functools import wraps
+
+from onnxruntime.capi import _pybind_state as _C
+
+from .kernel import *  # noqa: F403
+from .triton_op_executor import call_triton_by_name, call_triton_by_onnx, get_config
+
+
+def run_once_register_triton_op_executor(f):
+    """
+    Decorator to run a function only once.
+    :param f: function to be run only once during execution time despite the number of calls
+    :return: The original function with the params passed to it if it hasn't already been run before
+    """
+
+    @wraps(f)
+    def register_triton_op_executor_wrapper(*args, **kwargs):
+        if not register_triton_op_executor_wrapper.has_run:
+            with register_triton_op_executor_wrapper.lock:
+                if not register_triton_op_executor_wrapper.has_run:
+                    register_triton_op_executor_wrapper.has_run = True
+                    f(*args, **kwargs)
+
+    register_triton_op_executor_wrapper.lock = threading.Lock()
+    register_triton_op_executor_wrapper.has_run = False
+    return register_triton_op_executor_wrapper
+
+
+@run_once_register_triton_op_executor
+def register_triton_op_executor():
+    _C.register_triton_op_executor(get_config, call_triton_by_name, call_triton_by_onnx)

orttraining/orttraining/python/training/ort_triton/_cache.py

@@ -0,0 +1,79 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+# from torch/_inductor/codecache.py
+import base64
+import functools
+import getpass
+import hashlib
+import os
+import tempfile
+from types import ModuleType
+from typing import Tuple
+
+
+@functools.lru_cache(None)
+def _cache_dir():
+    return f"{tempfile.gettempdir()}/ort_triton_{getpass.getuser()}"
+
+
+def _code_hash(code):
+    return "c" + base64.b32encode(hashlib.sha256(code.encode("utf-8")).digest())[:51].decode("utf-8").lower()
+
+
+def _get_code_path(source_code, ext, extra):
+    basename = _code_hash(source_code + extra)
+    subdir = os.path.join(_cache_dir(), basename[1:3])
+    path = os.path.join(subdir, f"{basename}.{ext}")
+    return basename, subdir, path
+
+
+def _write_atomic(path: str, source_code: str):
+    # use a temp file for thread safety
+    fd, tmp_path = tempfile.mkstemp(dir=os.path.dirname(path))
+    with os.fdopen(fd, "w") as f:
+        f.write(source_code)
+    os.rename(tmp_path, path)
+
+
+def _write(source_code, ext, extra=""):
+    basename, subdir, path = _get_code_path(source_code, ext, extra)
+    if not os.path.exists(subdir):
+        os.makedirs(subdir, exist_ok=True)
+    if not os.path.exists(path):
+        _write_atomic(path, source_code)
+    return basename, path
+
+
+class PyCodeCache:
+    cache = dict()
+    clear = staticmethod(cache.clear)
+
+    @classmethod
+    def load(cls, source_code) -> ModuleType:
+        key, path = _write(source_code, "py")
+        if key not in cls.cache:
+            with open(path) as f:
+                code = compile(f.read(), path, "exec")
+                mod = ModuleType(f"{__name__}.{key}")
+                mod.__file__ = path
+                mod.key = key
+                exec(code, mod.__dict__, mod.__dict__)
+                # another thread might set this first
+                cls.cache.setdefault(key, mod)
+        return cls.cache[key]
+
+
+class ModuleCache:
+    cache = dict()
+    clear = staticmethod(cache.clear)
+
+    @classmethod
+    def load(cls, key_func, mod_func, *args) -> Tuple[str, ModuleType]:
+        key = key_func(*args)
+        if key not in cls.cache:
+            func_name, mod = mod_func(*args)
+            cls.cache[key] = (func_name, mod)
+        return cls.cache[key]

orttraining/orttraining/python/training/ort_triton/_codegen.py

@@ -0,0 +1,485 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+"""
+Generate code for each IR node.
+Mostly, Nodes are classified into two categories:
+    1. ElementwiseKernelNode: compute a tensor from other tensors, e.g. ElementwiseKernelNode
+    2. ReduceKernelNode: perform a reduction computation on a tensor, e.g. reduce_sum/max/min
+        one or more axes are supported
+
+"""
+
+from typing import Tuple
+
+import numpy as np
+import sympy
+import torch
+from sympy.codegen.rewriting import create_expand_pow_optimization
+
+from ._common import CodeBuffer, CodegenContext, NodeVisitor
+from ._ir import (
+    ComputeNode,
+    DropoutNode,
+    ElementwiseKernelNode,
+    IONode,
+    IRNode,
+    KernelNode,
+    ModuleNode,
+    OffsetCalculator,
+    ReduceForLoopEnd,
+    ReduceForLoopStart,
+    ReduceKernelNode,
+    ReduceNode,
+)
+from ._lowering import lower
+from ._sorted_graph import SortedGraph
+from ._sympy_utils import parse_shape, sympy_dot
+from ._utils import may_add_brackets
+
+
+class TritonCodegen(NodeVisitor):
+    """
+    Specialized codegen for Triton backend.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def codegen(self, node: IRNode, context: CodegenContext, code_buffer: CodeBuffer, indent: int):
+        func = getattr(self, node.__class__.__name__)
+        assert func is not None, "unimplemented node: %s" % node.__class__.__name__
+        func(node, context, code_buffer, indent)
+
+    def _get_elementwise_offset_mask(self, offset_calc: OffsetCalculator, arg_name: str) -> Tuple[str, str]:
+        if offset_calc.is_x_reduced(arg_name):
+            return "", ""
+        if offset_calc.is_same_x_shape(arg_name):
+            return "xindex", "xmask" if offset_calc.requires_x_mask else ""
+        strides = offset_calc.get_input_strides(arg_name)
+        idx_var = [f"x{idx}" for idx in range(len(strides))]
+        expand_opt = create_expand_pow_optimization(6)
+        offset_str = str(expand_opt(sympy_dot(parse_shape(idx_var), strides)))
+        return offset_str, "xmask" if offset_calc.requires_x_mask else ""
+
+    def _get_reduce_offset_mask(self, offset_calc: OffsetCalculator, arg_name: str) -> Tuple[str, str]:
+        offset_strs = []
+        mask_strs = []
+        if not offset_calc.is_x_reduced(arg_name):
+            x_strides = offset_calc.get_x_input_strides(arg_name)
+            if offset_calc.is_same_x_shape(arg_name):
+                xindex_str = "xindex" if x_strides[-1] == sympy.Integer(1) else f"xindex * {x_strides[-1]}"
+            else:
+                idx_var = [f"x{idx}" for idx in range(len(x_strides))]
+                expand_opt = create_expand_pow_optimization(6)
+                xindex_str = str(expand_opt(sympy_dot(parse_shape(idx_var), x_strides)))
+            offset_strs.append(xindex_str)
+            if offset_calc.requires_x_mask:
+                mask_strs.append("xmask")
+
+        if not offset_calc.is_r_reduced(arg_name):
+            r_strides = offset_calc.get_r_input_strides(arg_name)
+            if offset_calc.is_same_r_shape(arg_name):
+                rindex_str = "rindex" if r_strides[-1] == sympy.Integer(1) else f"rindex * {r_strides[-1]}"
+            else:
+                idx_var = [f"r{idx}" for idx in range(len(r_strides))]
+                expand_opt = create_expand_pow_optimization(6)
+                rindex_str = str(expand_opt(sympy_dot(parse_shape(idx_var), r_strides)))
+            offset_strs.append(rindex_str)
+            if offset_calc.requires_r_mask:
+                mask_strs.append("rmask")
+
+        return " + ".join(offset_strs), " & ".join(mask_strs)
+
+    def _get_offset_mask(self, node: OffsetCalculator, arg_name: str) -> Tuple[str, str]:
+        return (
+            self._get_reduce_offset_mask(node, arg_name)
+            if node.is_reduction
+            else self._get_elementwise_offset_mask(node, arg_name)
+        )
+
+    def IONode(self, node: IONode, context: CodegenContext, code_buffer: CodeBuffer, indent: int):  # noqa: N802
+        space_indent = " " * indent
+        name = node.tensor_arg.name
+        var_name = context.get_variable_name(name)
+        internal_var_name = context.get_internal_variable_name(name)
+        assert (
+            var_name != internal_var_name
+        ), f"variable name {var_name} and its internal variable name should not be the same."
+
+        offset_str, mask_str = self._get_offset_mask(node.offset_calc, node.tensor_arg.name)
+        if offset_str:
+            offset_str = f" + {offset_str}"
+        if mask_str:
+            mask_str = f", {mask_str}"
+        if node.is_load and mask_str:
+            mask_str += ", other=0.0"
+
+        if node.is_load:
+            code_buffer += f"{space_indent}{internal_var_name} = tl.load({var_name}{offset_str}{mask_str})\n"
+        else:
+            code_buffer += f"{space_indent}tl.store({var_name}{offset_str}, {internal_var_name}{mask_str})\n"
+
+    def _gen_kernel_signature(self, node: KernelNode, context: CodegenContext, code_buffer: CodeBuffer, indent: int):
+        is_reduction = node.offset_calc.is_reduction
+        space_indent = " " * indent
+        autotune_configs_str = ""
+        for config in node.offset_calc.autotune_configs.configs:
+            if is_reduction:
+                autotune_configs_str += (
+                    f'{space_indent}        triton.Config({{"XBLOCK": {config[0]}, "RBLOCK": {config[1]}}}, '
+                    f"num_warps={config[2]}),\n"
+                )
+            else:
+                autotune_configs_str += (
+                    f'{space_indent}        triton.Config({{"XBLOCK": {config[0]}}}, num_warps={config[2]}),\n'
+                )
+        keys_str = '"xnumel", "rnumel"' if is_reduction else '"xnumel"'
+        input_args = [context.get_variable_name(input.name) for input in node.inputs]
+        input_args_str = ", ".join(input_args)
+        if input_args_str:
+            input_args_str += ", "
+
+        output_args = [context.get_variable_name(output.name) for output in node.outputs]
+        output_args_str = ", ".join(output_args) + ", "
+
+        other_input_args = "seed_cuda, " if node.has_dropout else ""
+        # Support symbolic shape if any.
+        symbolic_shape_args_str = ", ".join(node.symbolic_shape_variables)
+        if symbolic_shape_args_str:
+            other_input_args += f"{symbolic_shape_args_str}, "
+
+        blocks_str = (
+            "xnumel, rnumel, XBLOCK: tl.constexpr, RBLOCK: tl.constexpr"
+            if is_reduction
+            else "xnumel, XBLOCK: tl.constexpr"
+        )
+
+        code_buffer += (
+            f"{space_indent}@triton.autotune(\n"
+            f"{space_indent}    configs=[\n"
+            f"{autotune_configs_str}"
+            f"{space_indent}    ],\n"
+            f"{space_indent}    key=[{keys_str}],\n"
+            f"{space_indent})\n"
+            f"{space_indent}@triton.jit\n"
+            f"{space_indent}def {node.name}({input_args_str}{output_args_str}{other_input_args}{blocks_str}):\n"
+        )
+
+    def ElementwiseKernelNode(  # noqa: N802
+        self, node: ElementwiseKernelNode, context: CodegenContext, code_buffer: CodeBuffer, indent: int
+    ):
+        self._gen_kernel_signature(node, context, code_buffer, indent)
+        offset_calc = node.offset_calc
+        indent += 4
+        space_indent = " " * indent
+        code_buffer += (
+            f"{space_indent}xnumel = {offset_calc.x_numel}\n"
+            f"{space_indent}xoffset = tl.program_id(0) * XBLOCK\n"
+            f"{space_indent}xindex = xoffset + tl.arange(0, XBLOCK)\n"
+        )
+        if offset_calc.requires_x_mask:
+            code_buffer += f"{space_indent}xmask = xindex < xnumel\n"
+        for idx in range(offset_calc.x_rank):
+            if idx in offset_calc.x_compute_dims:
+                div_str = (
+                    f" // {may_add_brackets(str(offset_calc.x_strides[idx]))}" if idx != offset_calc.x_rank - 1 else ""
+                )
+                mod_str = f" % {may_add_brackets(str(offset_calc.x_dims[idx]))}" if idx != 0 else ""
+                code_buffer += f"{space_indent}x{idx} = xindex{div_str}{mod_str}\n"
+        code_buffer += "\n"
+
+        if node.has_dropout:
+            code_buffer += (
+                f"{space_indent}t_seed_cuda = tl.load(seed_cuda)\n"
+                f"{space_indent}t_seed_cuda = tl.broadcast_to(t_seed_cuda, [XBLOCK])\n"
+            )
+
+        for ir_node in node.sub_nodes:
+            ir_node.codegen(self, context, code_buffer, indent)
+
+    def ReduceKernelNode(  # noqa: N802
+        self, node: ReduceKernelNode, context: CodegenContext, code_buffer: CodeBuffer, indent: int
+    ):
+        self._gen_kernel_signature(node, context, code_buffer, indent)
+        offset_calc = node.offset_calc
+        indent += 4
+        space_indent = " " * indent
+        code_buffer += (
+            f"{space_indent}xnumel = {offset_calc.x_numel}\n"
+            f"{space_indent}rnumel = {offset_calc.r_numel}\n"
+            f"{space_indent}xoffset = tl.program_id(0) * XBLOCK\n"
+            f"{space_indent}xindex = xoffset + tl.arange(0, XBLOCK)[:, None]\n"
+            f"{space_indent}rbase = tl.arange(0, RBLOCK)[None, :]\n"
+        )
+        if offset_calc.requires_x_mask:
+            code_buffer += f"{space_indent}xmask = xindex < xnumel\n"
+        for idx in range(offset_calc.x_rank):
+            if idx in offset_calc.x_compute_dims:
+                div_str = (
+                    f" // {may_add_brackets(str(offset_calc.x_strides[idx]))}" if idx != offset_calc.x_rank - 1 else ""
+                )
+                mod_str = f" % {may_add_brackets(str(offset_calc.x_dims[idx]))}" if idx != 0 else ""
+                code_buffer += f"{space_indent}x{idx} = xindex{div_str}{mod_str}\n"
+        code_buffer += "\n"
+
+        if node.has_dropout:
+            code_buffer += f"{space_indent}t_seed_cuda = tl.load(seed_cuda)\n"
+            code_buffer += f"{space_indent}t_seed_cuda = tl.broadcast_to(t_seed_cuda, [XBLOCK, RBLOCK])\n"
+
+        if not offset_calc.autotune_configs.requires_for_loop:
+            code_buffer += f"{space_indent}rindex = rbase\n"
+            if offset_calc.requires_r_mask:
+                code_buffer += f"{space_indent}rmask = rindex < rnumel\n"
+            for idx in range(offset_calc.r_rank):
+                if idx in offset_calc.r_compute_dims:
+                    div_str = (
+                        f" // {may_add_brackets(str(offset_calc.r_strides[idx]))}"
+                        if idx != offset_calc.r_rank - 1
+                        else ""
+                    )
+                    mod_str = f" % {may_add_brackets(str(offset_calc.r_dims[idx]))}" if idx != 0 else ""
+                    code_buffer += f"{space_indent}r{idx} = rindex{div_str}{mod_str}\n"
+
+        for ir_node in node.sub_nodes:
+            ir_node.codegen(self, context, code_buffer, indent)
+            if isinstance(ir_node, ReduceForLoopStart):
+                indent += 4
+            elif isinstance(ir_node, ReduceForLoopEnd):
+                indent -= 4
+
+    _COMPUTE_CODE_TEMPLATES = {
+        "Add": "{indent}{o0} = {i0} + {i1}\n",
+        "Sub": "{indent}{o0} = {i0} - {i1}\n",
+        "Mul": "{indent}{o0} = {i0} * {i1}\n",
+        "Div": "{indent}{o0} = {i0} / {i1}\n",
+        "Relu": "{indent}{o0} = tl.maximum({i0}, 0.0)\n",
+        "Pow": "{indent}{o0} = tl.math.pow({i0}, {i1})\n",
+        "Pow2": "{indent}{o0} = {i0} * {i0}\n",
+        "Pow3": "{indent}{o0} = {i0} * {i0} * {i0}\n",
+        "Sqrt": "{indent}{o0} = tl.sqrt({i0})\n",
+        "Rsqrt": "{indent}{o0} = 1.0 / tl.sqrt({i0})\n",
+        "Cast": "{indent}{o0} = {i0}.to(tl.{dtype})\n",
+        "CastBool": "{indent}{o0} = {i0} != 0\n",
+        "Erf": "{indent}{o0} = tl.libdevice.erf({i0})\n",
+        "Gelu": "{indent}{o0} = (tl.libdevice.erf({i0} / 1.41421356237) + 1.0) * 0.5\n",
+        "Exp": "{indent}{o0} = tl.exp({i0})\n",
+        "Tanh": "{indent}{o0} = tl.libdevice.tanh({i0})\n",
+        "Where": "{indent}{o0} = tl.where({i0}, {i1}, {i2})\n",
+        "Sigmoid": "{indent}{o0} = tl.sigmoid({i0})\n",
+        "Log": "{indent}{o0} = tl.log({i0})\n",
+        "DropoutGrad": "{indent}p = 1 - {i2}\n{indent}{o0} = tl.where({i1}, {i0} / p, 0.0)\n",
+        "Identity": "{indent}{o0} = {i0}\n",
+    }
+
+    def ComputeNode(  # noqa: N802
+        self, node: ComputeNode, context: CodegenContext, code_buffer: CodeBuffer, indent: int
+    ):
+        space_indent = " " * indent
+        kwargs = {}
+        for idx, input in enumerate(node.inputs):
+            kwargs[f"i{idx}"] = context.get_internal_variable_name(input.name)
+        for idx, output in enumerate(node.outputs):
+            kwargs[f"o{idx}"] = context.get_internal_variable_name(output.name)
+
+        op_type = node.op_type
+        if op_type == "Pow":
+            if kwargs["i1"] == 2:
+                op_type = "Pow2"
+            elif kwargs["i1"] == 3:
+                op_type = "Pow3"
+            elif kwargs["i1"] == 0.5:
+                op_type = "Sqrt"
+
+        if op_type == "Cast":
+            from_dtype = node.inputs[0].dtype.type
+            to_dtype = node.outputs[0].dtype.type
+            if from_dtype == to_dtype:
+                op_type = "Identity"
+            elif to_dtype == np.bool_:
+                op_type = "CastBool"
+            else:
+                kwargs["dtype"] = to_dtype.__name__
+
+        if op_type == "Sum":
+            output_var = kwargs["o0"]
+            formula = " + ".join([kwargs[f"i{idx}"] for idx in range(len(node.inputs))])
+            code_buffer += f"{space_indent}{output_var} = {formula}\n"
+            return
+
+        code_buffer += TritonCodegen._COMPUTE_CODE_TEMPLATES[op_type].format(indent=space_indent, **kwargs)
+
+    def ReduceNode(self, node: ReduceNode, context: CodegenContext, code_buffer: CodeBuffer, indent: int):  # noqa: N802
+        space_indent = " " * indent
+        input_var_name = context.get_internal_variable_name(node.inputs[0].name)
+        output_var_name = context.get_internal_variable_name(node.outputs[0].name)
+        masks = []
+        if node.offset_calc.requires_x_mask:
+            masks.append("xmask")
+        if node.offset_calc.requires_r_mask:
+            masks.append("rmask")
+        if len(masks) > 0:
+            masks_str = " & ".join(masks)
+            code_buffer += (
+                f"{space_indent}{input_var_name} = tl.where({masks_str}, {input_var_name}, {node.default_value})\n"
+            )
+        code_buffer += f"{space_indent}{output_var_name} = {node.triton_func}({input_var_name}, axis=1)[:, None]\n"
+
+    def ReduceForLoopStart(  # noqa: N802
+        self, node: ReduceForLoopStart, context: CodegenContext, code_buffer: CodeBuffer, indent: int
+    ):
+        space_indent = " " * indent
+        offset_calc = node.offset_calc
+        for reduce_node in node.reduce_nodes:
+            tmp_var_name = "tmp_" + context.get_internal_variable_name(reduce_node.outputs[0].name)
+            code_buffer += (
+                f"{space_indent}{tmp_var_name} = "
+                f"tl.zeros([XBLOCK, RBLOCK], tl.float32) + {reduce_node.default_value}\n"
+            )
+        code_buffer += (
+            f"{space_indent}for roffset in range(0, rnumel, RBLOCK):\n{space_indent}    rindex = rbase + roffset\n"
+        )
+        if offset_calc.requires_r_mask:
+            code_buffer += f"{space_indent}    rmask = rindex < rnumel\n"
+        for idx in range(offset_calc.r_rank):
+            if idx in offset_calc.r_compute_dims:
+                div_str = (
+                    f" // {may_add_brackets(str(offset_calc.r_strides[idx]))}" if idx != offset_calc.r_rank - 1 else ""
+                )
+                mod_str = f" % {may_add_brackets(str(offset_calc.r_dims[idx]))}" if idx != 0 else ""
+                code_buffer += f"{space_indent}    r{idx} = rindex{div_str}{mod_str}\n"
+
+    def ReduceForLoopEnd(  # noqa: N802
+        self, node: ReduceForLoopEnd, context: CodegenContext, code_buffer: CodeBuffer, indent: int
+    ):
+        space_indent = " " * indent
+        offset_calc = node.offset_calc
+        masks = []
+        if offset_calc.requires_x_mask:
+            masks.append("xmask")
+        if offset_calc.requires_r_mask:
+            masks.append("rmask")
+        masks_str = " & ".join(masks)
+        for reduce_node in node.reduce_nodes:
+            input_var_name = context.get_internal_variable_name(reduce_node.inputs[0].name)
+            output_var_name = context.get_internal_variable_name(reduce_node.outputs[0].name)
+            tmp_output_var_name = "tmp_" + output_var_name
+            if reduce_node.op_type == "ReduceSum":
+                if not masks_str:
+                    code_buffer += f"{space_indent}{tmp_output_var_name} = {tmp_output_var_name} + {input_var_name}\n"
+                else:
+                    code_buffer += (
+                        f"{space_indent}{tmp_output_var_name} = "
+                        f"tl.where({masks_str}, {tmp_output_var_name} + {input_var_name}, {tmp_output_var_name})\n"
+                    )
+            else:
+                op_str = " < " if reduce_node.op_type == "ReduceMax" else " > "
+                if masks_str:
+                    masks_str += " & "
+                code_buffer += (
+                    f"{space_indent}{tmp_output_var_name} = tl.where("
+                    f"{masks_str}({tmp_output_var_name}{op_str}{input_var_name}), "
+                    f"{input_var_name}, {tmp_output_var_name})\n"
+                )
+        space_indent_outside = " " * (indent - 4)
+        for reduce_node in node.reduce_nodes:
+            output_var_name = context.get_internal_variable_name(reduce_node.outputs[0].name)
+            input_var_name = "tmp_" + output_var_name
+            code_buffer += (
+                f"{space_indent_outside}{output_var_name} = "
+                f"{reduce_node.triton_func}({input_var_name}, axis=1)[:, None]\n"
+            )
+
+    def DropoutNode(  # noqa: N802
+        self, node: DropoutNode, context: CodegenContext, code_buffer: CodeBuffer, indent: int
+    ):
+        space_indent = " " * indent
+        input_var_name = context.get_internal_variable_name(node.inputs[0].name)
+        p_var_name = context.get_internal_variable_name(node.inputs[1].name)
+        output_var_name = context.get_internal_variable_name(node.outputs[0].name)
+        mask_var_name = (
+            context.get_internal_variable_name(node.outputs[1].name)
+            if len(node.outputs) >= 2
+            else "dropout_mask_output"
+        )
+        offset_str = f"{node.global_offset} + " if node.global_offset != sympy.Integer(0) else ""
+        offset_str += self._get_offset_mask(node.offset_calc, node.inputs[0].name)[0]
+        code_buffer += (
+            f"{space_indent}p = 1 - {p_var_name}\n"
+            f"{space_indent}random = tl.rand(t_seed_cuda, {offset_str})\n"
+            f"{space_indent}{mask_var_name} = random < p\n"
+            f"{space_indent}{output_var_name} = tl.where({mask_var_name}, {input_var_name} / p, 0.0)\n"
+        )
+
+    def ModuleNode(self, node: ModuleNode, context: CodegenContext, code_buffer: CodeBuffer, indent: int):  # noqa: N802
+        space_indent = " " * indent
+        code_buffer += (
+            f"{space_indent}import triton\n"
+            f"{space_indent}import triton.language as tl\n"
+            f"{space_indent}import torch\n"
+        )
+
+        for kernel_node in node.kernels:
+            code_buffer += "\n\n"
+            kernel_node.codegen(self, CodegenContext(kernel_node.var_map), code_buffer, indent)
+
+        input_args = ", ".join([context.get_variable_name(input.name) for input in node.inputs])
+        code_buffer += f"\n\n{space_indent}def {node.func_name}({input_args}):\n"
+
+        indent += 4
+        space_indent = " " * indent
+
+        if node.has_dropout:
+            code_buffer += (
+                f'{space_indent}seed_cuda = torch.randint(2**31, size=(), dtype=torch.int64, device="cuda")\n\n'
+            )
+
+        for idx, kernel_node in enumerate(node.kernels):
+            if idx != 0:
+                code_buffer += "\n"
+            # Allocate output tensor.
+            for output in kernel_node.outputs:
+                torch_dtype = torch.from_numpy(np.zeros(1, dtype=output.dtype)).dtype
+                # Workaround for DLPack which doesn't support bool.
+                if torch_dtype == torch.bool:
+                    torch_dtype = torch.uint8
+                code_buffer += (
+                    f"{space_indent}{context.get_variable_name(output.name)} = "
+                    f'torch.empty({tuple(output.shape)}, dtype={torch_dtype}, device="cuda")\n'
+                )
+            kernel_args_str = ", ".join([context.get_variable_name(input.name) for input in kernel_node.inputs])
+            if kernel_args_str:
+                kernel_args_str += ", "
+            kernel_args_str += ", ".join([context.get_variable_name(output.name) for output in kernel_node.outputs])
+            # TODO: support other kinds of variable args, such as symbolic shape variable.
+            if kernel_node.has_dropout:
+                kernel_args_str += ", seed_cuda"
+
+            if isinstance(kernel_node, ReduceKernelNode):
+                code_buffer += (
+                    f"{space_indent}x_numel = {kernel_node.offset_calc.x_numel}\n"
+                    f"{space_indent}r_numel = {kernel_node.offset_calc.r_numel}\n"
+                    f'{space_indent}grid = lambda meta: (triton.cdiv(x_numel, meta["XBLOCK"]),)\n'
+                    f"{space_indent}{kernel_node.name}[grid]({kernel_args_str}, x_numel, r_numel)\n"
+                )
+            else:
+                code_buffer += (
+                    f"{space_indent}n_elements = {kernel_node.offset_calc.x_numel}\n"
+                    f'{space_indent}grid = lambda meta: (triton.cdiv(n_elements, meta["XBLOCK"]),)\n'
+                    f"{space_indent}{kernel_node.name}[grid]({kernel_args_str}, n_elements)\n"
+                )
+
+            for name in node.cross_kernel_args_to_delete[idx]:
+                code_buffer += f"{space_indent}del {name}\n"
+
+        return_output_str = ", ".join([context.get_variable_name(output.name) for output in node.outputs])
+        code_buffer += f"\n{space_indent}return {return_output_str}\n"
+
+
+def codegen(func_name: str, sorted_graph: SortedGraph) -> str:
+    module_node = lower(func_name, sorted_graph)
+    code_buffer = CodeBuffer()
+    module_node.codegen(TritonCodegen(), CodegenContext(module_node.var_map), code_buffer)
+    return str(code_buffer)

orttraining/orttraining/python/training/ort_triton/_common.py

@@ -0,0 +1,188 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from abc import abstractmethod
+from typing import Any, Dict, List, Tuple
+
+import sympy
+from onnx import GraphProto, NodeProto, TensorProto
+
+from ._sympy_utils import parse_shape
+from ._utils import get_attribute, get_reduce_info, next_power_of_2
+
+
+class CodegenContext:
+    """
+    record variable name mapping in term of IRnodes.
+    """
+
+    def __init__(self, var_map: Dict[str, str]):
+        self._var_map: Dict[str, str] = {**var_map}
+
+    # Get variable name by the node arg name in ONNX graph.
+    def get_variable_name(self, name: str) -> str:
+        return self._var_map[name]
+
+    # For some operators such as data load/store, we need an internal variable name inside the kernel function.
+    def get_internal_variable_name(self, name: str) -> str:
+        var_name = self._var_map[name]
+        return self._var_map[var_name] if var_name in self._var_map else var_name
+
+
+class CodeBuffer:
+    def __init__(self):
+        self.buffer: List[str] = []
+
+    def __iadd__(self, other: str):
+        self.buffer.append(other)
+        return self
+
+    def __str__(self):
+        return "".join(self.buffer)
+
+
+class NodeVisitor:
+    @abstractmethod
+    def codegen(self, node: Any, context: CodegenContext, code_buffer: CodeBuffer, indent: int):
+        pass
+
+
+class TensorInfo:
+    """
+    Represent a input/output tensor of a node.
+    """
+
+    def __init__(self, dtype: TensorProto.DataType, shape: List[Any]):
+        self._dtype: TensorProto.DataType = dtype
+        self._shape: List[sympy.Expr] = parse_shape(shape)
+
+    @property
+    def dtype(self) -> TensorProto.DataType:
+        return self._dtype
+
+    @property
+    def shape(self) -> List[sympy.Expr]:
+        return self._shape
+
+
+def _infer_elementwise_shape(input_infos: List[TensorInfo]) -> List[sympy.Expr]:
+    max_len = max([len(input_info.shape) for input_info in input_infos])
+    output_shape: List[sympy.Expr] = [sympy.Integer(1)] * max_len
+    for input_info in input_infos:
+        offset = max_len - len(input_info.shape)
+        for i in range(len(input_info.shape)):
+            if not input_info.shape[i].is_number or input_info.shape[i] != 1:
+                output_shape[i + offset] = input_info.shape[i]
+    return output_shape
+
+
+def _infer_elementwise(node: NodeProto, input_infos: List[TensorInfo], graph: GraphProto) -> List[TensorInfo]:
+    return [TensorInfo(input_infos[0].dtype, _infer_elementwise_shape(input_infos))]
+
+
+def _infer_where(node: NodeProto, input_infos: List[TensorInfo], graph: GraphProto) -> List[TensorInfo]:
+    return [TensorInfo(input_infos[1].dtype, _infer_elementwise_shape(input_infos))]
+
+
+def _infer_reduction(node: NodeProto, input_infos: List[TensorInfo], graph: GraphProto) -> List[TensorInfo]:
+    input_rank = len(input_infos[0].shape)
+    keep_dims, axes = get_reduce_info(node, graph, input_rank)
+    axes = [axis + input_rank if axis < 0 else axis for axis in axes]
+    axes.sort()
+    shape = [input_infos[0].shape[i] for i in range(input_rank) if i not in axes]
+    if keep_dims:
+        for axis in axes:
+            shape.insert(axis, sympy.Integer(1))
+    return [TensorInfo(input_infos[0].dtype, shape)]
+
+
+def _infer_unary(node: NodeProto, input_infos: List[TensorInfo], graph: GraphProto) -> List[TensorInfo]:
+    return [input_infos[0]]
+
+
+def _infer_cast(node: NodeProto, input_infos: List[TensorInfo], graph: GraphProto) -> List[TensorInfo]:
+    dtype = get_attribute(node, "to", TensorProto.UNDEFINED)
+    assert dtype != TensorProto.UNDEFINED
+    return [TensorInfo(dtype, input_infos[0].shape)]
+
+
+def _infer_dropout(node: NodeProto, input_infos: List[TensorInfo], graph: GraphProto) -> List[TensorInfo]:
+    return [input_infos[0], TensorInfo(TensorProto.BOOL, input_infos[0].shape)]
+
+
+class TypeAndShapeInfer:
+    _INFER_FUNC_MAP = {
+        "Add": _infer_elementwise,
+        "Sub": _infer_elementwise,
+        "Mul": _infer_elementwise,
+        "Div": _infer_elementwise,
+        "Pow": _infer_elementwise,
+        "Sqrt": _infer_elementwise,
+        "Exp": _infer_elementwise,
+        "Where": _infer_where,
+        "Rsqrt": _infer_elementwise,
+        "Cast": _infer_cast,
+        "Dropout": _infer_dropout,
+        "DropoutGrad": _infer_unary,
+        "Identity": _infer_unary,
+        "ReduceSum": _infer_reduction,
+        "ReduceMax": _infer_reduction,
+        "ReduceMin": _infer_reduction,
+        "Sum": _infer_elementwise,
+    }
+
+    @classmethod
+    def infer(cls, node: NodeProto, input_infos: List[TensorInfo], graph: GraphProto) -> List[TensorInfo]:
+        if node.op_type not in cls._INFER_FUNC_MAP:
+            raise NotImplementedError(f"Unsupported op type: {node.op_type}")
+        return cls._INFER_FUNC_MAP[node.op_type](node, input_infos, graph)
+
+
+class AutotuneConfigs:
+    """
+    Generate all autotune configs for a kernel function by it's xnumel and rnumel.
+    A config is a tuple of (xblock, rblock, num_warps).
+    If it's elementwise kernel, the rnumel is 1.
+    If it's reduction kernel on last contiguous dimensions, the contiguous flag is True.
+    """
+
+    def __init__(self, x_numel: int, r_numel: int, contiguous: bool):
+        self.configs: List[Tuple[int, int, int]] = self._gen_autotune_configs(x_numel, r_numel, contiguous)
+        self.requires_for_loop: bool = any(config[1] < r_numel for config in self.configs)
+
+    def _num_warps(self, x: int, r: int) -> int:
+        return min(max(x * r // 256, 2), 8)
+
+    def _gen_config(self, xnp2: int, rnp2: int, x: int, r: int) -> Tuple[int, int, int]:
+        x = min(x, xnp2)
+        r = min(r, rnp2)
+        return x, r, self._num_warps(x, r)
+
+    # TODO: we need to tune more kernels to get more reasonable configs for better performance.
+    def _gen_autotune_configs(self, x_numel: int, r_numel: int, contiguous: bool) -> List[Tuple[int, int, int]]:
+        configs = []
+        xnp2 = next_power_of_2(x_numel)
+        if r_numel == 1:
+            configs.append(self._gen_config(xnp2, 1, 1024, 1))
+            if xnp2 > 1024:
+                configs.append(self._gen_config(xnp2, 1, 2048, 1))
+            return configs
+        rnp2 = next_power_of_2(r_numel)
+        if contiguous:
+            configs.append(self._gen_config(xnp2, rnp2, 1, 2048))
+            if rnp2 > 2048:
+                configs.append(self._gen_config(xnp2, rnp2, 1, 4096))
+            elif rnp2 <= 256:
+                x = min(xnp2, 256 // rnp2 * 2)
+                configs.append(self._gen_config(xnp2, rnp2, x, rnp2))
+        else:
+            config_set = {
+                self._gen_config(xnp2, rnp2, 1, 2048),
+                self._gen_config(xnp2, rnp2, 4, 512),
+                self._gen_config(xnp2, rnp2, 8, 512),
+                self._gen_config(xnp2, rnp2, 32, 128),
+            }
+            configs = list(config_set)
+        return configs

orttraining/orttraining/python/training/ort_triton/_decompose.py

@@ -0,0 +1,371 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+"""
+Decompose a complicated op into a series of simple ops.
+"simple ops" can be executed in one pass
+"""
+
+from typing import List
+
+import numpy as np
+import sympy
+from onnx import GraphProto, NodeProto, TensorProto, helper
+
+from ._utils import get_attribute, get_reduce_info, to_numpy_type
+
+
+def _is_half_dtype(dtype: int):
+    return dtype in [TensorProto.FLOAT16, TensorProto.BFLOAT16]
+
+
+class DecomposeDispatch:
+    """
+    A node does only responsible for a single computation or a type of triton ops.
+    For those compound Onnx nodes, like softmax/layernorm/groupnorm, etc., we need to decompose them into a series of
+    simple ops.
+    """
+
+    def __init__(self):
+        self.count = 0
+
+    def __call__(self, node: NodeProto, graph: GraphProto, **kwargs) -> List[NodeProto]:
+        op_type = node.op_type
+        if not hasattr(self, op_type):
+            raise NotImplementedError(f"Not implemented for op type: {op_type}")
+        return getattr(self, op_type)(node, graph, **kwargs)
+
+    def __contains__(self, node: NodeProto) -> bool:
+        return hasattr(self, node.op_type)
+
+    def _get_unique_var_name(self, prefix):
+        self.count += 1
+        return prefix + str(self.count)
+
+    def _new_node(self, node_name, op_type, inputs, outputs=None, **kwargs):
+        name = self._get_unique_var_name(f"{node_name}_{op_type}_")
+        if outputs is None:
+            outputs = [f"{name}_out"]
+        for idx, output in enumerate(outputs):
+            if output is None:
+                outputs[idx] = f"{name}_out{idx}"
+        return helper.make_node(op_type, inputs, outputs, name, **kwargs), *outputs
+
+    def _get_dtype_and_shape(self, arg_name: str, **kwargs):
+        node_arg_infos = kwargs["node_arg_infos"]
+        arg_info = node_arg_infos[arg_name]
+        return arg_info.dtype, arg_info.shape
+
+    def _decompose_elementwise_precision(self, node: NodeProto, graph: GraphProto, **kwargs):
+        x = node.input[0]
+        dtype, _ = self._get_dtype_and_shape(x, **kwargs)
+        if not _is_half_dtype(dtype):
+            return [node]
+        node_name = node.name
+        y = node.output[0]
+        op_type = node.op_type
+        inputs = [input for input in node.input]
+        cast_nodes = []
+        for idx, input in enumerate(inputs):
+            dtype, _ = self._get_dtype_and_shape(input, **kwargs)
+            if _is_half_dtype(dtype):
+                cast_node, cast_out = self._new_node(node_name, "Cast", [input], to=TensorProto.FLOAT)
+                cast_nodes.append(cast_node)
+                inputs[idx] = cast_out
+        op_node, op_out = self._new_node(node_name, op_type, inputs)
+        cast_node1, _ = self._new_node(node_name, "Cast", [op_out], outputs=[y], to=dtype)
+        return [*cast_nodes, op_node, cast_node1]
+
+    def Exp(self, node: NodeProto, graph: GraphProto, **kwargs):  # noqa: N802
+        return self._decompose_elementwise_precision(node, graph, **kwargs)
+
+    def Pow(self, node: NodeProto, graph: GraphProto, **kwargs):  # noqa: N802
+        return self._decompose_elementwise_precision(node, graph, **kwargs)
+
+    def Sqrt(self, node: NodeProto, graph: GraphProto, **kwargs):  # noqa: N802
+        return self._decompose_elementwise_precision(node, graph, **kwargs)
+
+    def LayerNormalization(self, node: NodeProto, graph: GraphProto, **kwargs):  # noqa: N802
+        node_name = node.name
+        x = node.input[0]
+        w = node.input[1]
+        b = node.input[2]
+        y = node.output[0]
+        mean = node.output[1] if len(node.output) > 1 and node.output[1] else None
+        inv_std_dev = node.output[2] if len(node.output) > 2 and node.output[2] else None
+        axis = get_attribute(node, "axis", -1)
+        epsilon = get_attribute(node, "epsilon", 1e-05)
+        xdtype, shape = self._get_dtype_and_shape(x, **kwargs)
+        wdtype, _ = self._get_dtype_and_shape(w, **kwargs)
+        is_x_half = _is_half_dtype(xdtype)
+        is_w_half = _is_half_dtype(wdtype)
+        if is_x_half or is_w_half:
+            decomposed_nodes = []
+            if is_x_half:
+                cast_node, x = self._new_node(node_name, "Cast", [x], to=TensorProto.FLOAT)
+                decomposed_nodes.append(cast_node)
+            if is_w_half:
+                cast_node1, w = self._new_node(node_name, "Cast", [w], to=TensorProto.FLOAT)
+                cast_node2, b = self._new_node(node_name, "Cast", [b], to=TensorProto.FLOAT)
+                decomposed_nodes.append(cast_node1)
+                decomposed_nodes.append(cast_node2)
+            outputs = [None] if is_w_half else [y]
+            if mean is not None:
+                outputs.append(mean)
+            if inv_std_dev is not None:
+                outputs.append(inv_std_dev)
+            layer_norm_node_outputs = self._new_node(
+                node_name, "LayerNormalization", [x, w, b], outputs=outputs, axis=axis, epsilon=epsilon
+            )
+            decomposed_nodes.append(layer_norm_node_outputs[0])
+            if is_w_half:
+                cast_node3, _ = self._new_node(node_name, "Cast", [layer_norm_node_outputs[1]], outputs=[y], to=wdtype)
+                decomposed_nodes.append(cast_node3)
+            return decomposed_nodes
+        rank = len(shape)
+        if axis < 0:
+            axis += rank
+        axes = list(range(axis, rank))
+        epsilon_tensor = helper.make_tensor(name="epsilon_const", data_type=xdtype, dims=(1,), vals=np.array([epsilon]))
+        const_node, const_out = self._new_node(node_name, "Constant", [], value=epsilon_tensor)
+        reducemean_node, reducemean_out = self._new_node(node_name, "ReduceMean", [x], outputs=[mean], axes=axes)
+        sub_node, sub_out = self._new_node(node_name, "Sub", [x, reducemean_out])
+        mul_node, mul_out = self._new_node(node_name, "Mul", [sub_out, sub_out])
+        reducemean_node1, reducemean_out1 = self._new_node(node_name, "ReduceMean", [mul_out], axes=axes)
+        add_node, add_out = self._new_node(node_name, "Add", [reducemean_out1, const_out])
+        rsqrt_node, rsqrt_out = self._new_node(node_name, "Rsqrt", [add_out], outputs=[inv_std_dev])
+        mul_node1, mul_out1 = self._new_node(node_name, "Mul", [sub_out, rsqrt_out])
+        mul_node2, mul_out2 = self._new_node(node_name, "Mul", [w, mul_out1])
+        add_node1, _ = self._new_node(node_name, "Add", [b, mul_out2], outputs=[y])
+        return [
+            const_node,
+            reducemean_node,
+            sub_node,
+            mul_node,
+            reducemean_node1,
+            add_node,
+            rsqrt_node,
+            mul_node1,
+            mul_node2,
+            add_node1,
+        ]
+
+    def LayerNormalizationGrad(self, node: NodeProto, graph: GraphProto, **kwargs):  # noqa: N802
+        node_name = node.name
+        dy = node.input[0]
+        x = node.input[1]
+        w = node.input[2]
+        mean = node.input[3]
+        inv_std_dev = node.input[4]
+        dx = node.output[0]
+        dw = node.output[1] if len(node.output) > 1 and node.output[1] else None
+        db = node.output[2] if len(node.output) > 2 and node.output[2] else None
+        axis = get_attribute(node, "axis", -1)
+        xdtype, shape = self._get_dtype_and_shape(x, **kwargs)
+        wdtype, _ = self._get_dtype_and_shape(w, **kwargs)
+        is_x_half = _is_half_dtype(xdtype)
+        is_w_half = _is_half_dtype(wdtype)
+        if is_x_half or is_w_half:
+            decomposed_nodes = []
+            if is_x_half:
+                cast_node, x = self._new_node(node_name, "Cast", [x], to=TensorProto.FLOAT)
+                decomposed_nodes.append(cast_node)
+            if is_w_half:
+                cast_node1, dy = self._new_node(node_name, "Cast", [dy], to=TensorProto.FLOAT)
+                cast_node2, w = self._new_node(node_name, "Cast", [w], to=TensorProto.FLOAT)
+                decomposed_nodes.append(cast_node1)
+                decomposed_nodes.append(cast_node2)
+            outputs = [None] if is_x_half else [dx]
+            if dw is not None:
+                outputs.append(None if is_w_half else dw)
+            if db is not None:
+                outputs.append(None if is_w_half else db)
+            layer_norm_grad_node_outputs = self._new_node(
+                node_name, "LayerNormalizationGrad", [dy, x, w, mean, inv_std_dev], outputs=outputs, axis=axis
+            )
+            decomposed_nodes.append(layer_norm_grad_node_outputs[0])
+            if is_x_half:
+                cast_node3, _ = self._new_node(
+                    node_name, "Cast", [layer_norm_grad_node_outputs[1]], outputs=[dx], to=xdtype
+                )
+                decomposed_nodes.append(cast_node3)
+            if dw is not None and is_w_half:
+                cast_node4, _ = self._new_node(
+                    node_name, "Cast", [layer_norm_grad_node_outputs[2]], outputs=[dw], to=wdtype
+                )
+                decomposed_nodes.append(cast_node4)
+            if db is not None and is_w_half:
+                cast_node5, _ = self._new_node(
+                    node_name, "Cast", [layer_norm_grad_node_outputs[3]], outputs=[db], to=wdtype
+                )
+                decomposed_nodes.append(cast_node5)
+            return decomposed_nodes
+        rank = len(shape)
+        if axis < 0:
+            axis += rank
+        axes = list(range(axis, rank))
+        sub_node, sub_out = self._new_node(node_name, "Sub", [x, mean])
+        mul_node, mul_out = self._new_node(node_name, "Mul", [sub_out, inv_std_dev])
+        mul_node1, mul_out1 = self._new_node(node_name, "Mul", [w, dy])
+        mul_node2, mul_out2 = self._new_node(node_name, "Mul", [mul_out, mul_out1])
+        reducemean_node, reducemean_out = self._new_node(node_name, "ReduceMean", [mul_out2], axes=axes)
+        reducemean_node1, reducemean_out1 = self._new_node(node_name, "ReduceMean", [mul_out1], axes=axes)
+        mul_node3, mul_out3 = self._new_node(node_name, "Mul", [reducemean_out, mul_out])
+        add_node, add_out = self._new_node(node_name, "Add", [mul_out3, reducemean_out1])
+        sub_node1, sub_out1 = self._new_node(node_name, "Sub", [mul_out1, add_out])
+        mul_node4, _ = self._new_node(node_name, "Mul", [sub_out1, inv_std_dev], outputs=[dx])
+        decomposed_nodes = [
+            sub_node,
+            mul_node,
+            mul_node1,
+            mul_node2,
+            reducemean_node,
+            reducemean_node1,
+            mul_node3,
+            add_node,
+            sub_node1,
+            mul_node4,
+        ]
+        dw_axes = list(range(axis))
+        if dw is not None:
+            mul_node5, mul_out5 = self._new_node(node_name, "Mul", [dy, mul_out])
+            reducesum_node, _ = self._new_node(
+                node_name, "ReduceSum", [mul_out5], outputs=[dw], axes=dw_axes, keepdims=0
+            )
+            decomposed_nodes.extend([mul_node5, reducesum_node])
+        if db is not None:
+            reducesum_node1, _ = self._new_node(node_name, "ReduceSum", [dy], outputs=[db], axes=dw_axes, keepdims=0)
+            decomposed_nodes.append(reducesum_node1)
+        return decomposed_nodes
+
+    def Softmax(self, node: NodeProto, graph: GraphProto, **kwargs):  # noqa: N802
+        node_name = node.name
+        x = node.input[0]
+        y = node.output[0]
+        axis = get_attribute(node, "axis", -1)
+        dtype, _ = self._get_dtype_and_shape(x, **kwargs)
+        if _is_half_dtype(dtype):
+            cast_node, x = self._new_node(node_name, "Cast", [x], to=TensorProto.FLOAT)
+            softmax_node, softmax_out = self._new_node(node_name, "Softmax", [x], axis=axis)
+            cast_node1, _ = self._new_node(node_name, "Cast", [softmax_out], outputs=[y], to=dtype)
+            return [cast_node, softmax_node, cast_node1]
+        max_node, max_out = self._new_node(node_name, "ReduceMax", [x], axes=[axis])
+        sub_node, sub_out = self._new_node(node_name, "Sub", [x, max_out])
+        exp_node, exp_out = self._new_node(node_name, "Exp", [sub_out])
+        sum_node, sum_out = self._new_node(node_name, "ReduceSum", [exp_out], axes=[axis])
+        div_node, _ = self._new_node(node_name, "Div", [exp_out, sum_out], outputs=[y])
+        return [max_node, sub_node, exp_node, sum_node, div_node]
+
+    def SoftmaxGrad_13(self, node: NodeProto, graph: GraphProto, **kwargs):  # noqa: N802
+        node_name = node.name
+        dy = node.input[0]
+        y = node.input[1]
+        dx = node.output[0]
+        axis = get_attribute(node, "axis", -1)
+        dtype, _ = self._get_dtype_and_shape(dy, **kwargs)
+        if _is_half_dtype(dtype):
+            cast_node, dy = self._new_node(node_name, "Cast", [dy], to=TensorProto.FLOAT)
+            cast_node1, y = self._new_node(node_name, "Cast", [y], to=TensorProto.FLOAT)
+            softmax_grad_node, softmax_grad_out = self._new_node(node_name, "SoftmaxGrad_13", [dy, y], axis=axis)
+            cast_node2, _ = self._new_node(node_name, "Cast", [softmax_grad_out], outputs=[dx], to=dtype)
+            return [cast_node, cast_node1, softmax_grad_node, cast_node2]
+        mul_node, mul_out = self._new_node(node_name, "Mul", [dy, y])
+        sum_node, sum_out = self._new_node(node_name, "ReduceSum", [mul_out], axes=[axis])
+        mul_node1, mul_out1 = self._new_node(node_name, "Mul", [y, sum_out])
+        sub_node, _ = self._new_node(node_name, "Sub", [mul_out, mul_out1], outputs=[dx])
+        return [mul_node, sum_node, mul_node1, sub_node]
+
+    # We support to codegen for reduce Ops that are reducing on contiguous axes. If it's not, we need to decompose
+    # it to multiple reduce Ops.
+    def _decompose_reduce_axes(self, node: NodeProto, graph: GraphProto, **kwargs):
+        node_name = node.name
+        op_type = node.op_type
+        x = node.input[0]
+        y = node.output[0]
+        _, shape = self._get_dtype_and_shape(x, **kwargs)
+        rank = len(shape)
+        keep_dims, axes = get_reduce_info(node, graph, rank)
+        if len(axes) == 0:
+            identity_node, _ = self._new_node(node_name, "Identity", [x], outputs=[y])
+            return [identity_node]
+        splited_axes = []
+        end = len(axes)
+        start = end - 1
+        while True:
+            while start > 0 and axes[start] == axes[start - 1] + 1:
+                start -= 1
+            splited_axes.append(axes[start:end])
+            if start == 0:
+                break
+            end = start
+            start = end - 1
+        if len(splited_axes) == 1:
+            if len(node.input) <= 1:
+                return [node]
+            reduce_node, _ = self._new_node(node_name, op_type, [x], outputs=[y], axes=axes, keepdims=keep_dims)
+            return [reduce_node]
+        result = []
+        for idx, axes in enumerate(splited_axes):
+            outputs = [y] if idx == len(splited_axes) - 1 else None
+            reduce_node, x = self._new_node(node_name, op_type, [x], outputs=outputs, axes=axes, keepdims=keep_dims)
+            result.append(reduce_node)
+        return result
+
+    def _decompose_reduce_precision(self, node: NodeProto, graph: GraphProto, **kwargs):
+        x = node.input[0]
+        dtype, shape = self._get_dtype_and_shape(x, **kwargs)
+        if not _is_half_dtype(dtype):
+            return [node]
+        node_name = node.name
+        rank = len(shape)
+        keep_dims, axes = get_reduce_info(node, graph, rank)
+        y = node.output[0]
+        cast_node, x = self._new_node(node_name, "Cast", [x], to=TensorProto.FLOAT)
+        reduce_node, reduce_out = self._new_node(node_name, node.op_type, [x], axes=axes, keepdims=keep_dims)
+        cast_node1, _ = self._new_node(node_name, "Cast", [reduce_out], outputs=[y], to=dtype)
+        return [cast_node, reduce_node, cast_node1]
+
+    def ReduceMax(self, node: NodeProto, graph: GraphProto, **kwargs):  # noqa: N802
+        return self._decompose_reduce_axes(node, graph, **kwargs)
+
+    def ReduceMin(self, node: NodeProto, graph: GraphProto, **kwargs):  # noqa: N802
+        return self._decompose_reduce_axes(node, graph, **kwargs)
+
+    def ReduceSum(self, node: NodeProto, graph: GraphProto, **kwargs):  # noqa: N802
+        precision_decompose_result = self._decompose_reduce_precision(node, graph, **kwargs)
+        # The decompose process will be called recursively, if it's already a decomposed result, just return.
+        if len(precision_decompose_result) != 1 or precision_decompose_result[0] != node:
+            return precision_decompose_result
+        return self._decompose_reduce_axes(node, graph, **kwargs)
+
+    def ReduceMean(self, node: NodeProto, graph: GraphProto, **kwargs):  # noqa: N802
+        precision_decompose_result = self._decompose_reduce_precision(node, graph, **kwargs)
+        # The decompose process will be called recursively, if it's already a decomposed result, just return.
+        if len(precision_decompose_result) != 1 or precision_decompose_result[0] != node:
+            return precision_decompose_result
+        axes_decompose_result = self._decompose_reduce_axes(node, graph, **kwargs)
+        # The decompose process will be called recursively, if it's already a decomposed result, just return.
+        if len(axes_decompose_result) != 1 or axes_decompose_result[0] != node:
+            return axes_decompose_result
+        node_name = node.name
+        x = node.input[0]
+        y = node.output[0]
+        dtype, shape = self._get_dtype_and_shape(x, **kwargs)
+        rank = len(shape)
+        keep_dims, axes = get_reduce_info(node, graph, rank)
+        sum_node, sum_out = self._new_node(node_name, "ReduceSum", [x], axes=axes, keepdims=keep_dims)
+        # If it's not concrete shape, we need add more Ops such as Shape, Gather to get the dim value,
+        # which is not supported yet.
+        assert all(shape[axis].is_number for axis in axes)
+        denominator = int(sympy.prod([shape[axis] for axis in axes]))
+        denominator_tensor = helper.make_tensor(
+            name=f"{node_name}_denominator",
+            dims=(),
+            data_type=dtype,
+            vals=np.array([denominator], dtype=to_numpy_type(dtype)),
+        )
+        denominator_node, denominator_out = self._new_node(node_name, "Constant", [], value=denominator_tensor)
+        div_node, _ = self._new_node(node_name, "Div", [sum_out, denominator_out], outputs=[y])
+        return [sum_node, denominator_node, div_node]

orttraining/orttraining/python/training/ort_triton/_ir.py

@@ -0,0 +1,363 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from abc import abstractmethod
+from collections import defaultdict
+from typing import Dict, List, Optional, Set, Tuple
+
+import numpy as np
+import sympy
+
+from ._common import AutotuneConfigs, CodeBuffer, CodegenContext, NodeVisitor, TensorInfo
+from ._sympy_utils import parse_shape
+from ._utils import gen_unique_name, gen_variable_name, sort_reduce_axes, to_numpy_type
+
+
+class TensorArg:
+    """
+    A TensorArg represents a tensor argument in the kernel function.
+    It contains a name (from ONNX graph), the data type, the shape.
+    If it's constant (initializer or constant node), it also contains the data in numpy array.
+    """
+
+    def __init__(self, name: str, tensor_info: Optional[TensorInfo] = None, data: Optional[np.ndarray] = None):
+        self._name: str = name
+        self._data: Optional[np.ndarray] = data
+        if data is not None:
+            self._dtype: np.dtype = data.dtype
+            self._shape: List[sympy.Expr] = parse_shape(list(data.shape))
+        else:
+            assert tensor_info is not None
+            self._dtype: np.dtype = to_numpy_type(tensor_info.dtype)
+            self._shape: List[sympy.Expr] = tensor_info.shape
+        self.cross_kernels: bool = False
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def dtype(self) -> np.dtype:
+        return self._dtype
+
+    @property
+    def shape(self) -> List[sympy.Expr]:
+        return self._shape
+
+    @property
+    def data(self) -> Optional[np.ndarray]:
+        return self._data
+
+
+class OffsetCalculator:
+    """
+    OffsetCalculator maps tensor arguments to the target shape of a kernel.
+    It' used to generate the offset code for data load/store for a tensor argument in a kernel with
+    specific target shape.
+    If the reduce_axes is not empty, it means the kernel is a reduction kernel, otherwise it's an element-wise kernel.
+    It requires the axes in reduce_axes are contiguous.
+    If a reduce node has non-contiguous axes, need to decompose it into multiple reduce nodes before code-gen.
+    """
+
+    def __init__(self, target_shape: List[sympy.Expr], reduce_axes: List[int]):
+        self.target_shape: List[sympy.Expr] = target_shape
+        self.is_reduction: bool = len(reduce_axes) > 0
+        self.rank = len(target_shape)
+        self.reduce_axes = sort_reduce_axes(reduce_axes, self.rank)
+        self.x_dims: List[sympy.Expr] = [target_shape[dim] for dim in range(self.rank) if dim not in self.reduce_axes]
+        self.x_rank: int = len(self.x_dims)
+        self.x_numel: sympy.Expr = sympy.prod(self.x_dims) if self.x_rank > 0 else sympy.Integer(1)
+        self.r_dims: List[sympy.Expr] = [target_shape[dim] for dim in self.reduce_axes]
+        self.r_rank: int = len(self.r_dims)
+        self.r_numel: sympy.Expr = sympy.prod(self.r_dims) if self.r_rank > 0 else sympy.Integer(1)
+        self.x_strides: List[sympy.Expr] = []
+        if self.x_rank > 0:
+            self.x_strides.append(sympy.Integer(1))
+            for i in range(self.x_rank - 2, -1, -1):
+                self.x_strides.insert(0, self.x_strides[0] * self.x_dims[i + 1])
+        # To avoid generating useless code for offset calculation, we use x_compute_dims and r_compute_dims to
+        # track the dimensions that need to be computed in the offset calculation. These 2 sets will be set in
+        # register_tensor_arg function below.
+        self.x_compute_dims: Set[int] = set()
+        self.r_strides: List[sympy.Expr] = []
+        if self.r_rank > 0:
+            self.r_strides.append(sympy.Integer(1))
+            for i in range(self.r_rank - 2, -1, -1):
+                self.r_strides.insert(0, self.r_strides[0] * self.r_dims[i + 1])
+        self.r_compute_dims: Set[int] = set()
+        self.input_strides: Dict[str, List[sympy.Expr]] = dict()
+        # Support concrete shape only for now.
+        assert self.x_numel.is_integer and self.r_numel.is_integer
+        self.autotune_configs: AutotuneConfigs = AutotuneConfigs(
+            int(self.x_numel), int(self.r_numel), not self.is_reduction or self.reduce_axes[-1] == self.rank - 1
+        )
+        self.requires_x_mask: bool = any(int(self.x_numel) % config[0] != 0 for config in self.autotune_configs.configs)
+        self.requires_r_mask: bool = any(int(self.r_numel) % config[1] != 0 for config in self.autotune_configs.configs)
+        self.reduced_args: Set[str] = set()
+
+    def get_input_strides(self, name: str) -> List[sympy.Expr]:
+        assert name in self.input_strides
+        return self.input_strides[name]
+
+    def get_x_input_strides(self, name: str) -> List[sympy.Expr]:
+        return [dim for idx, dim in enumerate(self.get_input_strides(name)) if idx not in self.reduce_axes]
+
+    def get_r_input_strides(self, name: str) -> List[sympy.Expr]:
+        return [dim for idx, dim in enumerate(self.get_input_strides(name)) if idx in self.reduce_axes]
+
+    # Whether the x shape of the tensor argument is contiguous and is same as the target shape.
+    def is_same_x_shape(self, name: str) -> bool:
+        if (
+            self.is_reduction
+            and self.reduce_axes[0] != 0
+            and self.reduce_axes[-1] != self.rank - 1
+            and not self.is_r_reduced(name)
+        ):
+            return False
+        return all(dim != sympy.Integer(0) for dim in self.get_x_input_strides(name))
+
+    # Whether the r shape of the tensor argument is same as the target shape (r shape dimensions are always contiguous).
+    def is_same_r_shape(self, name: str) -> bool:
+        return all(dim != sympy.Integer(0) for dim in self.get_r_input_strides(name))
+
+    def register_tensor_arg(self, tensor_arg: TensorArg):
+        if tensor_arg.name in self.input_strides:
+            return
+        strides = []
+        input_shape = tensor_arg.shape
+        if tensor_arg.name in self.reduced_args:
+            assert self.is_reduction
+            reduced_rank = len(input_shape) - len(self.reduce_axes)
+            if len(input_shape) < reduced_rank:
+                input_shape = [sympy.Integer(1)] * (reduced_rank - len(input_shape)) + input_shape
+            input_shape = (
+                input_shape[: self.reduce_axes[0]]
+                + ([sympy.Integer(1)] * len(self.reduce_axes))
+                + input_shape[self.reduce_axes[0] :]
+            )
+        elif len(input_shape) < len(self.target_shape):
+            input_shape = [sympy.Integer(1)] * (len(self.target_shape) - len(input_shape)) + input_shape
+        running_stride = sympy.Integer(1)
+        for i in range(len(self.target_shape) - 1, -1, -1):
+            if self.target_shape[i] == input_shape[i]:
+                strides.insert(0, running_stride)
+                running_stride = running_stride * input_shape[i]
+            else:
+                strides.insert(0, sympy.Integer(0))
+        self.input_strides[tensor_arg.name] = strides
+        if not self.is_same_x_shape(tensor_arg.name):
+            for idx, dim in enumerate(self.get_x_input_strides(tensor_arg.name)):
+                if dim != sympy.Integer(0):
+                    self.x_compute_dims.add(idx)
+        if not self.is_same_r_shape(tensor_arg.name):
+            for idx, dim in enumerate(self.get_r_input_strides(tensor_arg.name)):
+                if dim != sympy.Integer(0):
+                    self.r_compute_dims.add(idx)
+
+    def is_x_reduced(self, name: str) -> bool:
+        strides = self.get_input_strides(name)
+        return all(dim == sympy.Integer(0) for idx, dim in enumerate(strides) if idx not in self.reduce_axes)
+
+    def is_r_reduced(self, name: str) -> bool:
+        strides = self.get_input_strides(name)
+        return all(dim == sympy.Integer(0) for idx, dim in enumerate(strides) if idx in self.reduce_axes)
+
+
+class IRNode:
+    """
+    The base class for all IR nodes.
+    """
+
+    def __init__(self, inputs: List[TensorArg], outputs: List[TensorArg]):
+        self.inputs: List[TensorArg] = inputs
+        self.outputs: List[TensorArg] = outputs
+
+    @abstractmethod
+    def codegen(self, visitor: NodeVisitor, context: CodegenContext, code_buffer: CodeBuffer, indent: int = 0):
+        visitor.codegen(self, context, code_buffer, indent)
+
+
+class ComputeNode(IRNode):
+    """
+    Each operator is represented as a ComputeNode.
+    """
+
+    def __init__(self, op_type: str, inputs: List[TensorArg], outputs: List[TensorArg]):
+        super().__init__(inputs, outputs)
+        self._op_type: str = op_type
+
+    @property
+    def op_type(self):
+        return self._op_type
+
+
+class ReduceNode(ComputeNode):
+    def __init__(self, op_type: str, inputs: List[TensorArg], outputs: List[TensorArg], offset_calc: OffsetCalculator):
+        super().__init__(op_type, inputs, outputs)
+        assert op_type == "ReduceSum" or op_type == "ReduceMax" or op_type == "ReduceMin"
+        self.default_value: str = (
+            "0.0" if op_type == "ReduceSum" else ('float("-inf")' if op_type == "ReduceMax" else 'float("inf")')
+        )
+        self.triton_func: str = (
+            "tl.sum" if op_type == "ReduceSum" else ("tl.max" if op_type == "ReduceMax" else "tl.min")
+        )
+        self.offset_calc: OffsetCalculator = offset_calc
+
+
+class ReduceForLoopStart(ComputeNode):
+    """
+    For reduce kernels that need for loop to compute, ReduceForLoopStart and ReduceForLoopEnd are used to
+    represent the start and end of the for loop.
+
+    shared-memory declaration
+    """
+
+    def __init__(self, reduce_nodes: List[ReduceNode], offset_calc: OffsetCalculator):
+        super().__init__("", [], [])
+        self.reduce_nodes: List[ReduceNode] = reduce_nodes
+        self.offset_calc: OffsetCalculator = offset_calc
+
+
+class ReduceForLoopEnd(ComputeNode):
+    """
+    shared-memory reduction
+    """
+
+    def __init__(self, reduce_nodes: List[ReduceNode], offset_calc: OffsetCalculator):
+        super().__init__("", [], [])
+        self.reduce_nodes: List[ReduceNode] = reduce_nodes
+        self.offset_calc: OffsetCalculator = offset_calc
+
+
+class DropoutNode(ComputeNode):
+    """
+    DropoutNode is used to represent the dropout operator. It is special because we need to track the global offset
+    if there are more than one dropout operators in the subgraph.
+    """
+
+    def __init__(self, inputs: List[TensorArg], outputs: List[TensorArg], offset_calc: OffsetCalculator):
+        super().__init__("Dropout", inputs, outputs)
+        self.offset_calc: OffsetCalculator = offset_calc
+        self.offset_calc.register_tensor_arg(inputs[0])
+        self.global_offset: sympy.Expr = sympy.Integer(0)
+
+
+class IONode(IRNode):
+    """
+    The IONode is used to represent the input and output of the subgraph,
+    which is used to generate data load/store code.
+
+    """
+
+    def __init__(self, tensor_arg: TensorArg, offset_calc: OffsetCalculator, is_load: bool):
+        super().__init__([], [])
+        self.tensor_arg: TensorArg = tensor_arg
+        self.is_load: bool = is_load
+        self.offset_calc: OffsetCalculator = offset_calc
+        self.offset_calc.register_tensor_arg(tensor_arg)
+
+
+class KernelNode(IRNode):
+    """
+    The KernelNode is used to represent a single kernel. Each kernel has a unique target shape.
+
+    """
+
+    def __init__(self, inputs: List[TensorArg], outputs: List[TensorArg], target_shape: List, reduce_axes: List[int]):
+        super().__init__(inputs, outputs)
+        self.name: str = gen_unique_name("triton")
+        self.internal_args: Set[str] = set()
+        self.constants: Dict[str, TensorArg] = dict()
+        self.target_shape: List[sympy.Expr] = target_shape
+        self.sub_nodes: List[IRNode] = []
+        self.var_map: Dict[str, str] = dict()
+        self.symbolic_shape_variables: List[str] = []
+        self.has_dropout: bool = False
+        self.offset_calc: OffsetCalculator = OffsetCalculator(target_shape, reduce_axes)
+
+    def gen_variable_names(self):
+        existing_names = set()
+        for input in self.inputs:
+            name = gen_variable_name(input.name, "in", existing_names)
+            self.var_map[input.name] = name
+            self.var_map[name] = "t_" + name
+        for output in self.outputs:
+            name = gen_variable_name(output.name, "out", existing_names)
+            self.var_map[output.name] = name
+            self.var_map[name] = "t_" + name
+        for name in self.internal_args:
+            self.var_map[name] = gen_variable_name(name, "t", existing_names)
+        for constant_name in self.constants:
+            self.var_map[constant_name] = gen_variable_name(constant_name, "c", existing_names)
+            if self.constants[constant_name].data is not None:
+                value = self.constants[constant_name].data
+                if value is not None:
+                    assert value.size == 1, f"unsupported constant array {value}"
+                    variable_name = self.var_map[constant_name]
+                    assert variable_name not in self.var_map
+                    self.var_map[variable_name] = str(np.array(value.item(), value.dtype))
+
+        self.symbolic_shape_variables = [str(dim) for dim in self.target_shape if dim.is_symbol]
+
+
+class ElementwiseKernelNode(KernelNode):
+    def __init__(self, inputs: List[TensorArg], outputs: List[TensorArg], target_shape: List[sympy.Expr]):
+        super().__init__(inputs, outputs, target_shape, [])
+
+
+class ReduceKernelNode(KernelNode):
+    def __init__(
+        self,
+        inputs: List[TensorArg],
+        outputs: List[TensorArg],
+        target_shape: List[sympy.Expr],
+        reduce_axes: List[int],
+        reduced_args: Set[str],
+    ):
+        super().__init__(inputs, outputs, target_shape, reduce_axes)
+        self.offset_calc.reduced_args.update(reduced_args)
+
+
+class ModuleNode(IRNode):
+    """
+    The ModuleNode is used to represent the whole subgraph. It may contain multiple kernels.
+
+    """
+
+    def __init__(
+        self,
+        func_name: str,
+        inputs: List[TensorArg],
+        outputs: List[TensorArg],
+        constants: List[TensorArg],
+        cross_kernel_args: List[Tuple[TensorArg, int]],
+        kernels: List[KernelNode],
+    ):
+        super().__init__(inputs, outputs)
+        self.func_name: str = func_name
+        # Currently need inputs and outputs only. May need intermediate vars and constants later.
+        self.constants: List[TensorArg] = constants
+        self.kernels: List[KernelNode] = kernels
+        self.var_map: Dict[str, str] = dict()
+        existing_names = set()
+        for input in self.inputs:
+            name = gen_variable_name(input.name, "in", existing_names)
+            self.var_map[input.name] = name
+        for output in self.outputs:
+            name = gen_variable_name(output.name, "out", existing_names)
+            self.var_map[output.name] = name
+        self.cross_kernel_args_to_delete: Dict[int, Set[str]] = defaultdict(set)
+        for pair in cross_kernel_args:
+            name = gen_variable_name(pair[0].name, "buf", existing_names)
+            self.cross_kernel_args_to_delete[pair[1]].add(name)
+            self.var_map[pair[0].name] = name
+        running_offset = sympy.Integer(0)
+        self.has_dropout: bool = False
+        for kernel in self.kernels:
+            for ir_node in kernel.sub_nodes:
+                if isinstance(ir_node, DropoutNode):
+                    ir_node.global_offset = running_offset
+                    running_offset = running_offset + sympy.prod(ir_node.outputs[0].shape)
+                    self.has_dropout = True

orttraining/orttraining/python/training/ort_triton/_lowering.py

@@ -0,0 +1,559 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import itertools
+import warnings
+from collections import defaultdict
+from typing import Any, Dict, List, Set, Tuple
+
+import sympy
+from onnx import NodeProto
+
+from ._common import AutotuneConfigs, TensorInfo
+from ._ir import (
+    ComputeNode,
+    DropoutNode,
+    ElementwiseKernelNode,
+    IONode,
+    KernelNode,
+    ModuleNode,
+    OffsetCalculator,
+    ReduceForLoopEnd,
+    ReduceForLoopStart,
+    ReduceKernelNode,
+    ReduceNode,
+    TensorArg,
+)
+from ._op_config import is_reduction_node
+from ._sorted_graph import SortedGraph
+from ._utils import get_reduce_info, sort_reduce_axes, to_numpy_array
+
+
+class NodeGroup:
+    """
+    A NodeGroup contains nodes that can be lowered to a single Triton kernel node.
+
+    """
+
+    def __init__(self, node: NodeProto, reduce_axes: List[int], keep_dims: int, node_arg_infos: Dict[str, TensorInfo]):
+        self._node_arg_infos = node_arg_infos
+        self.nodes_groups: List[Any] = [node]
+        self.target_shape: List[sympy.Expr] = self._get_target_shape(node)
+        rank = len(self.target_shape)
+        self.reduce_axes: List[int] = sort_reduce_axes(reduce_axes, rank)
+        x_dims = [self.target_shape[dim] for dim in range(rank) if dim not in self.reduce_axes]
+        # x_numel is meant to hint how many rows of tensor will be processed by each kernel.
+        # x is same as CUDA block in X direction.
+        x_numel: sympy.Expr = sympy.prod(x_dims) if len(x_dims) > 0 else sympy.Integer(1)
+        r_dims: List[sympy.Expr] = [self.target_shape[dim] for dim in self.reduce_axes]
+        # r_numel is meant to hint how many elements in a row of tensor will be processed by each kernel.
+        # r is a abbreviation of reduction, so, it's only used for reduction nodes.
+        r_numel: sympy.Expr = sympy.prod(r_dims) if len(r_dims) > 0 else sympy.Integer(1)
+        # Support concrete shape only for now.
+        assert x_numel.is_integer and r_numel.is_integer
+        self.autotune_configs: AutotuneConfigs = AutotuneConfigs(
+            int(x_numel), int(r_numel), len(self.reduce_axes) == 0 or self.reduce_axes[-1] == rank - 1
+        )
+        self.reduced_args: Set[str] = set()
+        if keep_dims != 1:
+            self.reduced_args.add(node.output[0])
+
+    # Check if shape can be broadcasted to target_shape.
+    # For example, [1, 3, 1, 1] can be broadcasted to [1, 3, 5, 7].
+    # and we support `keepdims = false``, so [1, 3, 5, 7] is compatible with [1, 3, 5].
+    def _compatible_shape(self, shape: List[sympy.Expr], split_if_different: bool) -> bool:
+        if split_if_different:
+            return shape == self.target_shape
+        if len(shape) > len(self.target_shape):
+            return False
+        shape = [sympy.Integer(1)] * (len(self.target_shape) - len(shape)) + shape
+        for axis in range(len(shape)):
+            if shape[axis] != self.target_shape[axis] and (
+                not shape[axis].is_number or shape[axis] != sympy.Integer(1)
+            ):
+                return False
+        return True
+
+    # Only we consider reduction or elementwise nodes.
+    # target shape does effect how we block the tensor in a triton kernel
+    # for reduction, it's possible to set keepdims=False
+    # for element-wise, output shape is always the target shape.
+    def _get_target_shape(self, node):
+        name = node.input[0] if is_reduction_node(node) else node.output[0]
+        return self._node_arg_infos[name].shape
+
+    # Check if a node can be added to this group.
+    # a group represents a single kernel.
+    # Theoretically, we should fuse as more nodes as possible to benefit most from memory access pattern.
+    # But we have to consider the following factors:
+    #     1. We have to keep the order of nodes, so that we can't fuse nodes that are not adjacent.
+    #     2. The target shape of a group is determined by the first node in the group.
+    #       we call it dominators, and it determinate the partition strategy of X_numel/R_numel.
+    #       A group can't have multiple dominators.
+    def compatible(self, node: NodeProto, reduce_axes: List[int], keep_dims: int, split_if_different: bool) -> bool:
+        target_shape = self._get_target_shape(node)
+        if is_reduction_node(node):
+            # If the following nodes are all elementwise nodes on reduce output shape.
+            if len(self.reduce_axes) == 0 and self.target_shape == self._node_arg_infos[node.output[0]].shape:
+                return True
+            if keep_dims != 1:
+                return False
+            if (
+                len(self.reduce_axes) > 0 and self.reduce_axes != sort_reduce_axes(reduce_axes, len(target_shape))
+            ) or self.target_shape != target_shape:
+                return False
+            return True
+        return self._compatible_shape(target_shape, split_if_different)
+
+    # 1. Create a new group with the reduction node.
+    # 2. Add this node to the current group.
+    def add_node(self, node: NodeProto, reduce_axes: List[int], keep_dims: int):
+        if is_reduction_node(node):
+            group = NodeGroup(node, reduce_axes, keep_dims, self._node_arg_infos)
+            self.nodes_groups.append(group)
+            if len(self.reduce_axes) == 0:
+                self.target_shape = group.target_shape
+                self.reduce_axes = group.reduce_axes
+                self.autotune_configs = group.autotune_configs
+                if keep_dims != 1:
+                    for idx in range(len(self.nodes_groups) - 1):
+                        self.reduced_args.update(self.nodes_groups[idx].input)
+                        self.reduced_args.update(self.nodes_groups[idx].output)
+            return group
+        self.nodes_groups.append(node)
+        return self
+
+    def has_reduced_elementwise_nodes(self) -> bool:
+        return not is_reduction_node(self.nodes_groups[0]) and len(self.reduced_args) > 0
+
+    def dependent_nodes(self, keep_reduce_node: bool):
+        node_map = dict()
+        reduce_nodes = []
+        if not keep_reduce_node and self.has_reduced_elementwise_nodes():
+            for item in self.nodes_groups:
+                if not isinstance(item, NodeGroup):
+                    node_map[item.name] = item
+            return node_map, reduce_nodes
+        for item in self.nodes_groups:
+            if isinstance(item, NodeGroup):
+                node_map.update(item.dependent_nodes(keep_reduce_node)[0])
+            elif keep_reduce_node or not is_reduction_node(item):
+                node_map[item.name] = item
+            else:
+                reduce_nodes.append(item)
+        return node_map, reduce_nodes
+
+    # finalize the group, and return the flatten nodes
+    def flatten(self, sorted_nodes: List[NodeProto]) -> Tuple[List[NodeProto], List[List[int]]]:
+        if self.autotune_configs.requires_for_loop:
+            layers = []
+            group_layer = [self]
+            while len(group_layer) > 0:
+                node_map = dict()
+                reduce_node_map = dict()
+                next_layer = []
+                for group in group_layer:
+                    sub_node_map, reduce_nodes = group.dependent_nodes(False)
+                    node_map.update(sub_node_map)
+                    for node in reduce_nodes:
+                        reduce_node_map[node.name] = node
+                    next_layer.extend([item for item in group.nodes_groups if isinstance(item, NodeGroup)])
+                layers.append((node_map, reduce_node_map.values()))
+                group_layer = next_layer
+            nodes = []
+            layer_indices = []
+            for i in range(len(layers) - 1, -1, -1):
+                sub_nodes = list(layers[i][0].values())
+                sub_nodes.sort(key=sorted_nodes.index)
+                nodes.extend(sub_nodes)
+                sub_layer_indices = []
+                for node in layers[i][1]:
+                    nodes.append(node)
+                    sub_layer_indices.append(len(nodes) - 1)
+                layer_indices.append(sub_layer_indices)
+            return nodes, layer_indices
+        node_map, _ = self.dependent_nodes(True)
+        nodes = list(node_map.values())
+        nodes.sort(key=sorted_nodes.index)
+        return nodes, []
+
+    def try_merge(self, other) -> bool:
+        if (
+            self.target_shape != other.target_shape
+            or self.reduce_axes != other.reduce_axes
+            or self.has_reduced_elementwise_nodes() != other.has_reduced_elementwise_nodes()
+        ):
+            return False
+        self.nodes_groups.extend(other.nodes_groups)
+        self.reduced_args.update(other.reduced_args)
+        return True
+
+
+class KernelIO:
+    """
+    Used to represent the inputs and outputs of a kernel(triton kernel).
+    """
+
+    def __init__(self):
+        self.module_inputs: List[str] = []
+        self.cross_kernel_inputs: List[str] = []
+        self.constants: List[str] = []
+        self.module_outputs: List[str] = []
+        self.cross_kernel_outputs: [str] = []
+        self.internal_args: List[str] = []
+
+
+class GraphLowering:
+    """
+    GraphLowering does manager all steps of lowering onnx graph to triton irnode.
+    1. partition the graph into kernels (one or more kernels).
+        Manager to allocate inputs, outputs and buffers reuse between kernels.
+        we call it Module
+    2. convert kernel to irnodes.
+    3. analyze the IR relationship and buffer/Tensor inside a kernel.
+    4. Generate the Auto-Tune configs for each kernel, Tunning speed/Running faster depends.
+
+    we will end up getting a tree-liked IR structure with explicit input/output and intermediate buffer.
+
+    """
+
+    def __init__(self, sorted_graph: SortedGraph):
+        self._sorted_graph: SortedGraph = sorted_graph
+        self._node_arg_infos: Dict[str, TensorInfo] = sorted_graph.node_arg_infos
+        self._module_inputs: List[TensorArg] = []
+        self._module_outputs: List[TensorArg] = []
+        self._module_constants: List[TensorArg] = []
+        self._module_input_names: Set[str] = set()
+        self._module_output_names: Set[str] = set()
+        self._module_constant_names: Set[str] = set()
+        self._tensor_args: Dict[str, TensorArg] = {}
+        # Extract module inputs, outputs and constants.
+        self._extract_module_io()
+
+        # Group nodes into NodeGroups, each NodeGroup represents a kernel.
+        self._groups: List[NodeGroup] = []
+        self._group_nodes()
+
+        # Convert NodeGroups to KernelNodes.
+        self._kernel_nodes: List[KernelNode] = []
+        self._kernel_io_list: List[KernelIO] = []
+        self._lower()
+
+    # A module is map to a real onnx graph.
+    def _extract_module_io(self):
+        graph = self._sorted_graph.original_graph
+        self._module_inputs = [TensorArg(input.name, self._node_arg_infos[input.name]) for input in graph.input]
+        self._module_input_names = set(arg.name for arg in self._module_inputs)
+        self._module_outputs = [TensorArg(output.name, self._node_arg_infos[output.name]) for output in graph.output]
+        self._module_output_names = set(arg.name for arg in self._module_outputs)
+        for initializer in graph.initializer:
+            data = to_numpy_array(initializer)
+            self._module_constants.append(TensorArg(initializer.name, data=data))
+        for const_node in self._sorted_graph.const_nodes:
+            data = to_numpy_array(const_node)
+            self._module_constants.append(TensorArg(const_node.output[0], data=data))
+        self._module_constant_names = set(arg.name for arg in self._module_constants)
+        self._tensor_args = dict(
+            (arg.name, arg)
+            for arg in itertools.chain(self._module_inputs, self._module_outputs, self._module_constants)
+        )
+
+    def _get_reduce_info(self, node) -> Tuple[int, List[int]]:
+        assert is_reduction_node(node)
+        input_rank = len(self._node_arg_infos[node.input[0]].shape)
+        return get_reduce_info(node, self._sorted_graph.original_graph, input_rank)
+
+    def _process_node(self, node: NodeProto, precessors: Dict[str, List[NodeProto]], group: NodeGroup):
+        dependent_nodes = set()
+        dependent_nodes.add(node.name)
+        for precessor in precessors[node.name]:
+            if precessor.name in dependent_nodes:
+                continue
+            keep_dims = 1
+            reduce_axes = []
+            if is_reduction_node(precessor):
+                keep_dims, reduce_axes = self._get_reduce_info(precessor)
+            split_if_different = any(
+                output in self._sorted_graph.elementwise_graph_outputs for output in precessor.output
+            )
+            if group.compatible(precessor, reduce_axes, keep_dims, split_if_different):
+                next_group = group.add_node(precessor, reduce_axes, keep_dims)
+                dependent_nodes.update(self._process_node(precessor, precessors, next_group))
+        return dependent_nodes
+
+    def _group_nodes(self):
+        producers = dict()
+        precessors = defaultdict(list)
+        processed = set()
+        groups = []
+        sorted_nodes = self._sorted_graph.sorted_nodes
+        for node in sorted_nodes:
+            for output in node.output:
+                producers[output] = node
+            for input in node.input:
+                if input in producers:
+                    precessors[node.name].append(producers[input])
+        for _, value in precessors.items():
+            value.sort(key=sorted_nodes.index, reverse=True)
+        for idx in range(len(sorted_nodes) - 1, -1, -1):
+            node = sorted_nodes[idx]
+            if node.name not in processed:
+                reduce_axes = []
+                keep_dims = 1
+                if is_reduction_node(node):
+                    keep_dims, reduce_axes = self._get_reduce_info(node)
+                groups.append(NodeGroup(node, reduce_axes, keep_dims, self._node_arg_infos))
+                processed.update(self._process_node(node, precessors, groups[-1]))
+
+        # Merge groups with same target shape and reduce axes without dependency.
+        group_dependencies = defaultdict(set)
+        for i in range(len(groups) - 1):
+            group_inputs = set()
+            for node in groups[i].dependent_nodes(True)[0].values():
+                group_inputs.update(node.input)
+            for j in range(i + 1, len(groups)):
+                if any(output in group_inputs for output in groups[j].nodes_groups[0].output):
+                    group_dependencies[i].add(j)
+                    for k in range(0, i):
+                        if i in group_dependencies[k]:
+                            group_dependencies[k].add(j)
+
+        flag = set()
+        for i in range(len(groups)):
+            if i not in flag:
+                for j in range(i + 1, len(groups)):
+                    if j not in flag and j not in group_dependencies[i] and groups[i].try_merge(groups[j]):
+                        flag.add(j)
+                self._groups.append(groups[i])
+                flag.add(i)
+
+    def _get_node_io(self, node: NodeProto) -> Tuple[List[TensorArg], List[TensorArg]]:
+        input_args = []
+        for input in node.input:
+            if input in self._tensor_args:
+                input_args.append(self._tensor_args[input])
+            else:
+                input_args.append(TensorArg(input, self._node_arg_infos[input]))
+                self._tensor_args[input] = input_args[-1]
+        output_args = []
+        for output in node.output:
+            if output in self._tensor_args:
+                output_args.append(self._tensor_args[output])
+            else:
+                output_args.append(TensorArg(output, self._node_arg_infos[output]))
+                self._tensor_args[output] = output_args[-1]
+        return input_args, output_args
+
+    def _extract_kernel_io(self, nodes: List[NodeProto]) -> KernelIO:
+        kernel_io = KernelIO()
+        input_set = set()
+        output_set = set()
+        for node in nodes:
+            for input in node.input:
+                if input in input_set:
+                    continue
+                elif input in self._module_constant_names:
+                    kernel_io.constants.append(input)
+                elif input in self._module_input_names:
+                    kernel_io.module_inputs.append(input)
+                elif input not in output_set:
+                    kernel_io.cross_kernel_inputs.append(input)
+                input_set.add(input)
+            for output in node.output:
+                if output in output_set:
+                    continue
+                if output in self._module_output_names:
+                    kernel_io.module_outputs.append(output)
+                else:
+                    kernel_io.internal_args.append(output)
+                output_set.add(output)
+        return kernel_io
+
+    def _to_compute_node(self, node: NodeProto, offset_calc: OffsetCalculator):
+        inputs, outputs = self._get_node_io(node)
+        op_type = node.op_type
+        if op_type == "Dropout":
+            return DropoutNode(inputs, outputs, offset_calc)
+        if is_reduction_node(node):
+            return ReduceNode(op_type, inputs, outputs, offset_calc)
+        return ComputeNode(op_type, inputs, outputs)
+
+    def _analyze_kernel_io_list(self):
+        cross_kernel_inputs = set()
+        for kernel_io in self._kernel_io_list:
+            cross_kernel_inputs.update(kernel_io.cross_kernel_inputs)
+        for kernel_io in self._kernel_io_list:
+            kernel_io.cross_kernel_outputs = [arg for arg in kernel_io.internal_args if arg in cross_kernel_inputs]
+            kernel_io.internal_args = [
+                arg for arg in kernel_io.internal_args if arg not in kernel_io.cross_kernel_outputs
+            ]
+
+    def _insert_load_and_store(self, kernel_node: KernelNode):
+        input_names = [input.name for input in kernel_node.inputs]
+        output_name_map = dict()
+        for output in kernel_node.outputs:
+            output_name_map[output.name] = 0
+        for node in kernel_node.sub_nodes:
+            for output in node.outputs:
+                if output.name in output_name_map:
+                    output_name_map[output.name] += 1
+        sub_nodes = kernel_node.sub_nodes
+        new_sub_nodes = []
+        cur = 0
+        nxt = 0
+        reduce_store_nodes = []
+        while True:
+            while nxt < len(sub_nodes) and not isinstance(sub_nodes[nxt], ReduceForLoopEnd):
+                nxt += 1
+            load_cache = set()
+            load_nodes = []
+            store_nodes = []
+            for idx in range(cur, nxt):
+                for input in sub_nodes[idx].inputs:
+                    if input.name in kernel_node.constants or input.name in input_names:
+                        if (input.data is not None and input.data.size == 1) or input.name in load_cache:
+                            continue
+                        load_nodes.append(IONode(input, kernel_node.offset_calc, True))
+                        load_cache.add(input.name)
+                for output in sub_nodes[idx].outputs:
+                    if output.name in output_name_map:
+                        output_name_map[output.name] -= 1
+                        if output_name_map[output.name] == 0:
+                            store_nodes.append(IONode(output, kernel_node.offset_calc, False))
+            if isinstance(sub_nodes[cur], ReduceForLoopStart):
+                new_sub_nodes.append(sub_nodes[cur])
+                cur += 1
+            if nxt < len(sub_nodes):
+                assert isinstance(sub_nodes[nxt], ReduceForLoopEnd)
+                for reduce_node in sub_nodes[nxt].reduce_nodes:
+                    input = reduce_node.inputs[0]
+                    if input.name in kernel_node.constants or input.name in input_names:
+                        if (input.data is not None and input.data.size == 1) or input.name in load_cache:
+                            continue
+                        load_nodes.append(IONode(input, kernel_node.offset_calc, True))
+                        load_cache.add(input.name)
+            new_sub_nodes.extend(load_nodes)
+            new_sub_nodes.extend(sub_nodes[cur:nxt])
+            new_sub_nodes.extend(store_nodes)
+            if nxt < len(sub_nodes):
+                assert isinstance(sub_nodes[nxt], ReduceForLoopEnd)
+                for reduce_node in sub_nodes[nxt].reduce_nodes:
+                    if reduce_node.outputs[0].name in output_name_map:
+                        reduce_store_nodes.append(IONode(reduce_node.outputs[0], kernel_node.offset_calc, False))
+                new_sub_nodes.append(sub_nodes[nxt])
+                nxt += 1
+            cur = nxt
+            if cur >= len(sub_nodes):
+                break
+        new_sub_nodes.extend(reduce_store_nodes)
+        kernel_node.sub_nodes = new_sub_nodes
+
+    def _lower(self):
+        for group in self._groups:
+            is_reduction_kernel = len(group.reduce_axes) > 0
+            target_shape = group.target_shape
+            # The inputs and outputs will be initialized later.
+            kernel_node = (
+                ReduceKernelNode([], [], target_shape, group.reduce_axes, group.reduced_args)
+                if is_reduction_kernel
+                else ElementwiseKernelNode([], [], target_shape)
+            )
+            self._kernel_nodes.append(kernel_node)
+            sub_nodes = []
+            nodes, layer_indices = group.flatten(self._sorted_graph.sorted_nodes)
+            self._kernel_io_list.append(self._extract_kernel_io(nodes))
+            if group.autotune_configs.requires_for_loop:
+                start = 0
+                for layer_idx, indices in enumerate(layer_indices):
+                    need_for_loop = True
+                    if layer_idx == len(layer_indices) - 1 and group.has_reduced_elementwise_nodes():
+                        assert len(indices) == 0
+                        need_for_loop = False
+                    reduce_nodes = [self._to_compute_node(nodes[idx], kernel_node.offset_calc) for idx in indices]
+                    assert all(isinstance(node, ReduceNode) for node in reduce_nodes)
+                    if need_for_loop:
+                        sub_nodes.append(ReduceForLoopStart(reduce_nodes, kernel_node.offset_calc))
+                    end = indices[0] if len(indices) > 0 else len(nodes)
+                    for idx in range(start, end):
+                        node = nodes[idx]
+                        assert not is_reduction_node(node)
+                        sub_nodes.append(self._to_compute_node(node, kernel_node.offset_calc))
+                        if node.op_type == "Dropout":
+                            self._kernel_nodes[-1].has_dropout = True
+                    if len(indices) > 0:
+                        sub_nodes.append(ReduceForLoopEnd(reduce_nodes, kernel_node.offset_calc))
+                    start = indices[len(indices) - 1] + 1 if len(indices) > 0 else len(nodes)
+            else:
+                for node in nodes:
+                    sub_nodes.append(self._to_compute_node(node, kernel_node.offset_calc))
+                    if node.op_type == "Dropout":
+                        self._kernel_nodes[-1].has_dropout = True
+            self._kernel_nodes[-1].sub_nodes = sub_nodes
+
+        if any(kernel_node.has_dropout for kernel_node in self._kernel_nodes):
+            warnings.warn("Use triton's random for Dropout, ignore the random seed from ORT.", UserWarning)
+
+        self._analyze_kernel_io_list()
+        cross_kernel_arg_map = dict()
+        for idx, kernel_io in enumerate(self._kernel_io_list):
+            for output in itertools.chain(kernel_io.cross_kernel_outputs, kernel_io.module_outputs):
+                cross_kernel_arg_map[output] = idx
+        dependency = defaultdict(set)
+        for idx, kernel_io in enumerate(self._kernel_io_list):
+            for input in kernel_io.cross_kernel_inputs:
+                dependency[cross_kernel_arg_map[input]].add(idx)
+        visited = set()
+        sorted_indices = []
+
+        def _topological_sort_internal(idx):
+            visited.add(idx)
+            for next_idx in dependency[idx]:
+                if next_idx not in visited:
+                    _topological_sort_internal(next_idx)
+            sorted_indices.insert(0, idx)
+
+        for idx in range(len(self._kernel_io_list)):
+            if idx not in visited:
+                _topological_sort_internal(idx)
+
+        self._kernel_nodes = [self._kernel_nodes[idx] for idx in sorted_indices]
+        self._kernel_io_list = [self._kernel_io_list[idx] for idx in sorted_indices]
+        cross_kernel_arg_map.clear()
+        for idx, kernel_io in enumerate(self._kernel_io_list):
+            for arg in kernel_io.cross_kernel_inputs:
+                if arg not in self._module_output_names:
+                    cross_kernel_arg_map[arg] = idx
+
+        self._cross_kernel_args = [(self._tensor_args[key], value) for key, value in cross_kernel_arg_map.items()]
+
+        for idx, kernel_node in enumerate(self._kernel_nodes):
+            kernel_io = self._kernel_io_list[idx]
+            kernel_node.internal_args.update(kernel_io.internal_args)
+            kernel_node.inputs = [
+                self._tensor_args[name]
+                for name in itertools.chain(kernel_io.module_inputs, kernel_io.cross_kernel_inputs)
+            ]
+            kernel_node.outputs = [
+                self._tensor_args[name]
+                for name in itertools.chain(kernel_io.module_outputs, kernel_io.cross_kernel_outputs)
+            ]
+            for name in kernel_io.constants:
+                kernel_node.constants[name] = self._tensor_args[name]
+            self._insert_load_and_store(kernel_node)
+            kernel_node.gen_variable_names()
+
+    def module_node(self, func_name: str):
+        return ModuleNode(
+            func_name,
+            self._module_inputs,
+            self._module_outputs,
+            self._module_constants,
+            self._cross_kernel_args,
+            self._kernel_nodes,
+        )
+
+
+def lower(func_name: str, sorted_graph: SortedGraph) -> ModuleNode:
+    return GraphLowering(sorted_graph).module_node(func_name)

orttraining/orttraining/python/training/ort_triton/_op_config.py

@@ -0,0 +1,63 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""
+This file controls the Ops supported by Triton codegen.
+The Triton fusion on backend will try to extract subgraphs from the ONNX model which contain connected supported ops.
+For each supported op, it has the following attributes:
+  - domain: The domain of the op. If not specified, the op is from ONNX domain, which is empty string.
+  - versions: The supported opset versions.
+  - is_no_op: If the op is no-op. If not specified, the op is not no-op.
+  - conditions: define some conditions to control the fusion. For example, for Ops that has axis attribute,
+                a condition "axis": "-1" means only Ops that reduce on last dimension will be fused.
+                For Reduce* Ops, the "axes" condition can be list of ints, such as "axes": "[-1]",
+                or "axes": "single", means only support Ops reduce on single constant dimension,
+                or "axes": "constant", means the axes attribute or input must be constant.
+  - ignore_min_nodes: by default is False. If set to True, the fusion will ignore the min_nodes check for this Op.
+                      For example, if the min_nodes is 2, the fusion will only fuse the subgraphs with 2 or more
+                      non-no-op nodes. But if the ignore_min_nodes is True for ReduceSum, it's OK to fuse a single
+                      ReduceSum node to the subgraph.
+"""
+
+from onnx import NodeProto
+
+_ELEMENTWISE_OPS = {
+    "Add": {"versions": [13, 14]},
+    "Sub": {"versions": [13, 14]},
+    "Mul": {"versions": [13, 14]},
+    "Div": {"versions": [13, 14]},
+    "Pow": {"versions": [13, 15]},
+    "Sqrt": {"versions": [13]},
+    "Exp": {"versions": [13]},
+    "Where": {"versions": [9, 16]},
+    "Cast": {"versions": [13]},
+    "Dropout": {"versions": [13]},
+    "DropoutGrad": {"domain": "com.microsoft", "versions": [1]},
+    "Identity": {"versions": [13], "is_no_op": True},
+    "Sum": {"versions": [13]},
+}
+
+_REDUCTION_OPS = {
+    "ReduceMean": {"versions": [13], "conditions": {"axes": "[-1]"}},
+    "ReduceSum": {"versions": [13], "conditions": {"axes": "[-1]"}},
+    "ReduceMax": {"versions": [13], "conditions": {"axes": "[-1]"}},
+    "ReduceMin": {"versions": [13], "conditions": {"axes": "[-1]"}},
+    "Softmax": {"versions": [13]},
+    "SoftmaxGrad_13": {"domain": "com.microsoft", "versions": [1]},
+    # Disable LayerNormalization fusion for now as it's generated Triton code is inefficient compared to C++ kernel.
+    # "LayerNormalization": {"versions": [1]},
+    # "LayerNormalizationGrad": {"domain": "com.microsoft", "versions": [1]},
+}
+
+
+def is_elementwise_node(node: NodeProto) -> bool:
+    return node.op_type in _ELEMENTWISE_OPS
+
+
+def is_reduction_node(node: NodeProto) -> bool:
+    return node.op_type in _REDUCTION_OPS
+
+
+def get_supported_ops():
+    return {**_ELEMENTWISE_OPS, **_REDUCTION_OPS}

orttraining/orttraining/python/training/ort_triton/_sorted_graph.py

@@ -0,0 +1,210 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import copy
+import itertools
+from typing import Any, Dict, List, Set
+
+import numpy as np
+import onnx
+from onnx import GraphProto, ModelProto, NodeProto, helper
+
+from ._common import TensorInfo, TypeAndShapeInfer
+from ._decompose import DecomposeDispatch
+from ._op_config import is_elementwise_node
+from ._utils import get_attribute, to_numpy_array, topological_sort
+
+
+class SortedGraph:
+    """
+    This class is used to
+        1. decompose complex operators into preliminary operators,
+        2. sort the operators in topological order,
+        3. infer the type and shape of each node inputs and outputs.
+
+    input args:
+        model: the ONNX model.
+        input_shapes: the shapes of the model inputs. Can be numeric values or symbolic values.
+    """
+
+    def __init__(self, model: ModelProto, input_shapes: List[List[Any]]):
+        self._model: ModelProto = model
+        self._graph: GraphProto = model.graph
+        self._input_shapes: List[List[Any]] = input_shapes
+
+        # For elementwise graph outputs, when we group nodes to different kernels, if the target shape is different
+        # from other nodes' target shape, even it can be broadcasted, we still need to create a new kernel for it.
+        self._elementwise_graph_outputs: Set[str] = set()
+        for node in self._graph.node:
+            if is_elementwise_node(node):
+                self._elementwise_graph_outputs.update(node.output)
+
+        # Topological sort the nodes in the graph.
+        self._sorted_nodes: List[NodeProto] = topological_sort(
+            [input.name for input in self._graph.input] + [initializer.name for initializer in self._graph.initializer],
+            self._graph.node,
+        )
+
+        self._node_arg_infos: Dict[str, TensorInfo] = {}
+        for idx, input in enumerate(self._graph.input):
+            self._node_arg_infos[input.name] = TensorInfo(input.type.tensor_type.elem_type, self._input_shapes[idx])
+        for initializer in self._graph.initializer:
+            self._node_arg_infos[initializer.name] = TensorInfo(
+                initializer.data_type,
+                list(to_numpy_array(initializer).shape),
+            )
+
+        # Decompose complex operators.
+        self._decompose()
+
+        # Sort the initializers in reference order.
+        # We try to reuse Triton module for different ONNX models with same graph structure,
+        # even the node args names in the models are different.
+        # Sorting the initializers can help to generate same model key for different ONNX models.
+        initializers = {}
+        for initializer in self._graph.initializer:
+            initializers[initializer.name] = initializer
+        self._sorted_initializers: List[TensorInfo] = []
+        for node in self._sorted_nodes:
+            for input in node.input:
+                if input in initializers:
+                    self._sorted_initializers.append(initializers[input])
+                    initializers.pop(input)
+
+        # Split nodes to constant nodes and non-constant nodes.
+        self._const_nodes: List[NodeProto] = [node for node in self._sorted_nodes if node.op_type == "Constant"]
+        self._sorted_nodes: List[NodeProto] = [node for node in self._sorted_nodes if node.op_type != "Constant"]
+
+    def __str__(self):
+        """
+        Generate a unique key for the model based on the graph structure, ignoring the node args names.
+        We try to reuse Triton module for different ONNX models with same graph structure.
+        """
+        graph_inputs = []
+        name_map = {}
+        for idx, input in enumerate(self._graph.input):
+            shape_str = str(self._input_shapes[idx]).replace(" ", "")
+            graph_inputs.append(f"({str(input.type.tensor_type.elem_type)},{shape_str})")
+            name_map[input.name] = f"i{idx}"
+        graph_inputs_str = ",".join(graph_inputs)
+
+        constants = []
+        for idx, initializer in enumerate(self._sorted_initializers):
+            data_str = np.array2string(to_numpy_array(initializer), separator=",").replace("\n", "").replace(" ", "")
+            constants.append(f"({initializer.data_type},{data_str})")
+            name_map[initializer.name] = f"c{idx}"
+
+        for idx, node in enumerate(self._const_nodes):
+            value_attr = get_attribute(node, "value")
+            data_str = np.array2string(to_numpy_array(value_attr), separator=",").replace("\n", "").replace(" ", "")
+            constants.append(f"({value_attr.data_type},{data_str})")
+            name_map[node.output[0]] = f"c{idx + len(self._sorted_initializers)}"
+        constants_str = ",".join(constants)
+
+        for idx, output in enumerate(self._graph.output):
+            name_map[output.name] = f"o{idx}"
+
+        nodes = []
+        for node_idx, node in enumerate(self._sorted_nodes):
+            inputs = []
+            for input in node.input:
+                inputs.append(name_map.get(input, input))
+            inputs_str = ",".join(inputs)
+            outputs = []
+            for idx, output in enumerate(node.output):
+                if output in name_map:
+                    outputs.append(name_map[output])
+                else:
+                    name_map[output] = f"t{node_idx}_{idx}"
+                    outputs.append(name_map[output])
+            outputs_str = ",".join(outputs)
+            attributes = []
+            for attr in node.attribute:
+                fields = [str(f[1]) for f in attr.ListFields()]
+                attributes.append(f"{fields[0]}:{fields[2]}={fields[1]}")
+            attributes_str = ",".join(attributes)
+            nodes.append(f"{node.op_type}[{attributes_str}]({inputs_str})->({outputs_str})")
+        nodes_str = ",".join(nodes)
+        return f"{graph_inputs_str}|{str(len(self._graph.output))}|{constants_str}|{nodes_str}"
+
+    def __hash__(self):
+        return hash(str(self))
+
+    def __eq__(self, other):
+        return str(self) == str(other)
+
+    @property
+    def const_nodes(self) -> List[NodeProto]:
+        return self._const_nodes
+
+    @property
+    def sorted_nodes(self) -> List[NodeProto]:
+        return self._sorted_nodes
+
+    @property
+    def original_graph(self) -> GraphProto:
+        return self._graph
+
+    @property
+    def node_arg_infos(self) -> Dict[str, TensorInfo]:
+        return self._node_arg_infos
+
+    @property
+    def elementwise_graph_outputs(self) -> Set[str]:
+        return self._elementwise_graph_outputs
+
+    def _decompose(self):
+        dispatch = DecomposeDispatch()
+        pos = 0
+        # If a node is complex, decompose it and insert the decomposed nodes at the same position.
+        # All complex Ops are defined in DecomposeDispatch.
+        # It's possible that the decomposed nodes are also complex, so we need to do the decompose recursively.
+        # For example, decomposed nodes for "Softmax" contains "ReduceMean",
+        # which will be decomposed to "ReduceSum" and "Div" further.
+        while pos < len(self._sorted_nodes):
+            node = self._sorted_nodes[pos]
+            if node in dispatch:
+                new_nodes = dispatch(node, self._graph, node_arg_infos=self._node_arg_infos)
+                if len(new_nodes) != 1 or new_nodes[0] != node:
+                    new_nodes = topological_sort(node.input, new_nodes)
+                    self._sorted_nodes[pos : pos + 1] = new_nodes
+                    continue
+            if node.op_type == "Constant":
+                value_attr = get_attribute(node, "value")
+                self._node_arg_infos[node.output[0]] = TensorInfo(
+                    value_attr.data_type,
+                    list(to_numpy_array(value_attr).shape),
+                )
+            else:
+                input_infos = []
+                for input in node.input:
+                    input_infos.append(self._node_arg_infos[input])
+                output_infos = TypeAndShapeInfer.infer(node, input_infos, self._graph)
+                for idx, output in enumerate(node.output):
+                    self._node_arg_infos[output] = output_infos[idx]
+            pos += 1
+
+    # Save the ONNX graphs for debug purpose. The original ONNX graph is the subgraph from backend.
+    # The processed ONNX graph is the subgraph after decompose, it also contains the concrete shapes for each arg.
+    def save_onnx(self, file_path_prefix):
+        onnx.save(self._model, file_path_prefix + "_original.onnx")
+        processed_model = copy.deepcopy(self._model)
+        processed_model.graph.ClearField("node")
+        processed_model.graph.node.extend(self.const_nodes)
+        processed_model.graph.node.extend(self.sorted_nodes)
+        for node in itertools.chain(processed_model.graph.input, processed_model.graph.output):
+            node.type.tensor_type.shape.Clear()
+            for dim in self.node_arg_infos[node.name].shape:
+                node.type.tensor_type.shape.dim.add().dim_value = int(dim)
+        value_infos = []
+        for node in itertools.chain(self.const_nodes, self.sorted_nodes):
+            for output in node.output:
+                tensor_info = self.node_arg_infos[output]
+                value_infos.append(
+                    helper.make_tensor_value_info(output, tensor_info.dtype, [int(dim) for dim in tensor_info.shape])
+                )
+        processed_model.graph.ClearField("value_info")
+        processed_model.graph.value_info.extend(value_infos)
+        onnx.save(processed_model, file_path_prefix + "_processed.onnx")

orttraining/orttraining/python/training/ort_triton/_sympy_utils.py

@@ -0,0 +1,26 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import re
+from typing import Any, List
+
+import sympy
+
+
+def sympy_dot(seq1: List[sympy.Expr], seq2: List[sympy.Expr]) -> sympy.Expr:
+    assert len(seq1) == len(seq2)
+    return sympy.expand(sum(a * b for a, b in zip(seq1, seq2)))
+
+
+def parse_shape(shape: List[Any]) -> List[sympy.Expr]:
+    symbol_shapes = []
+    for dim in shape:
+        symbol_dim = dim
+        if isinstance(dim, str):
+            symbol_dim = sympy.Symbol(re.sub(r"[^a-zA-Z0-9_]+", "_", dim))
+        elif isinstance(dim, int):
+            symbol_dim = sympy.Integer(dim)
+        symbol_shapes.append(symbol_dim)
+    return symbol_shapes

orttraining/orttraining/python/training/ort_triton/_utils.py

@@ -0,0 +1,152 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import re
+import uuid
+from collections import defaultdict
+from typing import Any, List, Tuple
+
+import numpy as np
+from onnx import GraphProto, NodeProto, TensorProto, helper, numpy_helper
+from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE
+
+
+def gen_unique_name(prefix: str) -> str:
+    return prefix + "_" + uuid.uuid4().hex[:8]
+
+
+def _topological_sort_internal(node, visited, output_consumers, sorted_nodes):
+    visited.add(node.name)
+    for next_node in output_consumers[node.name]:
+        if next_node.name not in visited:
+            _topological_sort_internal(next_node, visited, output_consumers, sorted_nodes)
+
+    sorted_nodes.insert(0, node)
+
+
+# Topological sort of nodes given the input names. The list of nodes contain both constant and non-constant nodes.
+def topological_sort(inputs: List[str], nodes: List[NodeProto]) -> List[NodeProto]:
+    const_nodes = []
+    non_const_nodes = []
+    for node in nodes:
+        if not node.name:
+            node.name = gen_unique_name(node.op_type)
+        if node.op_type == "Constant":
+            inputs.append(node.output[0])
+            const_nodes.append(node)
+        else:
+            non_const_nodes.append(node)
+
+    # Build relationship between nodes.
+    graph_input_consumers = defaultdict(list)
+    output_consumers = defaultdict(list)
+    input_set = set(inputs)
+    for node in non_const_nodes:
+        for input in node.input:
+            if input in input_set:
+                graph_input_consumers[input].append(node)
+        for output in node.output:
+            if not output:
+                continue
+            for consumer in non_const_nodes:
+                if output in consumer.input:
+                    output_consumers[node.name].append(consumer)
+
+    # Topological sort.
+    visited = set()
+    sorted_nodes = []
+    for input in inputs:
+        for node in graph_input_consumers[input]:
+            if node.name not in visited:
+                _topological_sort_internal(node, visited, output_consumers, sorted_nodes)
+
+    return const_nodes + sorted_nodes
+
+
+# Get attribute value from node by attribute key.
+def get_attribute(node: NodeProto, attr_name: str, default_value: Any = None) -> Any:
+    found = [attr for attr in node.attribute if attr.name == attr_name]
+    if found:
+        return helper.get_attribute_value(found[0])
+    return default_value
+
+
+# Convert Constant node or TensorProto to numpy array.
+def to_numpy_array(node: Any) -> np.ndarray:
+    tensor = node
+    if isinstance(node, NodeProto):
+        tensor = get_attribute(node, "value")
+    assert isinstance(tensor, TensorProto)
+    return numpy_helper.to_array(tensor)
+
+
+def to_numpy_type(tensor_type: TensorProto.DataType) -> np.dtype:
+    return TENSOR_TYPE_TO_NP_TYPE[tensor_type] if not isinstance(tensor_type, np.dtype) else tensor_type
+
+
+# Generate a unique variable name based on the node arg name.
+def gen_variable_name(name: str, prefix: str, existing_names: set) -> str:
+    pos = name.rfind("/")
+    if pos != -1:
+        name = name[pos + 1 :]
+    pos = name.rfind(".")
+    if pos != -1:
+        name = name[pos + 1 :]
+    name = re.sub(r"[^a-zA-Z0-9]", "_", name)
+    if len(name) > 20:
+        name = name[-20:]
+
+    name = f"{prefix}_{name}"
+    while name in existing_names:
+        name = name + "_1"
+
+    existing_names.add(name)
+    return name
+
+
+def may_add_brackets(name: str) -> str:
+    if not re.match("^[A-Za-z0-9_.]*$", name):
+        return f"({name})"
+    return name
+
+
+def sort_reduce_axes(axes: List[int], rank: int, check_contiguous: bool = True) -> List[int]:
+    axes = [axis + rank if axis < 0 else axis for axis in axes]
+    axes.sort()
+    if check_contiguous:
+        for i in range(1, len(axes)):
+            assert axes[i] == axes[i - 1] + 1
+    return axes
+
+
+# Get the keep_dims attribute and reduce axes from a reduce node.
+def get_reduce_info(node: NodeProto, graph: GraphProto, input_rank: int) -> Tuple[int, List[int]]:
+    keep_dims = get_attribute(node, "keepdims", 1)
+    noop_with_empty_axes = get_attribute(node, "noop_with_empty_axes", 0)
+    axes = get_attribute(node, "axes", None)
+    if axes is None and len(node.input) > 1:
+        axes_initializer = None
+        for initializer in graph.initializer:
+            if initializer.name == node.input[1]:
+                axes_initializer = initializer
+                break
+        assert axes_initializer is not None
+        axes = to_numpy_array(axes_initializer).tolist()
+    if axes is None:
+        axes = list(range(input_rank)) if noop_with_empty_axes == 0 else []
+    axes = sort_reduce_axes(axes, input_rank, check_contiguous=False)
+    return keep_dims, axes
+
+
+def next_power_of_2(n: int) -> int:
+    assert n <= 2**32, "32-bit only"
+    n -= 1
+    n |= n >> 1
+    n |= n >> 2
+    n |= n >> 4
+    n |= n >> 8
+    n |= n >> 16
+    n += 1
+    return n

orttraining/orttraining/python/training/ort_triton/kernel/__init__.py

@@ -0,0 +1,17 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from ._mm import triton_gemm, triton_gemm_out, triton_matmul, triton_matmul_out
+from ._slice_scel import slice_scel, slice_scel_backward, transform_slice_scel
+
+__all__ = [
+    "triton_gemm",
+    "triton_gemm_out",
+    "triton_matmul",
+    "triton_matmul_out",
+    "slice_scel",
+    "slice_scel_backward",
+    "transform_slice_scel",
+]

orttraining/orttraining/python/training/ort_triton/kernel/_mm.py

@@ -0,0 +1,484 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import math
+import os
+from types import ModuleType
+from typing import Tuple
+
+import torch
+
+from .._cache import ModuleCache, PyCodeCache
+from .._utils import next_power_of_2
+
+_DEBUG_MODE = "ORTMODULE_TRITON_DEBUG" in os.environ and int(os.getenv("ORTMODULE_TRITON_DEBUG")) == 1
+
+
+_MM_TEMPLATE = """
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.autotune(
+    configs=[
+{autotune_configs}
+    ],
+    key=["M", "N", "K"],
+)
+@triton.jit
+def {kernel_name}(
+    A, B, OUT, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr
+):
+    M = {M}
+    N = {N}
+    K = {K}
+    stride_am = {stride_am}
+    stride_ak = {stride_ak}
+    stride_bk = {stride_bk}
+    stride_bn = {stride_bn}
+
+    # based on triton.ops.matmul
+    pid = tl.program_id(0)
+    grid_m = (M + BLOCK_M - 1) // BLOCK_M
+    grid_n = (N + BLOCK_N - 1) // BLOCK_N
+
+    # re-order program ID for better L2 performance
+    width = GROUP_M * grid_n
+    group_id = pid // width
+    group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
+    pid_m = group_id * GROUP_M + (pid % group_size)
+    pid_n = (pid % width) // (group_size)
+
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
+    rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
+    rk = tl.arange(0, BLOCK_K)
+    A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak)
+    B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn)
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(K, 0, -BLOCK_K):
+        if {even_k}:
+            a = tl.load(A)
+            b = tl.load(B)
+        else:
+            a = tl.load(A, mask=rk[None, :] < k, other=0.)
+            b = tl.load(B, mask=rk[:, None] < k, other=0.)
+        acc += tl.dot(a, b, allow_tf32={allow_tf32})
+        A += BLOCK_K * stride_ak
+        B += BLOCK_K * stride_bk
+
+    # rematerialize rm and rn to save registers
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    idx_m = rm[:, None]
+    idx_n = rn[None, :]
+    mask = (idx_m < M) & (idx_n < N)
+    OUT = OUT + (idx_m * N + idx_n)
+    tl.store(OUT, {post_process}, mask=mask)
+
+
+def {func_name}(a, b, out):
+    # 1D launch kernel where each block gets its own program.
+    grid = lambda META: (triton.cdiv({M}, META["BLOCK_M"]) * triton.cdiv({N}, META["BLOCK_N"]),)
+    {kernel_name}[grid](a, b, out, {M}, {N}, {K})
+"""
+
+
+_BMM_TEMPLATE = """
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.autotune(
+    configs=[
+{autotune_configs}
+    ],
+    key=["M", "N", "K"],
+)
+@triton.jit
+def {kernel_name}(
+    A, B, OUT, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr
+):
+    M = {M}
+    N = {N}
+    K = {K}
+    stride_aq = {stride_aq}
+    stride_am = {stride_am}
+    stride_ak = {stride_ak}
+    stride_bq = {stride_bq}
+    stride_bk = {stride_bk}
+    stride_bn = {stride_bn}
+    stride_cq = M * N
+
+    # based on triton.ops.matmul
+    pid = tl.program_id(0)
+    grid_m = (M + BLOCK_M - 1) // BLOCK_M
+    grid_n = (N + BLOCK_N - 1) // BLOCK_N
+
+    # re-order program ID for better L2 performance
+    width = GROUP_M * grid_n
+    group_id = pid // width
+    group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
+    pid_m = group_id * GROUP_M + (pid % group_size)
+    pid_n = (pid % width) // (group_size)
+
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
+    rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
+    rk = tl.arange(0, BLOCK_K)
+
+    idx_q = tl.program_id(1)  # batch dimension for BMM
+    A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak + idx_q * stride_aq)
+    B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn + idx_q * stride_bq)
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(K, 0, -BLOCK_K):
+        if {even_k}:
+            a = tl.load(A)
+            b = tl.load(B)
+        else:
+            a = tl.load(A, mask=rk[None, :] < k, other=0.)
+            b = tl.load(B, mask=rk[:, None] < k, other=0.)
+        acc += tl.dot(a, b, allow_tf32={allow_tf32})
+        A += BLOCK_K * stride_ak
+        B += BLOCK_K * stride_bk
+
+    # rematerialize rm and rn to save registers
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    idx_q = tl.program_id(1)  # batch dimension for BMM
+    idx_m = rm[:, None]
+    idx_n = rn[None, :]
+    mask = (idx_m < M) & (idx_n < N)
+    OUT = OUT + (idx_m * N + idx_n + idx_q * stride_cq)
+    tl.store(OUT, {post_process}, mask=mask)
+
+
+def {func_name}(a, b, out):
+    # 1D launch kernel where each block gets its own program.
+    grid = lambda META: (triton.cdiv({M}, META["BLOCK_M"]) * triton.cdiv({N}, META["BLOCK_N"]), {batch}, 1)
+    {kernel_name}[grid](a, b, out, {M}, {N}, {K})
+"""
+
+
+_GEMM_TEMPLATE = """
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.autotune(
+    configs=[
+{autotune_configs}
+    ],
+    key=["M", "N", "K"],
+)
+@triton.jit
+def {kernel_name}(
+    A, B, C, OUT, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr
+):
+    M = {M}
+    N = {N}
+    K = {K}
+    stride_am = {stride_am}
+    stride_ak = {stride_ak}
+    stride_bk = {stride_bk}
+    stride_bn = {stride_bn}
+    stride_cm = {stride_cm}
+    stride_cn = {stride_cn}
+
+    # based on triton.ops.matmul
+    pid = tl.program_id(0)
+    grid_m = (M + BLOCK_M - 1) // BLOCK_M
+    grid_n = (N + BLOCK_N - 1) // BLOCK_N
+
+    # re-order program ID for better L2 performance
+    width = GROUP_M * grid_n
+    group_id = pid // width
+    group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
+    pid_m = group_id * GROUP_M + (pid % group_size)
+    pid_n = (pid % width) // (group_size)
+
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
+    rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
+    rk = tl.arange(0, BLOCK_K)
+    A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak)
+    B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn)
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(K, 0, -BLOCK_K):
+        if {even_k}:
+            a = tl.load(A)
+            b = tl.load(B)
+        else:
+            a = tl.load(A, mask=rk[None, :] < k, other=0.)
+            b = tl.load(B, mask=rk[:, None] < k, other=0.)
+        acc += tl.dot(a, b, allow_tf32={allow_tf32})
+        A += BLOCK_K * stride_ak
+        B += BLOCK_K * stride_bk
+
+    # rematerialize rm and rn to save registers
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    idx_m = rm[:, None]
+    idx_n = rn[None, :]
+    mask = (idx_m < M) & (idx_n < N)
+    C = C + (idx_m * stride_cm + idx_n * stride_cn)
+    c = tl.load(C, mask=mask, other=0.)
+    OUT = OUT + (idx_m * N + idx_n)
+    tl.store(OUT, {post_process} + c * {beta}, mask=mask)
+
+
+def {func_name}(a, b, c, out):
+    # 1D launch kernel where each block gets its own program.
+    grid = lambda META: (triton.cdiv({M}, META["BLOCK_M"]) * triton.cdiv({N}, META["BLOCK_N"]),)
+    {kernel_name}[grid](a, b, c, out, {M}, {N}, {K})
+"""
+
+
+def _mm_configs(dtype, m, n, k, trans_a, trans_b, alpha, func_name):
+    condidates = [
+        # "BLOCK_M", "BLOCK_N", "BLOCK_K", "num_stages", "num_warps"
+        (32, 32, 16, 1, 2),
+        (64, 64, 16, 2, 4),
+        (64, 64, 32, 2, 4),
+        (64, 128, 32, 3, 4),
+        (128, 64, 32, 3, 4),
+        (64, 128, 32, 4, 8),
+        (128, 64, 32, 4, 8),
+        (64, 32, 32, 5, 8),
+        (32, 64, 32, 5, 8),
+        (128, 128, 32, 2, 8),
+        (64, 64, 64, 3, 8),
+        (32, 32, 128, 2, 4),
+    ]
+    tm = max(next_power_of_2(m), 16)
+    tn = max(next_power_of_2(n), 16)
+    tk = max(next_power_of_2(k), 16)
+    config_set = set()
+    config_strs = []
+    max_bk = 1
+    for bm, bn, bk, num_stages, num_warps in condidates:
+        new_bm = min(bm, tm)
+        new_bn = min(bn, tn)
+        new_bk = min(bk, tk)
+        if (new_bm, new_bn, new_bk) in config_set:
+            continue
+        config_set.add((new_bm, new_bn, new_bk))
+        if new_bk > max_bk:
+            max_bk = new_bk
+        new_num_warps = min(num_warps, new_bm * new_bn // 256)
+        config_strs.append(
+            f'        triton.Config({{"BLOCK_M": {new_bm}, "BLOCK_N": {new_bn}, "BLOCK_K": {new_bk}, "GROUP_M": 8}}, '
+            f"num_stages={num_stages}, num_warps={new_num_warps}),"
+        )
+    autotune_configs = "\n".join(config_strs)
+    post_process = "acc" if alpha == 1.0 else f"acc * {alpha}"
+    if dtype == torch.float16:
+        if alpha != 1.0:
+            post_process = f"({post_process})"
+        post_process = f"{post_process}.to(tl.float16)"
+    return dict(
+        autotune_configs=autotune_configs,
+        kernel_name=f"kernel_{func_name}",
+        M=m,
+        N=n,
+        K=k,
+        stride_am=(1 if trans_a else k),
+        stride_ak=(m if trans_a else 1),
+        stride_bk=(1 if trans_b else n),
+        stride_bn=(k if trans_b else 1),
+        even_k=(k % max_bk == 0),
+        allow_tf32=torch.backends.cuda.matmul.allow_tf32,
+        post_process=post_process,
+        func_name=func_name,
+    )
+
+
+def _gen_mm_key(dtype: torch.dtype, m: int, n: int, k: int, trans_a: bool, trans_b: bool, alpha: float) -> int:
+    return hash(f"mm|{dtype}|{m}|{n}|{k}|{trans_a}|{trans_b}|{alpha}") % (10**8)
+
+
+def _gen_mm_module(
+    dtype: torch.dtype, m: int, n: int, k: int, trans_a: bool, trans_b: bool, alpha: float
+) -> Tuple[str, ModuleType]:
+    func_name = f"mm_{_gen_mm_key(dtype, m, n, k, trans_a, trans_b, alpha)}"
+    kwargs = _mm_configs(dtype, m, n, k, trans_a, trans_b, alpha, func_name)
+    src_code = _MM_TEMPLATE.format(**kwargs)
+    if _DEBUG_MODE:
+        os.makedirs(os.path.dirname("triton_debug/"), exist_ok=True)
+        with open(f"triton_debug/{func_name}.py", "w") as f:
+            f.write(src_code)
+    return func_name, PyCodeCache().load(src_code)
+
+
+def _gen_gemm_key(
+    dtype: torch.dtype,
+    m: int,
+    n: int,
+    k: int,
+    stride_cm: int,
+    stride_cn: int,
+    trans_a: bool,
+    trans_b: bool,
+    alpha: float,
+    beta: float,
+) -> int:
+    return hash(f"gemm|{dtype}|{m}|{n}|{k}|{stride_cm}|{stride_cn}|{trans_a}|{trans_b}|{alpha}|{beta}") % (10**8)
+
+
+def _gen_gemm_module(
+    dtype: torch.dtype,
+    m: int,
+    n: int,
+    k: int,
+    stride_cm: int,
+    stride_cn: int,
+    trans_a: bool,
+    trans_b: bool,
+    alpha: float,
+    beta: float,
+) -> Tuple[str, ModuleType]:
+    func_name = f"gemm_{_gen_gemm_key(dtype, m, n, k, stride_cm, stride_cn, trans_a, trans_b, alpha, beta)}"
+    kwargs = _mm_configs(dtype, m, n, k, trans_a, trans_b, alpha, func_name)
+    kwargs["stride_cm"] = stride_cm
+    kwargs["stride_cn"] = stride_cn
+    kwargs["beta"] = beta
+    src_code = _GEMM_TEMPLATE.format(**kwargs)
+    if _DEBUG_MODE:
+        os.makedirs(os.path.dirname("triton_debug/"), exist_ok=True)
+        with open(f"triton_debug/{func_name}.py", "w") as f:
+            f.write(src_code)
+    return func_name, PyCodeCache().load(src_code)
+
+
+def _gen_bmm_key(
+    dtype: torch.dtype, m: int, n: int, k: int, batch_a: int, batch_b: int, trans_a: bool, trans_b: bool, alpha: float
+) -> int:
+    return hash(f"bmm|{dtype}|{m}|{n}|{k}|{batch_a}|{batch_b}|{trans_a}|{trans_b}|{alpha}") % (10**8)
+
+
+def _gen_bmm_module(
+    dtype: torch.dtype, m: int, n: int, k: int, batch_a: int, batch_b: int, trans_a: bool, trans_b: bool, alpha: float
+) -> Tuple[str, ModuleType]:
+    func_name = f"bmm_{_gen_bmm_key(dtype, m, n, k, batch_a, batch_b, trans_a, trans_b, alpha)}"
+    kwargs = _mm_configs(dtype, m, n, k, trans_a, trans_b, alpha, func_name)
+    batch = batch_a if batch_a >= batch_b else batch_b
+    kwargs["stride_aq"] = m * k if batch_a == batch else 0
+    kwargs["stride_bq"] = k * n if batch_b == batch else 0
+    kwargs["batch"] = batch
+    src_code = _BMM_TEMPLATE.format(**kwargs)
+    if _DEBUG_MODE:
+        os.makedirs(os.path.dirname("triton_debug/"), exist_ok=True)
+        with open(f"triton_debug/{func_name}.py", "w") as f:
+            f.write(src_code)
+    return func_name, PyCodeCache().load(src_code)
+
+
+def _matmul_internal(a, b, out, **kwargs):
+    rank_a = len(a.shape)
+    rank_b = len(b.shape)
+    assert rank_a >= 2 and rank_b >= 2
+    trans_a = kwargs.get("trans_a", False)
+    trans_b = kwargs.get("trans_b", False)
+    alpha = kwargs.get("alpha", 1.0)
+    m = a.shape[-1] if trans_a else a.shape[-2]
+    k = a.shape[-2] if trans_a else a.shape[-1]
+    bk = b.shape[-1] if trans_b else b.shape[-2]
+    assert k == bk
+    n = b.shape[-2] if trans_b else b.shape[-1]
+    assert rank_a == 2 or rank_b == 2 or a.shape[:-2] == b.shape[:-2]
+    batch_shape = a.shape[:-2] if rank_a >= 3 else b.shape[:-2]
+    if out is None:
+        out = torch.empty((*batch_shape, m, n), dtype=a.dtype, device=a.device)
+    else:
+        assert out.shape == (*batch_shape, m, n)
+    batch = math.prod(batch_shape)
+    dtype = a.dtype
+    if batch == 1 or (rank_a >= 3 and not trans_a and rank_b == 2):
+        if batch != 1:
+            m = batch * m
+        func_name, mod = ModuleCache.load(_gen_mm_key, _gen_mm_module, dtype, m, n, k, trans_a, trans_b, alpha)
+    else:
+        batch_a = batch if rank_a >= 3 else 1
+        batch_b = batch if rank_b >= 3 else 1
+        func_name, mod = ModuleCache.load(
+            _gen_bmm_key, _gen_bmm_module, dtype, m, n, k, batch_a, batch_b, trans_a, trans_b, alpha
+        )
+    func = getattr(mod, func_name)
+    func(a, b, out)
+    return out
+
+
+def triton_matmul(a, b, **kwargs):
+    """
+    Compute matrix multiplication of two tensors.
+    If trans_a is True in kwargs, input a will be transposed on last two dimensions before multiplication.
+    If trans_b is True in kwargs, input b will be transposed on last two dimensions before multiplication.
+    If alpha is specified in kwargs, the product will be multiplied by alpha.
+    The input tensors can be of rank 2 or higher, but currently support only simple broadcasting on batch dimensions.
+    I.e., the batch dimensions of a and b must either be equal, or one of them must be 1.
+    """
+
+    return _matmul_internal(a, b, None, **kwargs)
+
+
+def triton_matmul_out(a, b, out, **kwargs):
+    """
+    Same as triton_matmul, except that the output is allocated and passed from outside.
+    """
+
+    _matmul_internal(a, b, out, **kwargs)
+
+
+def _gemm_internal(a, b, c, out, **kwargs):
+    assert len(a.shape) == 2 and len(b.shape) == 2
+    trans_a = kwargs.get("trans_a", False)
+    trans_b = kwargs.get("trans_b", False)
+    alpha = kwargs.get("alpha", 1.0)
+    beta = kwargs.get("beta", 1.0)
+    m = a.shape[-1] if trans_a else a.shape[-2]
+    k = a.shape[-2] if trans_a else a.shape[-1]
+    bk = b.shape[-1] if trans_b else b.shape[-2]
+    assert k == bk
+    n = b.shape[-2] if trans_b else b.shape[-1]
+    stride_cm = c.shape[-1] if len(c.shape) == 2 and c.shape[-2] == m else 0
+    stride_cn = 1 if len(c.shape) >= 1 and c.shape[-1] == n else 0
+    if out is None:
+        out = torch.empty(m, n, dtype=a.dtype, device=a.device)
+    else:
+        assert out.shape == (m, n)
+    dtype = a.dtype
+    func_name, mod = ModuleCache.load(
+        _gen_gemm_key, _gen_gemm_module, dtype, m, n, k, stride_cm, stride_cn, trans_a, trans_b, alpha, beta
+    )
+    func = getattr(mod, func_name)
+    func(a, b, c, out)
+    return out
+
+
+def triton_gemm(a, b, c, **kwargs):
+    """
+    Compute alpha * a @ b + beta * c. Here a and b are 2D tensors, and c is broadcastable to the result.
+    If trans_a is True in kwargs, input a will be transposed before multiplication.
+    If trans_b is True in kwargs, input b will be transposed before multiplication.
+    """
+
+    return _gemm_internal(a, b, c, None, **kwargs)
+
+
+def triton_gemm_out(a, b, c, out, **kwargs):
+    """
+    Same as triton_gemm, except that the output is allocated and passed from outside.
+    """
+
+    _gemm_internal(a, b, c, out, **kwargs)

orttraining/orttraining/python/training/ort_triton/kernel/_slice_scel.py

@@ -0,0 +1,367 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import itertools
+import sys
+
+import torch
+import triton
+import triton.language as tl
+from onnx import TensorProto, helper
+
+from onnxruntime.training.ortmodule import register_graph_transformer
+
+from .._utils import get_attribute, to_numpy_array
+
+
+@triton.jit
+def _triton_slice_log_softmax(log_prob, logit, d: tl.constexpr, c: tl.constexpr, RBLOCK: tl.constexpr):
+    xoffset = tl.program_id(0)
+    logit_xoffset = (xoffset // d * (d + 1) + xoffset % d) * c
+    rbase = tl.arange(0, RBLOCK)
+    logit_max_row = tl.zeros([RBLOCK], tl.float32) + float("-inf")
+    for roffset in range(0, c, RBLOCK):
+        rindex = rbase + roffset
+        rmask = rindex < c
+        logit_row = tl.load(logit + logit_xoffset + rindex, mask=rmask, other=0.0).to(tl.float32)
+        logit_max_row = tl.where(rmask & (logit_max_row < logit_row), logit_row, logit_max_row)
+    logit_max_reduced = tl.max(logit_max_row, axis=0)
+    exp_sum_row = tl.zeros([RBLOCK], tl.float32)
+    for roffset in range(0, c, RBLOCK):
+        rindex = rbase + roffset
+        rmask = rindex < c
+        logit_row = tl.load(logit + logit_xoffset + rindex, mask=rmask, other=0.0).to(tl.float32)
+        exp_sum_row = tl.where(rmask, exp_sum_row + tl.exp(logit_row - logit_max_reduced), exp_sum_row)
+    reduced_log_sum = tl.log(tl.sum(exp_sum_row, axis=0)) + logit_max_reduced
+    for roffset in range(0, c, RBLOCK):
+        rindex = rbase + roffset
+        rmask = rindex < c
+        logit_row = tl.load(logit + logit_xoffset + rindex, mask=rmask, other=0.0).to(tl.float32)
+        output_row = logit_row - reduced_log_sum
+        tl.store(log_prob + xoffset * c + rindex, output_row, mask=rmask)
+
+
+@triton.jit
+def _triton_slice_scel(
+    loss,
+    factor,
+    log_prob,
+    label,
+    ignore_index,
+    d: tl.constexpr,
+    c: tl.constexpr,
+    n_cols: tl.constexpr,
+    RBLOCK: tl.constexpr,
+):
+    rbase = tl.arange(0, RBLOCK)
+    neg_sum_row = tl.zeros([RBLOCK], tl.float32)
+    factor_row = tl.zeros([RBLOCK], tl.float32)
+    for roffset in range(0, n_cols, RBLOCK):
+        rindex = rbase + roffset
+        rmask = rindex < n_cols
+        label_row = tl.load(label + (rindex // d) * (d + 1) + rindex % d + 1, mask=rmask, other=0.0).to(tl.int32)
+        mask = rmask & (label_row != ignore_index)
+        log_prob_row = tl.load(log_prob + rindex * c + label_row, mask=mask, other=0.0)
+        neg_sum_row = tl.where(mask, neg_sum_row - log_prob_row, neg_sum_row)
+        factor_row = tl.where(mask, factor_row + 1.0, factor_row)
+    reduced_neg_sum = tl.sum(neg_sum_row, axis=0)
+    reduced_factor = tl.sum(factor_row, axis=0)
+    loss_value = reduced_neg_sum / reduced_factor
+    tl.store(loss, loss_value)
+    tl.store(factor, reduced_factor)
+
+
+def slice_scel(logit, label, ignore_index):
+    ignore_index_value = ignore_index.item()
+    c = logit.shape[-1]
+    logit_d = logit.shape[-2]
+    d = logit_d - 1
+    n = logit.numel() // (logit_d * c)
+    log_prob_shape = list(logit.shape)[:-2] + [d, c]
+    log_prob = torch.empty(log_prob_shape, dtype=torch.float, device=logit.device)
+    rblock = 4096 if c > 4096 else triton.next_power_of_2(c)
+    num_warps = 16 if rblock >= 4096 else (8 if rblock >= 2048 else 4)
+    _triton_slice_log_softmax[(n * d,)](log_prob, logit, d, c, num_warps=num_warps, RBLOCK=rblock)
+    loss = torch.empty([], dtype=logit.dtype, device=logit.device)
+    factor = torch.empty([], dtype=torch.float, device=logit.device)
+    n_cols = n * d
+    rblock = 1024 if n_cols > 1024 else triton.next_power_of_2(n_cols)
+    _triton_slice_scel[(1,)](loss, factor, log_prob, label, ignore_index_value, d, c, n_cols, RBLOCK=rblock)
+    return loss, log_prob, factor
+
+
+@triton.jit
+def _triton_slice_scel_backward(
+    dlogit,
+    dloss,
+    log_prob,
+    label,
+    factor,
+    d: tl.constexpr,
+    c: tl.constexpr,
+    n_elements: tl.constexpr,
+    XBLOCK: tl.constexpr,
+):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)
+    xmask = xindex < n_elements
+    nd_index = xindex // c
+    dlogit_nd_index = (nd_index // d) * (d + 1) + nd_index % d
+    label_nd_index = dlogit_nd_index + 1
+    c_index = xindex % c
+    dloss_value = tl.load(dloss).to(tl.float32)
+    log_prob_row = tl.load(log_prob + xindex, mask=xmask, other=0.0)
+    label_row = tl.load(label + label_nd_index, mask=xmask, other=0.0).to(tl.int32)
+    factor_value = tl.load(factor)
+    dlogit_row = dloss_value * (tl.exp(log_prob_row) - tl.where(c_index == label_row, 1.0, 0.0)) / factor_value
+    tl.store(dlogit + dlogit_nd_index * c + c_index, dlogit_row, mask=xmask)
+
+
+@triton.jit
+def _triton_slice_scel_bias_backward(
+    dlogit,
+    dloss,
+    log_prob,
+    label,
+    factor,
+    bias,
+    dlogit_d: tl.constexpr,
+    c: tl.constexpr,
+    n_elements: tl.constexpr,
+    XBLOCK: tl.constexpr,
+):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)
+    xmask = xindex < n_elements
+    dlogit_nd_index = xindex // c
+    dlogit_n_index = dlogit_nd_index // dlogit_d
+    dlogit_d_index = dlogit_nd_index % dlogit_d
+    nd_index = dlogit_n_index * (dlogit_d - 1) + dlogit_d_index
+    nd_mask = xmask & (dlogit_d_index != dlogit_d - 1)
+    c_index = xindex % c
+    dloss_value = tl.load(dloss).to(tl.float32)
+    log_prob_row = tl.load(log_prob + nd_index * c + c_index, mask=nd_mask, other=0.0)
+    label_row = tl.load(label + dlogit_nd_index + 1, mask=nd_mask, other=0.0).to(tl.int32)
+    factor_value = tl.load(factor)
+    bias_row = tl.load(bias + xindex, mask=xmask, other=0.0).to(tl.float32)
+    dlogit_row = dloss_value * (tl.exp(log_prob_row) - tl.where(c_index == label_row, 1.0, 0.0)) / factor_value
+    dlogit_row = tl.where(nd_mask, dlogit_row, 0.0) + bias_row
+    tl.store(dlogit + xindex, dlogit_row, mask=xmask)
+
+
+def slice_scel_backward(dloss, log_prob, label, factor, bias):
+    c = log_prob.shape[-1]
+    d = log_prob.shape[-2]
+    dlogit_d = d + 1
+    dlogit_shape = list(log_prob.shape)[:-2] + [dlogit_d, c]
+    dlogit = (
+        torch.empty(dlogit_shape, dtype=dloss.dtype, device=dloss.device)
+        if bias is not None
+        else torch.zeros(dlogit_shape, dtype=dloss.dtype, device=dloss.device)
+    )
+    n_elements = dlogit.numel() if bias is not None else log_prob.numel()
+    xblock = 1024 if n_elements > 1024 else triton.next_power_of_2(n_elements)
+
+    def grid(meta):
+        return (triton.cdiv(n_elements, meta["XBLOCK"]),)
+
+    if bias is not None:
+        _triton_slice_scel_bias_backward[grid](
+            dlogit, dloss, log_prob, label, factor, bias, dlogit_d, c, n_elements, XBLOCK=xblock
+        )
+    else:
+        _triton_slice_scel_backward[grid](dlogit, dloss, log_prob, label, factor, d, c, n_elements, XBLOCK=xblock)
+    return dlogit
+
+
+def _get_producer(graph, arg, op_type):
+    for node in graph.node:
+        if node.op_type == op_type:
+            for output in node.output:
+                if output == arg:
+                    return node
+    return None
+
+
+def _get_consumer(graph, arg, op_type):
+    for node in graph.node:
+        if node.op_type == op_type:
+            for input in node.input:
+                if input == arg:
+                    return node
+    return None
+
+
+def _get_arg_info(graph, arg):
+    value_info = None
+    for info in itertools.chain(graph.input, graph.output, graph.value_info):
+        if info.name == arg:
+            value_info = info
+    if value_info is None or value_info.type is None:
+        return None, None
+    tensor_type = value_info.type.tensor_type
+    return tensor_type.elem_type, tensor_type.shape
+
+
+def _get_constant(graph, arg):
+    initializer = None
+    for init in graph.initializer:
+        if init.name == arg:
+            initializer = init
+    if initializer is None:
+        return None
+    return to_numpy_array(initializer)
+
+
+def _check_slice(graph, node, start, end, axis, step):
+    _, shape = _get_arg_info(graph, node.input[0])
+    if shape is None:
+        return False
+    rank = len(shape.dim)
+    if axis < 0:
+        axis += rank
+    for idx, value in enumerate([start, end, axis, step]):
+        constant = _get_constant(graph, node.input[idx + 1])
+        if constant is None or constant.size != 1:
+            return False
+        constant_value = constant.item()
+        if idx == 2 and constant_value < 0:
+            constant_value += rank
+        if constant_value != value:
+            return False
+    return True
+
+
+def _get_shape_related_nodes(graph, start_arg, sub_graph_nodes):
+    args = [start_arg]
+    while len(args) > 0:
+        arg = args.pop(0)
+        for node in graph.node:
+            if arg in node.input and node not in sub_graph_nodes:
+                sub_graph_nodes.append(node)
+                for output in node.output:
+                    if output not in args:
+                        args.append(output)
+
+
+@register_graph_transformer(devices="cuda")
+def transform_slice_scel(graph):
+    remove_nodes = []
+    triton_nodes = []
+    value_infos = []
+    idx = 0
+    for node in graph.node:
+        if node.op_type != "SoftmaxCrossEntropyLossInternal":
+            continue
+        # Weight input not supported for now, support reduction=mean only for now.
+        # It's required that the output_type is same as the logit dtype because currently we cannot pass the attribute
+        # value from TritonOp. We can add support of this in the future.
+        if len(node.input) > 2 and node.input[2]:
+            continue
+        reduction_attr = get_attribute(node, "reduction", "mean")
+        if isinstance(reduction_attr, bytes):
+            reduction_attr = reduction_attr.decode()
+        if reduction_attr != "mean":
+            continue
+        elem_type, _ = _get_arg_info(graph, node.input[0])
+        output_type_attr = get_attribute(node, "output_type", elem_type)
+        if output_type_attr != elem_type:
+            continue
+        reshape0 = _get_producer(graph, node.input[0], "Reshape")
+        if reshape0 is None:
+            continue
+        _, shape0 = _get_arg_info(graph, reshape0.input[0])
+        _, shape1 = _get_arg_info(graph, reshape0.output[0])
+        if shape0 is None or shape1 is None or shape0.dim[-1] != shape1.dim[-1]:
+            continue
+        slice0 = _get_producer(graph, reshape0.input[0], "Slice")
+        if slice0 is None or not _check_slice(graph, slice0, 0, -1, -2, 1):
+            continue
+        reshape1 = _get_producer(graph, node.input[1], "Reshape")
+        if reshape1 is None:
+            continue
+        slice1 = _get_producer(graph, reshape1.input[0], "Slice")
+        if slice1 is None or not _check_slice(graph, slice1, 1, sys.maxsize, -1, 1):
+            continue
+        scel_grad = _get_consumer(graph, node.output[1], "SoftmaxCrossEntropyLossInternalGrad")
+        if scel_grad is None:
+            continue
+        reshape2 = _get_consumer(graph, scel_grad.output[0], "Reshape")
+        if reshape2 is None:
+            continue
+        slice_grad = _get_consumer(graph, reshape2.output[0], "SliceGrad")
+        if slice_grad is None:
+            continue
+        shape_node = _get_producer(graph, slice_grad.input[1], "Shape")
+        if shape_node is None:
+            continue
+        bias_arg = ""
+        sum_node = _get_consumer(graph, slice_grad.output[0], "Sum")
+        if sum_node is not None and len(sum_node.input) == 2:
+            _, shape0 = _get_arg_info(graph, sum_node.input[0])
+            _, shape1 = _get_arg_info(graph, sum_node.input[1])
+            if shape0 is not None and shape0 == shape1:
+                bias_arg = sum_node.input[0] if sum_node.input[1] == slice_grad.output[0] else sum_node.input[1]
+        sub_graph_nodes = [node, reshape0, slice0, reshape1, slice1, scel_grad, reshape2, slice_grad, shape_node]
+        _get_shape_related_nodes(graph, slice0.output[0], sub_graph_nodes)
+        if bias_arg:
+            sub_graph_nodes.append(sum_node)
+        remove_nodes.extend(sub_graph_nodes)
+        forward_inputs = [slice0.input[0], slice1.input[0]]
+        if len(node.input) > 3:
+            forward_inputs.append(node.input[3])
+        else:
+            ignore_index_arg = "ignore_index_" + str(idx)
+            ignore_index_node = helper.make_node(
+                "Constant",
+                inputs=[],
+                outputs=[ignore_index_arg],
+                value=helper.make_tensor(
+                    name=ignore_index_arg,
+                    data_type=TensorProto.INT64,
+                    dims=(),
+                    vals=[-100],
+                ),
+            )
+            forward_inputs.append(ignore_index_arg)
+            triton_nodes.append(ignore_index_node)
+        # Use float for log_prob, which has better precision, but may consume more memory.
+        log_prob_arg = helper.make_tensor_value_info("scel_log_prob_" + str(idx), TensorProto.FLOAT, None)
+        factor_arg = helper.make_tensor_value_info("scel_factor_" + str(idx), TensorProto.FLOAT, None)
+        value_infos.extend([log_prob_arg, factor_arg])
+        triton_fw_node = helper.make_node(
+            "TritonOp",
+            forward_inputs,
+            [node.output[0], log_prob_arg.name, factor_arg.name],
+            "TritonOp_Slice_SCEL_" + str(idx),
+            None,
+            "com.microsoft",
+            func_name="slice_scel",
+        )
+        triton_nodes.append(triton_fw_node)
+        backward_outputs = [sum_node.output[0] if bias_arg else slice_grad.output[0]]
+        triton_bw_node = helper.make_node(
+            "TritonOp",
+            [scel_grad.input[0], log_prob_arg.name, slice1.input[0], factor_arg.name, bias_arg],
+            backward_outputs,
+            "TritonOp_Slice_SCEL_Backward_" + str(idx),
+            None,
+            "com.microsoft",
+            func_name="slice_scel_backward",
+        )
+        triton_nodes.append(triton_bw_node)
+        idx += 1
+
+    all_nodes = []
+    for node in graph.node:
+        if node not in remove_nodes:
+            all_nodes.append(node)
+
+    for node in triton_nodes:
+        all_nodes.append(node)
+
+    graph.ClearField("node")
+    graph.node.extend(all_nodes)
+    graph.value_info.extend(value_infos)

orttraining/orttraining/python/training/ort_triton/triton_op_executor.py

@@ -0,0 +1,98 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import functools
+import json
+import os
+import sys
+from types import ModuleType
+from typing import List, Tuple
+
+import onnx
+from torch._C import _from_dlpack
+from torch.utils.dlpack import to_dlpack
+
+from ._cache import ModuleCache, PyCodeCache
+from ._codegen import codegen
+from ._op_config import get_supported_ops
+from ._sorted_graph import SortedGraph
+from ._sympy_utils import parse_shape
+from ._utils import gen_unique_name
+
+_DEBUG_MODE = "ORTMODULE_TRITON_DEBUG" in os.environ and int(os.getenv("ORTMODULE_TRITON_DEBUG")) == 1
+
+
+@functools.lru_cache(None)
+def _gen_module_internal(sorted_graph: SortedGraph) -> Tuple[str, str, ModuleType]:
+    func_name = gen_unique_name("func")
+    src_code = codegen(func_name, sorted_graph)
+    return func_name, src_code, PyCodeCache().load(src_code)
+
+
+def _gen_key(onnx_key: int, onnx_str: bytes, shapes: List[List[int]]) -> int:
+    return hash(f"{onnx_key}|{str(shapes).replace(' ', '')}") % (10**8)
+
+
+def _gen_module(onnx_key: int, onnx_str: bytes, shapes: List[List[int]]) -> Tuple[str, ModuleType]:
+    model = onnx.load_model_from_string(onnx_str)
+    sorted_graph = SortedGraph(model, [parse_shape(shape) for shape in shapes])
+    if _DEBUG_MODE:
+        os.makedirs(os.path.dirname("triton_debug/"), exist_ok=True)
+        sorted_graph.save_onnx(f"triton_debug/{onnx_key}")
+    func_name, src_code, mod = _gen_module_internal(sorted_graph)
+    if _DEBUG_MODE:
+        py_file_path = f"triton_debug/{func_name}_{onnx_key}.py"
+        with open(py_file_path, "w") as f:
+            f.write(src_code)
+    return func_name, mod
+
+
+def get_config() -> str:
+    """
+    Get the supported ops and other configs in JSON format to control the Triton fusion on backend side.
+    All supported ops are from _op_config.py. The Triton fusion will try to fuse subgraphs with connected supported ops.
+    The initializer value can be "none", "scalar", and "all".
+        "none": no initializer will be added to subgraphs.
+        "scalar": only related scalar initializers will be added to subgraphs.
+        "all": all related initializers will be added to subgraphs.
+    The min_nodes is used to control the minimum number of non-no-op nodes in a subgraph.
+    """
+
+    config = {"ops": get_supported_ops(), "initializer": "scalar", "min_nodes": 2}
+    return json.dumps(config)
+
+
+def call_triton_by_name(func_name: str, *tensors, **kwargs):
+    """
+    Call triton kernel by function name. It's expected that there are functions and kernels registered manually
+    with that func_name (normally in .kernel sub-module), this function try to get the Python function by name
+    and execute it with the given tensors.
+    """
+
+    torch_tensors = [_from_dlpack(tensor) if tensor is not None else None for tensor in tensors]
+    func = getattr(sys.modules[".".join(__name__.split(".")[:-1])], func_name)
+    output = func(*torch_tensors, **kwargs)
+    if output is not None:
+        if isinstance(output, tuple):
+            return tuple([to_dlpack(tensor) for tensor in output])
+        return to_dlpack(output)
+    return None
+
+
+def call_triton_by_onnx(onnx_key: int, onnx_str: bytes, *tensors):
+    """
+    Call triton kernel by ONNX model. Load the ONNX model from onnx_str, generate the Triton function and kernels,
+    and execute the function with the given tensors.
+    """
+
+    assert all(tensor is not None for tensor in tensors)
+    torch_tensors = [_from_dlpack(tensor) for tensor in tensors]
+    concrete_shapes = [list(tensor.size()) for tensor in torch_tensors]
+    func_name, mod = ModuleCache.load(_gen_key, _gen_module, onnx_key, onnx_str, concrete_shapes)
+    func = getattr(mod, func_name)
+    output = func(*torch_tensors)
+    if isinstance(output, tuple):
+        return tuple([to_dlpack(tensor) for tensor in output])
+    return to_dlpack(output)

orttraining/orttraining/python/training/ortmodule/__init__.py

@@ -120,6 +120,7 @@ def _are_deterministic_algorithms_enabled():
     return ORTMODULE_IS_DETERMINISTIC
 
 
+from .graph_transformer_registry import register_graph_transformer  # noqa: E402, F401
 from .options import DebugOptions, LogLevel  # noqa: E402, F401
 
 # ORTModule must be loaded only after all validation passes

orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py

@@ -134,6 +134,11 @@ class GraphExecutionManager(GraphExecutionInterface):
         # Assign self._torch_alloc and self._torch_free if self._use_external_gpu_allocator is True
         self._get_torch_gpu_allocator_function_addresses()
 
+        if self._runtime_options.enable_triton:
+            from onnxruntime.training.ort_triton import register_triton_op_executor
+
+            register_triton_op_executor()
+
     def _get_torch_gpu_allocator_function_addresses(self):
         if self._runtime_options.use_external_gpu_allocator and torch.cuda.is_available():
             # CPP extension to get torch GPU allocator's alloc and free function addresses
@@ -197,6 +202,15 @@ class GraphExecutionManager(GraphExecutionInterface):
                 provider_option_map["cudnn_conv_algo_search"] = self._runtime_options.conv_algo_search
                 provider_option_map["cudnn_conv_use_max_workspace"] = "1"
                 provider_option_map["cudnn_conv1d_pad_to_nc1d"] = "1"
+                if self._runtime_options.enable_tuning:
+                    provider_option_map["tunable_op_enable"] = "1"
+                    provider_option_map["tunable_op_tuning_enable"] = "1"
+                    if self._runtime_options.max_tuning_duration_ms:
+                        provider_option_map["tunable_op_max_tuning_duration_ms"] = str(
+                            self._runtime_options.max_tuning_duration_ms
+                        )
+                elif self._runtime_options.tuning_results_path:
+                    provider_option_map["tunable_op_enable"] = "1"
             if self._runtime_options.use_external_gpu_allocator:
                 provider_option_map["gpu_external_alloc"] = str(self._torch_alloc)
                 provider_option_map["gpu_external_free"] = str(self._torch_free)
@@ -592,6 +606,29 @@ class GraphExecutionManager(GraphExecutionInterface):
             )
         )
 
+        if self._runtime_options.enable_triton:
+            feature_map.append(
+                (
+                    "TritonOp Enabled",
+                    True,
+                    "ORT will switch to Triton for executing some ops to further accelerate training.",
+                )
+            )
+
+        if self._runtime_options.enable_tuning:
+            desc = "Enable tunning Ops online"
+            if self._runtime_options.tuning_results_path:
+                desc += f", save tuning results to {self._runtime_options.tuning_results_path}"
+            feature_map.append(("Online Op Tuning", True, desc))
+        elif self._runtime_options.tuning_results_path:
+            feature_map.append(
+                (
+                    "Offline Op Tuning",
+                    True,
+                    f"Use offline tuning results from {self._runtime_options.tuning_results_path}",
+                )
+            )
+
         mode = "training" if self._export_mode == torch.onnx.TrainingMode.TRAINING else "inference"
         mode = f"{_logger.LogColor.UNDERLINE}{mode}{_logger.LogColor.ENDC}"
 

orttraining/orttraining/python/training/ortmodule/_inference_manager.py

@@ -16,6 +16,7 @@ from ._execution_agent import InferenceAgent
 from ._fallback import ORTModuleFallbackException, _FallbackManager, _FallbackPolicy
 from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo
 from ._logger import TimeTrackerPhase, TrackTime
+from ._utils import save_tuning_results, set_tuning_results
 from .options import DebugOptions, _SkipCheck
 
 
@@ -175,6 +176,15 @@ class InferenceManager(GraphExecutionManager):
                 *prepared_input_list,
             )
 
+            if (
+                create_execution_session
+                and self._runtime_options.enable_tuning
+                and self._runtime_options.tuning_results_path
+            ):
+                save_tuning_results(
+                    self._execution_agent._inference_session, False, self._runtime_options.tuning_results_path
+                )
+
             return _io.unflatten_user_output(self._module_output_schema, user_outputs)
         except ORTModuleFallbackException as e:
             # Exceptions subject to fallback are handled here
@@ -212,3 +222,8 @@ class InferenceManager(GraphExecutionManager):
         self._execution_agent = InferenceAgent(
             self._onnx_models.optimized_model.SerializeToString(), session_options, providers, provider_options
         )
+
+        if not self._runtime_options.enable_tuning and self._runtime_options.tuning_results_path:
+            set_tuning_results(
+                self._execution_agent._inference_session, False, self._runtime_options.tuning_results_path
+            )

orttraining/orttraining/python/training/ortmodule/_training_manager.py

@@ -20,6 +20,8 @@ from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo
 from ._io import _FlattenedModule, _InputInfo
 from ._logger import TimeTrackerPhase, TrackTime
 from ._runtime_inspector import Phase
+from ._utils import save_tuning_results, set_tuning_results
+from .graph_transformer_registry import GraphTransformerRegistry
 from .options import DebugOptions, _SkipCheck
 
 
@@ -316,10 +318,21 @@ class TrainingManager(GraphExecutionManager):
                 self._runtime_inspector,
             )
 
-            return _io.unflatten_user_output(
+            outputs = _io.unflatten_user_output(
                 self._module_output_schema,
                 self._forward_class.apply(*prepared_input_list),
             )
+
+            if (
+                create_execution_session
+                and self._runtime_options.enable_tuning
+                and self._runtime_options.tuning_results_path
+            ):
+                save_tuning_results(
+                    self._execution_agent._inference_session, True, self._runtime_options.tuning_results_path
+                )
+
+            return outputs
         except ORTModuleFallbackException as e:
             # Exceptions subject to fallback are handled here
             self._fallback_manager.handle_exception(exception=e, log_level=self._debug_options.logging.log_level)
@@ -345,6 +358,15 @@ class TrainingManager(GraphExecutionManager):
         self._onnx_models.optimized_pre_grad_model = onnx.load_model_from_string(
             self._graph_builder.get_forward_model()
         )
+
+        # Apply registered graph transformers to the optimized model
+        device_type = self._device.type
+        if device_type == "cuda" and self.is_rocm_pytorch:
+            device_type = "rocm"
+        GraphTransformerRegistry.transform_all(
+            type(self._flattened_module._original_module).__name__, device_type, self._onnx_models.optimized_model.graph
+        )
+
         if self._debug_options.save_onnx_models.save:
             self._onnx_models.save_optimized_model(
                 self._debug_options.save_onnx_models.path,
@@ -417,6 +439,11 @@ class TrainingManager(GraphExecutionManager):
             local_device_rank,
         )
 
+        if not self._runtime_options.enable_tuning and self._runtime_options.tuning_results_path:
+            set_tuning_results(
+                self._execution_agent._inference_session, True, self._runtime_options.tuning_results_path
+            )
+
     def _reinitialize_graph_builder(self, input_info: _InputInfo):
         """Return true if the module graph builder was reinitialized"""
 

orttraining/orttraining/python/training/ortmodule/_utils.py

@@ -7,6 +7,7 @@
 import copy
 import functools
 import inspect
+import json
 import logging
 import os
 import random
@@ -423,3 +424,23 @@ def get_runtime_pytorch_version():
 
 def check_function_has_param(function: Callable, param_name: str) -> bool:
     return param_name in inspect.signature(function).parameters
+
+
+def save_tuning_results(session, is_training, tuning_results_path):
+    """Save the online Op tuning results to a json file in the specified path."""
+
+    os.makedirs(tuning_results_path, exist_ok=True)
+    suffix = "training" if is_training else "inference"
+    tuning_result_file = os.path.join(tuning_results_path, f"tuning_results_{suffix}.json")
+    with open(tuning_result_file, "w", encoding="utf-8") as f:
+        json.dump(session.get_tuning_results(), f, indent=4)
+
+
+def set_tuning_results(session, is_training, tuning_results_path):
+    """Set the offline Op tuning results from a json file in the specified path."""
+
+    suffix = "training" if is_training else "inference"
+    tuning_result_file = os.path.join(tuning_results_path, f"tuning_results_{suffix}.json")
+    if os.path.isfile(tuning_result_file):
+        with open(tuning_result_file, encoding="utf-8") as f:
+            session.set_tuning_results(json.load(f))

orttraining/orttraining/python/training/ortmodule/graph_transformer_registry.py

@@ -0,0 +1,47 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from typing import Callable
+
+from onnx.onnx_ml_pb2 import GraphProto
+
+
+class GraphTransformerRegistry:
+    _TRANSFORMER_FUNCS = {}
+
+    @classmethod
+    def register(cls, target_modules: str, devices: str, priority: int, fn: Callable[[GraphProto], None]):
+        modules = []
+        if target_modules == "all":
+            modules.append("all")
+        else:
+            modules = target_modules.split("|")
+        for module in modules:
+            if module in cls._TRANSFORMER_FUNCS:
+                cls._TRANSFORMER_FUNCS[module].append((fn, devices, priority))
+            else:
+                cls._TRANSFORMER_FUNCS[module] = [(fn, devices, priority)]
+
+    @classmethod
+    def transform_all(cls, module_name: str, device: str, graph: GraphProto):
+        transformers_to_apply = []
+        if "all" in cls._TRANSFORMER_FUNCS:
+            transformers_to_apply.extend(cls._TRANSFORMER_FUNCS["all"])
+        if module_name in cls._TRANSFORMER_FUNCS:
+            transformers_to_apply.extend(cls._TRANSFORMER_FUNCS[module_name])
+        transformers_to_apply = [x for x in transformers_to_apply if x[1] == "all" or device in x[1]]
+        transformers_to_apply.sort(key=lambda x: x[2], reverse=True)
+        for fn, _, _ in transformers_to_apply:
+            fn(graph)
+
+
+# target_modules can be multiple module names separated by "|", or "all" means apply to all modules.
+# devices can be multiple device types separated by "|" or "all" means apply to all devices.
+def register_graph_transformer(target_modules: str = "all", devices: str = "all", priority: int = 0):
+    def graph_transformer_wrapper(fn):
+        GraphTransformerRegistry.register(target_modules, devices, priority, fn)
+        return fn
+
+    return graph_transformer_wrapper

orttraining/orttraining/python/training/ortmodule/options.py

@@ -236,6 +236,12 @@ class _RuntimeOptions:
             _SkipCheck.SKIP_CHECK_DEVICE | _SkipCheck.SKIP_CHECK_BUILD_GRADIENT | _SkipCheck.SKIP_CHECK_EXECUTION_AGENT
         )
 
+        # Triton support.
+        self.enable_triton = False
+        self.enable_tuning = False
+        self.max_tuning_duration_ms = 0
+        self.tuning_results_path = ""
+
         # Override the feature config if it exists in os env.
         self._override_from_env_vars()
 
@@ -285,3 +291,26 @@ class _RuntimeOptions:
                 lambda x, y: x | y,
                 [_SkipCheck[name] for name in parse_os_env_skip_check_flags("ORTMODULE_SKIPCHECK_POLICY")],
             )
+
+        # Configuration for Triton.
+        # Enable Triton op executor if Triton is installed, backend has support and environment variable is set.
+        if (
+            "ORTMODULE_USE_TRITON" in os.environ
+            and int(os.getenv("ORTMODULE_USE_TRITON")) == 1
+            and C.is_triton_enabled()
+        ):
+            try:
+                import triton  # noqa: F401
+            except ImportError:
+                pass
+            else:
+                self.enable_triton = True
+
+        if "ORTMODULE_ENABLE_TUNING" in os.environ and int(os.getenv("ORTMODULE_ENABLE_TUNING")) == 1:
+            self.enable_tuning = True
+        if "ORTMODULE_MAX_TUNING_DURATION_MS" in os.environ:
+            max_tuning_duration_ms = int(os.getenv("ORTMODULE_MAX_TUNING_DURATION_MS"))
+            if max_tuning_duration_ms > 0:
+                self.max_tuning_duration_ms = max_tuning_duration_ms
+        if "ORTMODULE_TUNING_RESULTS_PATH" in os.environ:
+            self.tuning_results_path = os.getenv("ORTMODULE_TUNING_RESULTS_PATH")

orttraining/orttraining/test/optimizer/graph_transform_test.cc

@@ -29,6 +29,9 @@
 #include "orttraining/core/optimizer/sce_loss_grad_bias_fusion.h"
 #include "orttraining/core/optimizer/qdq_fusion.h"
 #include "orttraining/core/optimizer/lstm_replacement.h"
+#ifdef ENABLE_TRITON
+#include "orttraining/core/optimizer/triton_fusion.h"
+#endif
 
 #include <random>
 
@@ -1212,11 +1215,188 @@ TEST_F(GraphTransformationTests, MegatronSelfAttentionPartitionCorrectnessTest)
 TEST_F(GraphTransformationTests, MegatronBARTSelfAttentionPartitionCorrectnessTest) {
   RunPartitionCorrectnessTest("testdata/transform/model_parallel/bart_self_attention_megatron_basic_test", *logger_, 2, {"input"}, {{6, 8, 4}});
 }
+
 // end of USE_CUDA
 #endif
 
 // end of DISABLE_CONTRIB_OPS
 #endif
 
+#ifdef ENABLE_TRITON
+TEST_F(GraphTransformationTests, TritonFusion) {
+  auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
+  std::shared_ptr<Model> model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
+  Graph& graph = model->MainGraph();
+  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+  ASSERT_TRUE(op_to_count["Add"] == 17);
+  ASSERT_TRUE(op_to_count["Sub"] == 1);
+  ASSERT_TRUE(op_to_count["Mul"] == 7);
+  ASSERT_TRUE(op_to_count["Div"] == 3);
+  ASSERT_TRUE(op_to_count["Cast"] == 7);
+  ASSERT_TRUE(op_to_count["Dropout"] == 4);
+  ASSERT_TRUE(op_to_count["Softmax"] == 1);
+  ASSERT_TRUE(op_to_count["LayerNormalization"] == 4);
+
+  {
+    auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
+    std::shared_ptr<Model> model;
+    ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
+    Graph& graph = model->MainGraph();
+    const char* config = R"(
+      {
+        "ops": {
+          "Add": { "versions": [13, 14] },
+          "Sub": { "versions": [13, 14] },
+          "Mul": { "versions": [13, 14] },
+          "Div": { "versions": [13, 14] },
+          "Cast": { "versions": [13] },
+          "Dropout": { "versions": [13] },
+          "Softmax": { "versions": [13], "conditions": { "axis": "-1" } },
+          "LayerNormalization": { "versions": [1], "conditions": { "axis": "-1" } }
+        },
+        "initializer": "scalar",
+        "min_nodes": 2
+      }
+    )";
+
+    std::unique_ptr<GraphTransformer> transformer = std::make_unique<TritonFusion>(config);
+    onnxruntime::GraphTransformerManager graph_transformation_mgr{1};
+    ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(transformer), TransformerLevel::Level1));
+    ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
+
+    op_to_count = CountOpsInGraph(graph);
+    ASSERT_TRUE(op_to_count["Add"] == 6);
+    ASSERT_TRUE(op_to_count["Sub"] == 0);
+    ASSERT_TRUE(op_to_count["Mul"] == 0);
+    ASSERT_TRUE(op_to_count["Div"] == 0);
+    ASSERT_TRUE(op_to_count["Cast"] == 4);
+    ASSERT_TRUE(op_to_count["Dropout"] == 0);
+    ASSERT_TRUE(op_to_count["Softmax"] == 0);
+    ASSERT_TRUE(op_to_count["LayerNormalization"] == 0);
+    ASSERT_TRUE(op_to_count["com.microsoft.TritonOp"] == 10);
+  }
+
+  // No Dropout.
+  {
+    auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
+    std::shared_ptr<Model> model;
+    ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
+    Graph& graph = model->MainGraph();
+    const char* config = R"(
+      {
+        "ops": {
+          "Add": { "versions": [13, 14] },
+          "Sub": { "versions": [13, 14] },
+          "Mul": { "versions": [13, 14] },
+          "Div": { "versions": [13, 14] },
+          "Cast": { "versions": [13] },
+          "Softmax": { "versions": [13], "conditions": { "axis": "-1" } },
+          "LayerNormalization": { "versions": [1], "conditions": { "axis": "-1" } }
+        },
+        "initializer": "scalar",
+        "min_nodes": 2
+      }
+    )";
+
+    std::unique_ptr<GraphTransformer> transformer = std::make_unique<TritonFusion>(config);
+    onnxruntime::GraphTransformerManager graph_transformation_mgr{1};
+    ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(transformer), TransformerLevel::Level1));
+    ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
+
+    op_to_count = CountOpsInGraph(graph);
+    ASSERT_TRUE(op_to_count["Add"] == 8);
+    ASSERT_TRUE(op_to_count["Sub"] == 0);
+    ASSERT_TRUE(op_to_count["Mul"] == 0);
+    ASSERT_TRUE(op_to_count["Div"] == 0);
+    ASSERT_TRUE(op_to_count["Cast"] == 4);
+    ASSERT_TRUE(op_to_count["Dropout"] == 4);
+    ASSERT_TRUE(op_to_count["Softmax"] == 0);
+    ASSERT_TRUE(op_to_count["LayerNormalization"] == 0);
+    ASSERT_TRUE(op_to_count["com.microsoft.TritonOp"] == 10);
+  }
+
+  // Ignore min nodes.
+  {
+    auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
+    std::shared_ptr<Model> model;
+    ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
+    Graph& graph = model->MainGraph();
+    const char* config = R"(
+      {
+        "ops": {
+          "Add": { "versions": [13, 14] },
+          "Sub": { "versions": [13, 14] },
+          "Mul": { "versions": [13, 14] },
+          "Div": { "versions": [13, 14] },
+          "Cast": { "versions": [13], "ignore_min_nodes": true },
+          "Dropout": { "versions": [13] },
+          "Softmax": { "versions": [13], "conditions": { "axis": "-1" } },
+          "LayerNormalization": { "versions": [1], "conditions": { "axis": "-1" } }
+        },
+        "initializer": "scalar",
+        "min_nodes": 2
+      }
+    )";
+
+    std::unique_ptr<GraphTransformer> transformer = std::make_unique<TritonFusion>(config);
+    onnxruntime::GraphTransformerManager graph_transformation_mgr{1};
+    ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(transformer), TransformerLevel::Level1));
+    ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
+
+    op_to_count = CountOpsInGraph(graph);
+    ASSERT_TRUE(op_to_count["Add"] == 6);
+    ASSERT_TRUE(op_to_count["Sub"] == 0);
+    ASSERT_TRUE(op_to_count["Mul"] == 0);
+    ASSERT_TRUE(op_to_count["Div"] == 0);
+    ASSERT_TRUE(op_to_count["Cast"] == 0);
+    ASSERT_TRUE(op_to_count["Dropout"] == 0);
+    ASSERT_TRUE(op_to_count["Softmax"] == 0);
+    ASSERT_TRUE(op_to_count["LayerNormalization"] == 0);
+    ASSERT_TRUE(op_to_count["com.microsoft.TritonOp"] == 14);
+  }
+
+  // Exclude Softmax using axis attribute.
+  {
+    auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
+    std::shared_ptr<Model> model;
+    ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
+    Graph& graph = model->MainGraph();
+    const char* config = R"(
+      {
+        "ops": {
+          "Add": { "versions": [13, 14] },
+          "Sub": { "versions": [13, 14] },
+          "Mul": { "versions": [13, 14] },
+          "Div": { "versions": [13, 14] },
+          "Cast": { "versions": [13]},
+          "Dropout": { "versions": [13] },
+          "Softmax": { "versions": [13], "conditions": { "axis": "1" } },
+          "LayerNormalization": { "versions": [1], "conditions": { "axis": "-1" } }
+        },
+        "initializer": "scalar",
+        "min_nodes": 2
+      }
+    )";
+
+    std::unique_ptr<GraphTransformer> transformer = std::make_unique<TritonFusion>(config);
+    onnxruntime::GraphTransformerManager graph_transformation_mgr{1};
+    ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(transformer), TransformerLevel::Level1));
+    ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
+
+    op_to_count = CountOpsInGraph(graph);
+    ASSERT_TRUE(op_to_count["Add"] == 6);
+    ASSERT_TRUE(op_to_count["Sub"] == 0);
+    ASSERT_TRUE(op_to_count["Mul"] == 0);
+    ASSERT_TRUE(op_to_count["Div"] == 0);
+    ASSERT_TRUE(op_to_count["Cast"] == 4);
+    ASSERT_TRUE(op_to_count["Dropout"] == 1);
+    ASSERT_TRUE(op_to_count["Softmax"] == 1);
+    ASSERT_TRUE(op_to_count["LayerNormalization"] == 0);
+    ASSERT_TRUE(op_to_count["com.microsoft.TritonOp"] == 10);
+  }
+}
+#endif
+
 }  // namespace test
 }  // namespace onnxruntime

orttraining/orttraining/test/python/orttraining_test_ortmodule_triton.py

@@ -0,0 +1,779 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import copy
+import json
+import os
+import random
+
+import _test_helpers
+import onnx
+import pytest
+import torch
+from onnx import TensorProto, helper
+from torch._C import _from_dlpack
+from torch.utils.dlpack import to_dlpack
+
+from onnxruntime.training.ort_triton import call_triton_by_name, call_triton_by_onnx
+from onnxruntime.training.ortmodule import DebugOptions, ORTModule
+
+pytest.importorskip("triton")
+
+DEVICE = "cuda"
+
+
+@pytest.fixture(scope="session", autouse=True)
+def run_before_test_session(request):
+    def insert_use_triton_env():
+        os.environ["ORTMODULE_USE_TRITON"] = "1"
+
+    def remove_use_triton_env():
+        del os.environ["ORTMODULE_USE_TRITON"]
+
+    insert_use_triton_env()
+    request.addfinalizer(remove_use_triton_env)
+
+
+def _onnx_dtype_to_torch_dtype(onnx_dtype):
+    if onnx_dtype == TensorProto.FLOAT:
+        return torch.float32
+    elif onnx_dtype == TensorProto.FLOAT16:
+        return torch.float16
+    else:
+        raise RuntimeError("Unsupported ONNX dtype")
+
+
+def _torch_add(input1, input2):
+    return input1 + input2
+
+
+def _torch_sub(input1, input2):
+    return input1 - input2
+
+
+def _torch_mul(input1, input2):
+    return input1 * input2
+
+
+def _torch_div(input1, input2):
+    return input1 / input2
+
+
+def _torch_pow(input, exponent):
+    return torch.pow(input, exponent)
+
+
+def _torch_sqrt(input):
+    return torch.sqrt(input)
+
+
+def _torch_exp(input):
+    return torch.exp(input)
+
+
+def _torch_cast(input, **kwargs):
+    return input.to(_onnx_dtype_to_torch_dtype(kwargs.get("to")))
+
+
+def _torch_where(condition, x, y):
+    return torch.where(condition, x, y)
+
+
+def _torch_sum(*inputs):
+    result = inputs[0]
+    for input in inputs[1:]:
+        result += input
+    return result
+
+
+def _troch_dropout_gard(dy, mask, **kwargs):
+    ratio = kwargs.get("ratio", 0.5)
+    return torch.where(mask, dy / (1.0 - ratio), 0.0)
+
+
+def _torch_softmax(input, **kwargs):
+    axis = kwargs.get("axis", -1)
+    return torch.softmax(input, axis)
+
+
+def _torch_reduce(input, func, **kwargs):
+    rank = len(input.shape)
+    axes = kwargs.get("axes", [idx for idx in range(rank)])
+    keepdims = kwargs.get("keepdims", True)
+    axes = [axis if axis >= 0 else rank + axis for axis in axes]
+    axes.sort(reverse=True)
+    result = input
+    for axis in axes:
+        result = func(result, dim=axis, keepdim=keepdims)
+        if func == torch.max or func == torch.min:
+            result = result[0]
+    return result
+
+
+def _torch_reduce_sum(input, **kwargs):
+    return _torch_reduce(input, torch.sum, **kwargs)
+
+
+def _torch_reduce_mean(input, **kwargs):
+    return _torch_reduce(input, torch.mean, **kwargs)
+
+
+def _torch_reduce_max(input, **kwargs):
+    return _torch_reduce(input, torch.max, **kwargs)
+
+
+def _torch_reduce_min(input, **kwargs):
+    return _torch_reduce(input, torch.min, **kwargs)
+
+
+def _torch_layer_norm(input, weight, bias, **kwargs):
+    rank = len(input.shape)
+    axis = kwargs.get("axis", -1)
+    if axis < 0:
+        axis += rank
+    normalized_shape = input.shape[axis:]
+    return torch.nn.functional.layer_norm(input, normalized_shape, weight, bias)
+
+
+class TorchFuncExecutor:
+    _INFER_FUNC_MAP = {
+        "Add": _torch_add,
+        "Sub": _torch_sub,
+        "Mul": _torch_mul,
+        "Div": _torch_div,
+        "Pow": _torch_pow,
+        "Sqrt": _torch_sqrt,
+        "Exp": _torch_exp,
+        "Cast": _torch_cast,
+        "Where": _torch_where,
+        "Sum": _torch_sum,
+        "DropoutGrad": _troch_dropout_gard,
+        "ReduceSum": _torch_reduce_sum,
+        "ReduceMean": _torch_reduce_mean,
+        "ReduceMax": _torch_reduce_max,
+        "ReduceMin": _torch_reduce_min,
+        "Softmax": _torch_softmax,
+        "LayerNormalization": _torch_layer_norm,
+    }
+
+    @classmethod
+    def run(cls, op_type, *torch_tensors, **kwargs):
+        if op_type not in cls._INFER_FUNC_MAP:
+            raise NotImplementedError(f"Unsupported op type: {op_type}")
+        return cls._INFER_FUNC_MAP[op_type](*torch_tensors, **kwargs)
+
+
+def _run_op_test(op_type, onnx_dtype, create_model_func, gen_inputs_func, **kwargs):
+    rtol = kwargs.get("rtol", 1e-03 if onnx_dtype == TensorProto.FLOAT16 else 1e-04)
+    atol = kwargs.get("atol", 1e-03 if onnx_dtype == TensorProto.FLOAT16 else 1e-05)
+    pt_inputs = gen_inputs_func(_onnx_dtype_to_torch_dtype(onnx_dtype))
+    ort_inputs = copy.deepcopy(pt_inputs)
+    ort_inputs = [tensor.to(torch.uint8) if tensor.dtype == torch.bool else tensor for tensor in ort_inputs]
+    pt_outputs = TorchFuncExecutor.run(op_type, *pt_inputs, **kwargs)
+    model_str = create_model_func(op_type, onnx_dtype, **kwargs).SerializeToString()
+    ort_outputs = call_triton_by_onnx(hash(model_str), model_str, *[to_dlpack(tensor) for tensor in ort_inputs])
+    if isinstance(pt_outputs, tuple):
+        assert isinstance(ort_outputs, tuple)
+        assert len(pt_outputs) == len(ort_outputs)
+        for pt_output, ort_output in zip(pt_outputs, ort_outputs):
+            _test_helpers.assert_values_are_close(pt_output, _from_dlpack(ort_output), rtol=rtol, atol=atol)
+    else:
+        _test_helpers.assert_values_are_close(pt_outputs, _from_dlpack(ort_outputs), rtol=rtol, atol=atol)
+
+
+def _run_step(model, *tensors):
+    prediction = model(*tensors)
+    loss = prediction.sum()
+    loss.backward()
+    return prediction
+
+
+def _run_module_test(module_cls, dtype, gen_inputs_func, triton_op_count, **kwargs):
+    pt_model = module_cls().to(DEVICE).to(dtype)
+    ort_model = ORTModule(copy.deepcopy(pt_model), DebugOptions(save_onnx=True, onnx_prefix="triton_model"))
+    rtol = kwargs.get("rtol", 1e-03 if dtype == torch.float16 else 1e-04)
+    atol = kwargs.get("atol", 1e-03 if dtype == torch.float16 else 1e-05)
+    for _ in range(10):
+        pt_inputs = gen_inputs_func(dtype)
+        ort_inputs = copy.deepcopy(pt_inputs)
+        pt_output = _run_step(pt_model, *pt_inputs)
+        ort_output = _run_step(ort_model, *ort_inputs)
+        _test_helpers.assert_values_are_close(pt_output, ort_output, rtol=rtol, atol=atol)
+        _test_helpers.assert_gradients_match_and_reset_gradient(pt_model, ort_model, rtol=rtol, atol=atol)
+        for i in range(len(pt_inputs)):
+            if pt_inputs[i].requires_grad:
+                _test_helpers.assert_values_are_close(pt_inputs[i].grad, ort_inputs[i].grad, rtol=rtol, atol=atol)
+
+    assert os.path.exists(os.path.join(os.getcwd(), "triton_model_torch_exported_training.onnx"))
+    assert os.path.exists(os.path.join(os.getcwd(), "triton_model_optimized_training.onnx"))
+    assert os.path.exists(os.path.join(os.getcwd(), "triton_model_optimized_pre_grad_training.onnx"))
+    assert os.path.exists(os.path.join(os.getcwd(), "triton_model_execution_model_training.onnx"))
+    model = onnx.load(os.path.join(os.getcwd(), "triton_model_execution_model_training.onnx"))
+    actual_triton_op_count = 0
+    for node in model.graph.node:
+        if node.op_type == "TritonOp":
+            actual_triton_op_count += 1
+    assert actual_triton_op_count == triton_op_count
+    os.remove(os.path.join(os.getcwd(), "triton_model_torch_exported_training.onnx"))
+    os.remove(os.path.join(os.getcwd(), "triton_model_optimized_training.onnx"))
+    os.remove(os.path.join(os.getcwd(), "triton_model_optimized_pre_grad_training.onnx"))
+    os.remove(os.path.join(os.getcwd(), "triton_model_execution_model_training.onnx"))
+
+
+def _run_tunable_op_test(module_cls, dtype, gen_inputs_func, tunable_op, impl_count, **kwargs):
+    pt_model = module_cls().to(DEVICE).to(dtype)
+    ort_model = ORTModule(copy.deepcopy(pt_model))
+    rtol = kwargs.get("rtol", 1e-03 if dtype == torch.float16 else 1e-04)
+    atol = kwargs.get("atol", 1e-03 if dtype == torch.float16 else 1e-05)
+    os.environ["ORTMODULE_ENABLE_TUNING"] = "1"
+    os.environ["ORTMODULE_TUNING_RESULTS_PATH"] = "./"
+    for _ in range(5):
+        pt_inputs = gen_inputs_func(dtype)
+        ort_inputs = copy.deepcopy(pt_inputs)
+        pt_output = _run_step(pt_model, *pt_inputs)
+        ort_output = _run_step(ort_model, *ort_inputs)
+        _test_helpers.assert_values_are_close(pt_output, ort_output, rtol=rtol, atol=atol)
+        _test_helpers.assert_gradients_match_and_reset_gradient(pt_model, ort_model, rtol=rtol, atol=atol)
+    tunable_results_file = os.path.join(os.getcwd(), "tuning_results_training.json")
+    assert os.path.exists(tunable_results_file)
+    with open(tunable_results_file) as f:
+        tunable_results = json.load(f)
+    assert tunable_op in str(tunable_results)
+    del os.environ["ORTMODULE_ENABLE_TUNING"]
+    for i in range(impl_count - 1):
+        new_tunable_results = copy.deepcopy(tunable_results)
+        for k, v in new_tunable_results[0]["results"].items():
+            if tunable_op in k:
+                for param, impl in v.items():
+                    v[param] = (impl + 1 + i) % impl_count
+        with open(tunable_results_file, "w") as f:
+            json.dump(new_tunable_results, f)
+        ort_model = ORTModule(copy.deepcopy(pt_model))
+        for _ in range(5):
+            pt_inputs = gen_inputs_func(dtype)
+            ort_inputs = copy.deepcopy(pt_inputs)
+            pt_output = _run_step(pt_model, *pt_inputs)
+            ort_output = _run_step(ort_model, *ort_inputs)
+            _test_helpers.assert_values_are_close(pt_output, ort_output, rtol=rtol, atol=atol)
+            _test_helpers.assert_gradients_match_and_reset_gradient(pt_model, ort_model, rtol=rtol, atol=atol)
+    os.remove(tunable_results_file)
+    del os.environ["ORTMODULE_TUNING_RESULTS_PATH"]
+
+
+@pytest.mark.parametrize("op_type", ["Add", "Sub", "Mul", "Div"])
+@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT, TensorProto.FLOAT16])
+@pytest.mark.parametrize("input_shapes", [([1024, 2], [1024, 2]), ([2, 3, 3, 3], [3, 1, 3]), ([2049], [1])])
+def test_binary_elementwise_op(op_type, onnx_dtype, input_shapes):
+    def _create_model(op_type, onnx_dtype):
+        graph = helper.make_graph(
+            [helper.make_node(op_type, ["X", "Y"], ["Z"], name="test")],
+            "test",
+            [
+                helper.make_tensor_value_info("X", onnx_dtype, None),
+                helper.make_tensor_value_info("Y", onnx_dtype, None),
+            ],
+            [helper.make_tensor_value_info("Z", onnx_dtype, None)],
+        )
+        return helper.make_model(graph, producer_name="test")
+
+    def _gen_inputs(dtype):
+        return [
+            torch.randn(*input_shapes[0], dtype=dtype, device=DEVICE),
+            torch.randn(*input_shapes[1], dtype=dtype, device=DEVICE),
+        ]
+
+    _run_op_test(op_type, onnx_dtype, _create_model, _gen_inputs)
+
+
+@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT, TensorProto.FLOAT16])
+@pytest.mark.parametrize("input_shapes", [([1024, 2], [1024, 2]), ([2, 3, 3, 3], [3, 1, 3], [2, 1, 3, 1])])
+def test_sum_op(onnx_dtype, input_shapes):
+    def _create_model(op_type, onnx_dtype):
+        graph = helper.make_graph(
+            [helper.make_node(op_type, [f"I{idx}" for idx in range(len(input_shapes))], ["O"], name="test")],
+            "test",
+            [helper.make_tensor_value_info(f"I{idx}", onnx_dtype, None) for idx in range(len(input_shapes))],
+            [helper.make_tensor_value_info("O", onnx_dtype, None)],
+        )
+        return helper.make_model(graph, producer_name="test")
+
+    def _gen_inputs(dtype):
+        return [torch.randn(*input_shape, dtype=dtype, device=DEVICE) for input_shape in input_shapes]
+
+    _run_op_test("Sum", onnx_dtype, _create_model, _gen_inputs)
+
+
+@pytest.mark.parametrize("op_type", ["Sqrt", "Exp"])
+@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT, TensorProto.FLOAT16])
+@pytest.mark.parametrize("input_shape", [[1024, 4], [2, 3, 3, 3], [2049, 1]])
+def test_unary_elementwise_op(op_type, onnx_dtype, input_shape):
+    def _create_model(op_type, onnx_dtype):
+        graph = helper.make_graph(
+            [helper.make_node(op_type, ["X"], ["Y"], name="test")],
+            "test",
+            [helper.make_tensor_value_info("X", onnx_dtype, None)],
+            [helper.make_tensor_value_info("Y", onnx_dtype, None)],
+        )
+        return helper.make_model(graph, producer_name="test")
+
+    def _gen_inputs(dtype):
+        return [torch.rand(*input_shape, dtype=dtype, device=DEVICE)]
+
+    _run_op_test(op_type, onnx_dtype, _create_model, _gen_inputs)
+
+
+@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT, TensorProto.FLOAT16])
+@pytest.mark.parametrize("input_shape_and_exponent", [([1024, 2], 2.0), ([2, 3, 3, 3], 0.5), ([2049], 3.0)])
+def test_pow_op(onnx_dtype, input_shape_and_exponent):
+    def _create_model(op_type, onnx_dtype):
+        graph = helper.make_graph(
+            [helper.make_node(op_type, ["X", "Y"], ["Z"], name="test")],
+            "test",
+            [
+                helper.make_tensor_value_info("X", onnx_dtype, None),
+                helper.make_tensor_value_info("Y", onnx_dtype, None),
+            ],
+            [helper.make_tensor_value_info("Z", onnx_dtype, None)],
+        )
+        return helper.make_model(graph, producer_name="test")
+
+    def _gen_inputs(dtype):
+        return [
+            torch.rand(*input_shape_and_exponent[0], dtype=dtype, device=DEVICE),
+            torch.tensor(input_shape_and_exponent[1], dtype=dtype, device=DEVICE),
+        ]
+
+    _run_op_test("Pow", onnx_dtype, _create_model, _gen_inputs)
+
+
+@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT, TensorProto.FLOAT16])
+@pytest.mark.parametrize("input_shape", [[1024, 2], [2, 3, 3, 3], [1, 2050]])
+def test_cast_op(onnx_dtype, input_shape):
+    def _create_model(op_type, onnx_dtype, **kwargs):
+        graph = helper.make_graph(
+            [helper.make_node(op_type, ["X"], ["Y"], name="test", **kwargs)],
+            "test",
+            [
+                helper.make_tensor_value_info("X", onnx_dtype, None),
+            ],
+            [helper.make_tensor_value_info("Y", kwargs.get("to"), None)],
+        )
+        return helper.make_model(graph, producer_name="test")
+
+    def _gen_inputs(dtype):
+        return [torch.randn(*input_shape, dtype=dtype, device=DEVICE)]
+
+    kwargs = {"to": TensorProto.FLOAT16 if onnx_dtype == TensorProto.FLOAT else TensorProto.FLOAT}
+    _run_op_test("Cast", onnx_dtype, _create_model, _gen_inputs, **kwargs)
+
+
+@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT, TensorProto.FLOAT16])
+@pytest.mark.parametrize("input_shapes", [([2, 1024], [2, 1024], [2, 1024]), ([2, 1, 3, 1], [2, 3, 3, 3], [3, 1, 3])])
+def test_where_op(onnx_dtype, input_shapes):
+    def _create_model(op_type, onnx_dtype):
+        graph = helper.make_graph(
+            [helper.make_node(op_type, ["C", "X", "Y"], ["Z"], name="test")],
+            "test",
+            [
+                helper.make_tensor_value_info("C", TensorProto.BOOL, None),
+                helper.make_tensor_value_info("X", onnx_dtype, None),
+                helper.make_tensor_value_info("Y", onnx_dtype, None),
+            ],
+            [helper.make_tensor_value_info("Z", onnx_dtype, None)],
+        )
+        return helper.make_model(graph, producer_name="test")
+
+    def _gen_inputs(dtype):
+        return [
+            torch.rand(*input_shapes[0], dtype=dtype, device=DEVICE) < 0.5,
+            torch.randn(*input_shapes[1], dtype=dtype, device=DEVICE),
+            torch.randn(*input_shapes[2], dtype=dtype, device=DEVICE),
+        ]
+
+    _run_op_test("Where", onnx_dtype, _create_model, _gen_inputs)
+
+
+@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT, TensorProto.FLOAT16])
+@pytest.mark.parametrize("input_shape_and_ratio", [([1024, 2], 0.2), ([25, 75], 0.6), ([1, 2049], 0.5)])
+def test_dropout_op(onnx_dtype, input_shape_and_ratio):
+    graph = helper.make_graph(
+        [
+            helper.make_node("Add", ["X1", "X2"], ["T1"], name="add1"),
+            helper.make_node("Dropout", ["T1", "ratio"], ["Y1", "mask1"], name="dropout1"),
+            helper.make_node("Add", ["Y1", "X3"], ["T2"], name="add2"),
+            helper.make_node("Dropout", ["T2", "ratio"], ["Y2", "mask2"], name="dropout2"),
+        ],
+        "test",
+        [
+            helper.make_tensor_value_info("X1", onnx_dtype, None),
+            helper.make_tensor_value_info("X2", onnx_dtype, None),
+            helper.make_tensor_value_info("X3", onnx_dtype, None),
+        ],
+        [
+            helper.make_tensor_value_info("Y1", onnx_dtype, None),
+            helper.make_tensor_value_info("mask1", TensorProto.BOOL, None),
+            helper.make_tensor_value_info("Y2", onnx_dtype, None),
+            helper.make_tensor_value_info("mask2", TensorProto.BOOL, None),
+        ],
+        initializer=[helper.make_tensor("ratio", TensorProto.FLOAT, (), [input_shape_and_ratio[1]])],
+    )
+    model_str = helper.make_model(graph, producer_name="test").SerializeToString()
+    torch_dtype = _onnx_dtype_to_torch_dtype(onnx_dtype)
+    input_tensor = [
+        torch.randn(*input_shape_and_ratio[0], dtype=torch_dtype, device=DEVICE),
+        torch.randn(*input_shape_and_ratio[0], dtype=torch_dtype, device=DEVICE),
+        torch.randn(*input_shape_and_ratio[0], dtype=torch_dtype, device=DEVICE),
+    ]
+    outputs = call_triton_by_onnx(hash(model_str), model_str, *[to_dlpack(t) for t in input_tensor])
+    y1, mask1, y2, mask2 = tuple([_from_dlpack(o).detach().cpu().numpy().flatten() for o in outputs])
+    x1 = (input_tensor[0] + input_tensor[1]).detach().cpu().numpy().flatten()
+    x2 = y1 * mask1 + input_tensor[2].detach().cpu().numpy().flatten()
+
+    def _check_output(x, y, mask, ratio):
+        all_count = 0
+        masked_count = 0
+        for x_value, y_value, mask_value in zip(x, y, mask):
+            if mask_value:
+                assert abs(y_value - x_value / (1.0 - ratio)) < 0.05
+            else:
+                assert y_value == 0
+            if not mask_value:
+                masked_count += 1
+            all_count += 1
+        assert abs(masked_count / all_count - ratio) < 0.05
+
+    _check_output(x1, y1, mask1, input_shape_and_ratio[1])
+    _check_output(x2, y2, mask2, input_shape_and_ratio[1])
+    assert any(mask1[i] != mask2[i] for i in range(len(mask1)))
+
+
+@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT, TensorProto.FLOAT16])
+@pytest.mark.parametrize("input_shape_and_ratio", [([1024, 2], 0.2), ([25, 75], 0.6), ([1, 2049], 0.5)])
+def test_dropout_grad_op(onnx_dtype, input_shape_and_ratio):
+    def _create_model(op_type, onnx_dtype, **kwargs):
+        graph = helper.make_graph(
+            [
+                helper.make_node(op_type, ["dY", "mask", "ratio"], ["dX"], name="test", domain="com.microsoft"),
+            ],
+            "test",
+            [
+                helper.make_tensor_value_info("dY", onnx_dtype, None),
+                helper.make_tensor_value_info("mask", TensorProto.BOOL, None),
+            ],
+            [helper.make_tensor_value_info("dX", onnx_dtype, None)],
+            initializer=[helper.make_tensor("ratio", TensorProto.FLOAT, (), [input_shape_and_ratio[1]])],
+        )
+        return helper.make_model(graph, producer_name="test")
+
+    def _gen_inputs(dtype):
+        return [
+            torch.randn(*input_shape_and_ratio[0], dtype=dtype, device=DEVICE),
+            torch.rand(*input_shape_and_ratio[0], dtype=dtype, device=DEVICE) >= input_shape_and_ratio[1],
+        ]
+
+    kwargs = {"ratio": input_shape_and_ratio[1]}
+    _run_op_test("DropoutGrad", onnx_dtype, _create_model, _gen_inputs, **kwargs)
+
+
+@pytest.mark.parametrize("op_type", ["ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"])
+@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT, TensorProto.FLOAT16])
+@pytest.mark.parametrize(
+    "input_shape_and_reduce_info",
+    [
+        ([2, 1024], [-1], True),
+        ([1050, 3], [0], False),
+        ([2, 3, 3, 3], [1, 2], True),
+        ([123, 4, 5, 6], [2], False),
+        ([16, 8, 16, 8], [1, 3], True),
+        ([16, 8, 16, 8], [0, 2], False),
+    ],
+)
+def test_reduce_op(op_type, onnx_dtype, input_shape_and_reduce_info):
+    def _create_model(op_type, onnx_dtype, **kwargs):
+        reduce_inputs = ["X"]
+        initializer = []
+        if input_shape_and_reduce_info[1] is not None:
+            reduce_inputs.append("axes")
+            initializer.append(
+                helper.make_tensor(
+                    "axes",
+                    onnx.TensorProto.INT64,
+                    [len(input_shape_and_reduce_info[1])],
+                    input_shape_and_reduce_info[1],
+                )
+            )
+        node = (
+            helper.make_node(op_type, reduce_inputs, ["Y"], name="test", keepdims=input_shape_and_reduce_info[2])
+            if input_shape_and_reduce_info[2] is not None
+            else helper.make_node(op_type, reduce_inputs, ["Y"], name="test")
+        )
+        graph = helper.make_graph(
+            [node],
+            "test",
+            [helper.make_tensor_value_info("X", onnx_dtype, None)],
+            [helper.make_tensor_value_info("Y", onnx_dtype, None)],
+            initializer=initializer,
+        )
+        return helper.make_model(graph, producer_name="test")
+
+    def _gen_inputs(dtype):
+        return [torch.randn(*input_shape_and_reduce_info[0], dtype=dtype, device=DEVICE)]
+
+    kwargs = {}
+    if input_shape_and_reduce_info[1] is not None:
+        kwargs["axes"] = input_shape_and_reduce_info[1]
+    if input_shape_and_reduce_info[2] is not None:
+        kwargs["keepdims"] = input_shape_and_reduce_info[2]
+    if onnx_dtype == TensorProto.FLOAT16:
+        kwargs["atol"] = 1e-2
+        kwargs["rtol"] = 1e-2
+    _run_op_test(op_type, onnx_dtype, _create_model, _gen_inputs, **kwargs)
+
+
+@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT, TensorProto.FLOAT16])
+@pytest.mark.parametrize("input_shape_and_axis", [([2, 1024], -1), ([2, 3, 3, 3], 1), ([4, 2049], 1)])
+def test_softmax_op(onnx_dtype, input_shape_and_axis):
+    def _create_model(op_type, onnx_dtype, **kwargs):
+        graph = helper.make_graph(
+            [helper.make_node(op_type, ["X"], ["Y"], name="test", **kwargs)],
+            "test",
+            [helper.make_tensor_value_info("X", onnx_dtype, None)],
+            [helper.make_tensor_value_info("Y", onnx_dtype, None)],
+        )
+        return helper.make_model(graph, producer_name="test")
+
+    def _gen_inputs(dtype):
+        return [torch.randn(*input_shape_and_axis[0], dtype=dtype, device=DEVICE)]
+
+    kwargs = {"axis": input_shape_and_axis[1]}
+    _run_op_test("Softmax", onnx_dtype, _create_model, _gen_inputs, **kwargs)
+
+
+@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT, TensorProto.FLOAT16])
+@pytest.mark.parametrize("input_shape_and_axis", [([2, 1024], -1), ([2, 3, 3, 3], 2), ([4, 2049], 1)])
+def test_layer_norm_op(onnx_dtype, input_shape_and_axis):
+    def _create_model(op_type, onnx_dtype, **kwargs):
+        graph = helper.make_graph(
+            [helper.make_node(op_type, ["X", "W", "B"], ["Y"], name="test", **kwargs)],
+            "test",
+            [
+                helper.make_tensor_value_info("X", onnx_dtype, None),
+                helper.make_tensor_value_info("W", onnx_dtype, None),
+                helper.make_tensor_value_info("B", onnx_dtype, None),
+            ],
+            [helper.make_tensor_value_info("Y", onnx_dtype, None)],
+        )
+        return helper.make_model(graph, producer_name="test")
+
+    def _gen_inputs(dtype):
+        return [
+            torch.randn(*input_shape_and_axis[0], dtype=dtype, device=DEVICE),
+            torch.randn(*input_shape_and_axis[0][input_shape_and_axis[1] :], dtype=dtype, device=DEVICE),
+            torch.randn(input_shape_and_axis[0][input_shape_and_axis[1] :], dtype=dtype, device=DEVICE),
+        ]
+
+    kwargs = {"axis": input_shape_and_axis[1]}
+    _run_op_test("LayerNormalization", onnx_dtype, _create_model, _gen_inputs, **kwargs)
+
+
+@pytest.mark.parametrize("dtype", [torch.float, torch.float16])
+@pytest.mark.parametrize(
+    "input_info",
+    [
+        ([32, 64], False, [64, 16], False, 1.0),
+        ([33, 68], False, [18, 68], True, 0.5),
+        ([128, 64], True, [128, 32], False, 0.5),
+        ([123, 234], True, [345, 123], True, -1.0),
+        ([22, 33, 44], False, [44, 55], False, 1.0),
+        ([22, 33, 44], False, [666, 44], True, 0.2),
+        ([22, 33, 44], True, [33, 666], False, -0.2),
+        ([64, 128], False, [16, 64, 128], True, 0.5),
+        ([16, 32, 64], False, [16, 64, 32], False, 1.0),
+        ([8, 16, 32, 16], True, [8, 16, 32, 32], True, 1.0),
+    ],
+)
+def test_matmul(dtype, input_info):
+    pt_inputs = [
+        torch.rand(*input_info[0], dtype=dtype, device=DEVICE),
+        torch.rand(*input_info[2], dtype=dtype, device=DEVICE),
+    ]
+    ort_inputs = copy.deepcopy(pt_inputs)
+    kwargs = {}
+    if input_info[1]:
+        pt_inputs[0] = pt_inputs[0].transpose(-1, -2)
+        kwargs["trans_a"] = True
+    if input_info[3]:
+        pt_inputs[1] = pt_inputs[1].transpose(-1, -2)
+        kwargs["trans_b"] = True
+    if input_info[4] != 1.0:
+        kwargs["alpha"] = input_info[4]
+    pt_output = torch.matmul(*pt_inputs) * input_info[4]
+    alloc_out = random.choice([True, False])
+    if alloc_out:
+        ort_output = torch.empty(pt_output.shape, dtype=dtype, device=DEVICE)
+        ort_inputs.append(ort_output)
+        call_triton_by_name("triton_matmul_out", *[to_dlpack(tensor) for tensor in ort_inputs], **kwargs)
+    else:
+        ort_output = _from_dlpack(
+            call_triton_by_name("triton_matmul", *[to_dlpack(tensor) for tensor in ort_inputs], **kwargs)
+        )
+    rtol = 1e-02 if dtype == torch.float16 else 1e-04
+    atol = 1e-02 if dtype == torch.float16 else 1e-05
+    _test_helpers.assert_values_are_close(pt_output, ort_output, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("dtype", [torch.float, torch.float16])
+@pytest.mark.parametrize(
+    "input_info",
+    [
+        ([64, 32], False, [32, 64], False, [64, 64], 1.0, 0.0),
+        ([65, 129], False, [65, 129], True, [65, 1], 0.5, -0.5),
+        ([127, 63], True, [127, 127], False, [127], -1.0, 0.2),
+        ([256, 64], True, [128, 256], True, [1], 0.2, 0.5),
+    ],
+)
+def test_gemm(dtype, input_info):
+    pt_inputs = [
+        torch.rand(*input_info[0], dtype=dtype, device=DEVICE),
+        torch.rand(*input_info[2], dtype=dtype, device=DEVICE),
+        torch.rand(*input_info[4], dtype=dtype, device=DEVICE),
+    ]
+    ort_inputs = copy.deepcopy(pt_inputs)
+    kwargs = {}
+    if input_info[1]:
+        pt_inputs[0] = pt_inputs[0].transpose(-1, -2)
+        kwargs["trans_a"] = True
+    if input_info[3]:
+        pt_inputs[1] = pt_inputs[1].transpose(-1, -2)
+        kwargs["trans_b"] = True
+    if input_info[5] != 1.0:
+        kwargs["alpha"] = input_info[5]
+    if input_info[6] != 1.0:
+        kwargs["beta"] = input_info[6]
+    pt_output = torch.matmul(pt_inputs[0], pt_inputs[1]) * input_info[5] + pt_inputs[2] * input_info[6]
+    alloc_out = random.choice([True, False])
+    if alloc_out:
+        ort_output = torch.empty(pt_output.shape, dtype=dtype, device=DEVICE)
+        ort_inputs.append(ort_output)
+        call_triton_by_name("triton_gemm_out", *[to_dlpack(tensor) for tensor in ort_inputs], **kwargs)
+    else:
+        ort_output = _from_dlpack(
+            call_triton_by_name("triton_gemm", *[to_dlpack(tensor) for tensor in ort_inputs], **kwargs)
+        )
+    rtol = 1e-02 if dtype == torch.float16 else 1e-04
+    atol = 1e-02 if dtype == torch.float16 else 1e-05
+    _test_helpers.assert_values_are_close(pt_output, ort_output, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
+def test_elementwise_module(dtype):
+    n, d, h, w = 8, 768, 12, 64
+
+    class NeuralNetElementwise(torch.nn.Module):
+        def forward(self, input1, input2, input3, input4):
+            return input1 + input2 - input3 * input4
+
+    def _gen_inputs(dtype):
+        return [
+            torch.rand(n, d, h, w, dtype=dtype, device=DEVICE, requires_grad=True),
+            torch.rand(w, dtype=dtype, device=DEVICE, requires_grad=True),
+            torch.rand(d, 1, 1, dtype=dtype, device=DEVICE, requires_grad=True),
+            torch.rand(n, 1, h, w, dtype=dtype, device=DEVICE, requires_grad=True),
+        ]
+
+    _run_module_test(NeuralNetElementwise, dtype, _gen_inputs, 1)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
+@pytest.mark.parametrize("input_shapes_and_axis", [([2, 3, 3, 3], [3, 3], 2), ([2, 1024], [2, 1024], -1)])
+def test_softmax_module(dtype, input_shapes_and_axis):
+    class NeuralNetSoftmax(torch.nn.Module):
+        def forward(self, input1, input2):
+            return torch.softmax(input1 * input2, dim=input_shapes_and_axis[2])
+
+    def _gen_inputs(dtype):
+        return [
+            torch.randn(*input_shapes_and_axis[0], dtype=dtype, device=DEVICE, requires_grad=True),
+            torch.randn(*input_shapes_and_axis[1], dtype=dtype, device=DEVICE, requires_grad=True),
+        ]
+
+    _run_module_test(NeuralNetSoftmax, dtype, _gen_inputs, 2)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
+@pytest.mark.parametrize(
+    "input_shapes_and_axis", [([2, 1024], [2, 1024], -1), ([2, 2049], [2, 1], -1), ([2, 3, 3, 3], [3, 3], 2)]
+)
+def test_layer_norm_module(dtype, input_shapes_and_axis):
+    pytest.skip("LayerNorm is disabled for now due to perf issue.")
+
+    class NeuralNetLayerNorm(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.layer_norm = torch.nn.LayerNorm(
+                input_shapes_and_axis[0][input_shapes_and_axis[2] :], device=DEVICE, dtype=dtype
+            )
+
+        def forward(self, input1, input2):
+            return self.layer_norm(input1 * input2)
+
+    def _gen_inputs(dtype):
+        return [
+            torch.randn(*input_shapes_and_axis[0], dtype=dtype, device=DEVICE, requires_grad=True),
+            torch.randn(*input_shapes_and_axis[1], dtype=dtype, device=DEVICE, requires_grad=True),
+        ]
+
+    _run_module_test(NeuralNetLayerNorm, dtype, _gen_inputs, 2)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
+@pytest.mark.parametrize("has_sum", [True, False])
+def test_slice_scel_module(dtype, has_sum):
+    class NeuralNetSliceScel(torch.nn.Module):
+        def forward(self, logits, labels):
+            shift_logits = logits[..., :-1, :].contiguous()
+            labels = labels[..., 1:].contiguous()
+            loss_fct = torch.nn.CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))
+            return logits + loss if has_sum else loss
+
+    def _gen_inputs(dtype):
+        return [
+            (torch.rand(4, 8, 16) * 0.01).to(dtype=dtype, device=DEVICE).requires_grad_(True),
+            torch.randint(0, 16, (4, 8), dtype=torch.int64, device=DEVICE),
+        ]
+
+    _run_module_test(NeuralNetSliceScel, dtype, _gen_inputs, 2)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
+@pytest.mark.parametrize("input_shapes", [([128, 64], [64, 64]), ([16, 64, 128], [16, 128, 64])])
+def test_matmul_tunable_op(dtype, input_shapes):
+    class NeuralNetMatmul(torch.nn.Module):
+        def forward(self, input1, input2):
+            return torch.matmul(input1, input2)
+
+    def _gen_inputs(dtype):
+        return [
+            torch.rand(*input_shapes[0], dtype=dtype, device=DEVICE, requires_grad=True),
+            torch.rand(*input_shapes[1], dtype=dtype, device=DEVICE, requires_grad=True),
+        ]
+
+    _run_tunable_op_test(NeuralNetMatmul, dtype, _gen_inputs, "MatMulTunableOp", 2)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
+@pytest.mark.parametrize("m_n_k", [(64, 64, 64)])
+def test_gemm_tunable_op(dtype, m_n_k):
+    class NeuralNetGemm(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = torch.nn.Linear(m_n_k[2], m_n_k[1])
+
+        def forward(self, input):
+            return self.linear(input)
+
+    def _gen_inputs(dtype):
+        return [torch.rand(m_n_k[0], m_n_k[2], dtype=dtype, device=DEVICE, requires_grad=True)]
+
+    _run_tunable_op_test(NeuralNetGemm, dtype, _gen_inputs, "GemmTunableOp", 2)

orttraining/orttraining/training_ops/cpu/triton/triton_op.cc

@@ -0,0 +1,74 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#ifdef ENABLE_TRITON
+
+#include "orttraining/training_ops/cpu/triton/triton_op.h"
+
+#ifndef SHARED_PROVIDER
+#include "core/framework/op_kernel_context_internal.h"
+#endif
+#include "orttraining/core/framework/triton/triton_op_executor.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+InlinedHashSet<size_t> TritonOp::GetBoolOutputs(size_t output_size) const {
+  InlinedHashSet<size_t> bool_outputs;
+  for (size_t i = 0; i < output_size; ++i) {
+    ORT_ENFORCE(i < Node().OutputDefs().size(), "Output index out of range.");
+    if (Node().OutputDefs()[i]->TypeAsProto()->tensor_type().elem_type() ==
+        ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_BOOL) {
+      bool_outputs.insert(i);
+    }
+  }
+  return bool_outputs;
+}
+
+Status TritonOp::Compute(OpKernelContext* context) const {
+  auto* p_ctx_internal = reinterpret_cast<OpKernelContextInternal*>(context);
+  size_t input_size = static_cast<size_t>(p_ctx_internal->InputCount());
+  size_t output_size = static_cast<size_t>(p_ctx_internal->OutputCount());
+  InlinedVector<const OrtValue*> inputs;
+  for (size_t i = 0; i < input_size; ++i) {
+    inputs.emplace_back(p_ctx_internal->GetInputMLValue(static_cast<int>(i)));
+  }
+  InlinedVector<OrtValue> outputs;
+  InlinedHashSet<size_t> bool_outputs = GetBoolOutputs(output_size);
+  auto& executor = training::framework::triton::TritonOpExecutor::Instance();
+  if (func_name_ != "") {
+    executor.ExecuteByFuncName(func_name_, inputs, outputs, bool_outputs);
+  } else {
+    executor.ExecuteByOnnx(onnx_key_, onnx_string_, inputs, outputs, bool_outputs);
+  }
+  ORT_ENFORCE(output_size == outputs.size());
+  for (size_t i = 0; i < output_size; ++i) {
+    ORT_THROW_IF_ERROR(p_ctx_internal->SetOutputMLValue(static_cast<int>(i), outputs[i]));
+  }
+  return Status::OK();
+}
+
+bool IsTritonOpExecutorInitialized() {
+  return training::framework::triton::TritonOpExecutor::Instance().IsInitialized();
+}
+
+Status ExecuteTritonOpByFuncName(OpKernelContext* p_ctx, const std::string& func_name, size_t input_count,
+                                 size_t output_count,
+                                 const InlinedHashMap<std::string, std::pair<std::string, int>>& kwargs) {
+  auto* p_ctx_internal = reinterpret_cast<OpKernelContextInternal*>(p_ctx);
+  InlinedVector<const OrtValue*> inputs;
+  for (size_t i = 0; i < input_count; ++i) {
+    inputs.emplace_back(p_ctx_internal->GetInputMLValue(static_cast<int>(i)));
+  }
+  for (size_t i = 0; i < output_count; ++i) {
+    inputs.emplace_back(p_ctx_internal->GetOutputMLValue(static_cast<int>(i)));
+  }
+  InlinedVector<OrtValue> outputs;
+  training::framework::triton::TritonOpExecutor::Instance().ExecuteByFuncName(func_name, inputs, outputs, {}, kwargs);
+  return Status::OK();
+}
+
+}  // namespace contrib
+}  // namespace onnxruntime
+
+#endif

orttraining/orttraining/training_ops/cpu/triton/triton_op.h

@@ -0,0 +1,41 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#ifdef ENABLE_TRITON
+
+#pragma once
+
+#ifndef SHARED_PROVIDER
+#include "core/framework/op_kernel.h"
+#endif
+
+namespace onnxruntime {
+namespace contrib {
+
+class TritonOp final : public OpKernel {
+ public:
+  TritonOp(const OpKernelInfo& info) : OpKernel(info) {
+    ORT_THROW_IF_ERROR(info.GetAttr("func_name", &func_name_));
+    ORT_THROW_IF_ERROR(info.GetAttr("onnx_key", &onnx_key_));
+    ORT_THROW_IF_ERROR(info.GetAttr("onnx_string", &onnx_string_));
+  }
+
+  Status Compute(OpKernelContext* context) const override;
+
+ private:
+  InlinedHashSet<size_t> GetBoolOutputs(size_t output_size) const;
+
+  std::string func_name_;
+  int64_t onnx_key_;
+  std::string onnx_string_;
+};
+
+bool IsTritonOpExecutorInitialized();
+Status ExecuteTritonOpByFuncName(OpKernelContext* p_ctx, const std::string& func_name, size_t input_count,
+                                 size_t output_count,
+                                 const InlinedHashMap<std::string, std::pair<std::string, int>>& kwargs);
+
+}  // namespace contrib
+}  // namespace onnxruntime
+
+#endif

orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc

@@ -216,6 +216,10 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Reco
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, WaitEvent);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, YieldOp);
 
+#ifdef ENABLE_TRITON
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, TritonOp);
+#endif
+
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, GistBinarizeEncoder);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, GistBinarizeEncoder);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, GistBinarizeEncoder);
@@ -456,6 +460,10 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) {
     BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, WaitEvent)>,
     BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, YieldOp)>,
 
+#ifdef ENABLE_TRITON
+    BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, TritonOp)>,
+#endif
+
     BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, GistBinarizeEncoder)>,
     BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, GistBinarizeEncoder)>,
     BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, GistBinarizeEncoder)>,

orttraining/orttraining/training_ops/cuda/triton/triton_op.cc

@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#ifdef ENABLE_TRITON
+
+#include "core/providers/shared_library/provider_api.h"
+#include "orttraining/training_ops/cpu/triton/triton_op.h"
+#include "core/providers/cuda/cuda_fwd.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+ONNX_OPERATOR_KERNEL_EX(TritonOp, kMSDomain, 1, kCudaExecutionProvider,
+                        (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()),
+                        onnxruntime::contrib::TritonOp);
+
+}  // namespace cuda
+}  // namespace onnxruntime
+
+#endif

setup.py

@@ -529,6 +529,8 @@ if enable_training or enable_training_apis:
                 "onnxruntime.training.ortmodule.torch_cpp_extensions.cpu.torch_interop_utils",
                 "onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.torch_gpu_allocator",
                 "onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.fused_ops",
+                "onnxruntime.training.ort_triton",
+                "onnxruntime.training.ort_triton.kernel",
                 "onnxruntime.training.utils.data",
                 "onnxruntime.training.utils.hooks",
             ]