Integrate eager mode source code into onnxruntime repo (#8584)

* integrate eager mode source codde; build with cmake and integrate the python test

* Adding the python path for importing libraries in the Eager mode

* fix clang break;check if training and python enabled

* handling the linking of torch libraries across multiple platforms

* merge and fix the naming

* add build instruction

Co-authored-by: Abhishek Jindal <abjindal@OrtTrainingDev0.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
Co-authored-by: ajindal1 <abjindal@microsoft.com>

.gitignore

@@ -51,3 +51,4 @@ onnxruntime/python/version_info.py
 /tools/perf_util/send_perf_metrics.iml
 /tools/perf_util/target/classes
 /tools/perf_util/src/main/resources
+/orttraining/orttraining/eager/ort_aten.g.cpp

cmake/CMakeLists.txt

@@ -156,6 +156,9 @@ option(onnxruntime_BUILD_OPSCHEMA_LIB "Build op schema library" ON)
 # option to enable custom operators in onnxruntime-extensions
 option(onnxruntime_ENABLE_EXTENSION_CUSTOM_OPS "Enable custom operators in onnxruntime-extensions" OFF)
 
+# pre-build python path
+option(onnxruntime_PREBUILT_PYTORCH_PATH "Path to pytorch installation dir")
+
 # Single output director for all binaries
 set (RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin CACHE PATH "Single output directory for all binaries.")
 
@@ -1658,6 +1661,12 @@ endif()
 #names in this var must match the directory names under onnxruntime/core/providers
 set(ONNXRUNTIME_TARGETS onnxruntime_common onnxruntime_graph onnxruntime_framework onnxruntime_util onnxruntime_providers onnxruntime_optimizer onnxruntime_session onnxruntime_mlas onnxruntime_flatbuffers)
 if(onnxruntime_ENABLE_EAGER_MODE)
+  if(NOT onnxruntime_ENABLE_TRAINING OR NOT onnxruntime_ENABLE_PYTHON)
+    message(
+        FATAL_ERROR
+        "Option onnxruntime_ENABLE_EAGER_MODE can only be used when onnxruntime_ENABLE_TRAINING and onnxruntime_ENABLE_PYTHON are enabled")
+  endif()
+  add_compile_definitions(ENABLE_EAGER_MODE)
   list(APPEND ONNXRUNTIME_TARGETS onnxruntime_eager)
 endif()
 foreach(target_name ${ONNXRUNTIME_TARGETS})

cmake/onnxruntime_eager.cmake

@@ -20,3 +20,5 @@ set_target_properties(onnxruntime_eager PROPERTIES FOLDER "ONNXRuntime")
 if (onnxruntime_ENABLE_TRAINING)
   target_include_directories(onnxruntime_session PRIVATE ${ORTTRAINING_ROOT})
 endif()
+
+

cmake/onnxruntime_python.cmake

@@ -24,6 +24,24 @@ if(NOT onnxruntime_PYBIND_EXPORT_OPSCHEMA)
   list(REMOVE_ITEM onnxruntime_pybind_srcs  ${ONNXRUNTIME_ROOT}/python/onnxruntime_pybind_schema.cc)
 endif()
 
+if (onnxruntime_ENABLE_EAGER_MODE)
+  list(APPEND CMAKE_PREFIX_PATH ${onnxruntime_PREBUILT_PYTORCH_PATH})
+  find_package(Torch REQUIRED)
+  find_library(TORCH_PYTHON_LIBRARY torch_python PATHS "${TORCH_INSTALL_PREFIX}/lib")
+
+  file(GLOB onnxruntime_eager_extension_srcs CONFIGURE_DEPENDS
+    "${ORTTRAINING_ROOT}/orttraining/eager/*.cpp"
+    )
+
+  if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
+    list(APPEND onnxruntime_eager_extension_srcs 
+              "${ORTTRAINING_ROOT}/orttraining/core/framework/torch/dlpack_python.cc")
+  endif()
+
+  list(APPEND onnxruntime_pybind_srcs 
+              ${onnxruntime_eager_extension_srcs})
+endif()
+
 onnxruntime_add_shared_library_module(onnxruntime_pybind11_state ${onnxruntime_pybind_srcs})
 if(MSVC)
   target_compile_options(onnxruntime_pybind11_state PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:SHELL:--compiler-options /utf-8>" "$<$<NOT:$<COMPILE_LANGUAGE:CUDA>>:/utf-8>")
@@ -75,6 +93,19 @@ if (onnxruntime_ENABLE_TRAINING)
   target_link_libraries(onnxruntime_pybind11_state PRIVATE onnxruntime_training)
 endif()
 
+if (onnxruntime_ENABLE_EAGER_MODE)
+  # todo: this is because the prebuild pytorch may use a different version of protobuf headers.
+  # force the build to find the protobuf headers ort using.
+  target_include_directories(onnxruntime_pybind11_state PRIVATE "${REPO_ROOT}/cmake/external/protobuf/src")
+  target_link_libraries(onnxruntime_pybind11_state PRIVATE onnxruntime_eager ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY})
+  # the ort_aten.g.cpp is generated from tools. currently it has some limitations.
+  # todo: fix this
+  set_source_files_properties("${ORTTRAINING_ROOT}/orttraining/eager/ort_aten.g.cpp" PROPERTIES COMPILE_FLAGS -Wno-unused-parameter)
+  set_source_files_properties("${ORTTRAINING_ROOT}/orttraining/eager/ort_aten.cpp" PROPERTIES COMPILE_FLAGS -Wno-unused-parameter)
+  set_source_files_properties("${ORTTRAINING_ROOT}/orttraining/eager/ort_guard.cpp" PROPERTIES COMPILE_FLAGS -Wno-unused-parameter)
+  set_source_files_properties("${ORTTRAINING_ROOT}/orttraining/eager/ort_tensor.cpp" PROPERTIES COMPILE_FLAGS -Wno-unused-parameter)
+endif()
+
 target_link_libraries(onnxruntime_pybind11_state PRIVATE
     onnxruntime_session
     ${onnxruntime_libs}
@@ -308,6 +339,7 @@ add_custom_command(
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/quantization
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/transformers
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/transformers/test_data/models
+  COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/eager_test
   COMMAND ${CMAKE_COMMAND} -E copy
       ${ONNXRUNTIME_ROOT}/__init__.py
       $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/
@@ -410,6 +442,18 @@ if (onnxruntime_BUILD_UNIT_TESTS)
   )
 endif()
 
+if (onnxruntime_BUILD_UNIT_TESTS AND onnxruntime_ENABLE_EAGER_MODE)
+  file(GLOB onnxruntime_eager_test_srcs CONFIGURE_DEPENDS
+      "${ORTTRAINING_ROOT}/orttraining/eager/test/*.py"
+  )
+  add_custom_command(
+    TARGET onnxruntime_pybind11_state POST_BUILD
+    COMMAND ${CMAKE_COMMAND} -E copy
+        ${onnxruntime_eager_test_srcs}
+        $<TARGET_FILE_DIR:${build_output_target}>/eager_test/
+  )
+endif()
+
 if (onnxruntime_ENABLE_TRAINING)
   add_custom_command(
     TARGET onnxruntime_pybind11_state POST_BUILD

onnxruntime/python/onnxruntime_pybind_state.cc

@@ -1480,6 +1480,10 @@ static struct {
 void addObjectMethodsForTraining(py::module& m);
 #endif
 
+#ifdef ENABLE_EAGER_MODE
+void addObjectMethodsForEager(py::module& m);
+#endif
+
 void CreatePybindStateModule(py::module& m) {
   m.doc() = "pybind11 stateful interface to ONNX runtime";
   RegisterExceptions(m);
@@ -1560,6 +1564,11 @@ void CreatePybindStateModule(py::module& m) {
   addOpSchemaSubmodule(m);
   addOpKernelSubmodule(m);
 #endif
+
+#ifdef ENABLE_EAGER_MODE
+  addObjectMethodsForEager(m);
+#endif
+
 }
 
 // static variable used to create inference session and training session.

orttraining/orttraining/eager/README.md

@@ -0,0 +1,11 @@
+# ONNX Runtime Eager Mode Support for PyTorch
+
+## Build Instrcutions
+
+* Build Pytorch with commit: 0834a368181d18c8fd2429614fafca50fb412ce1
+* Install the pytorch build to your current environment
+* Build onnxruntime, make sure you enable training and python build together with eager mode:
+
+```bash
+./build.sh --enable_training --enalbe_pybind --build_eager_mode
+```

orttraining/orttraining/eager/opgen/onnxgen.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os.path as path
+from sys import argv
+from onnx import defs
+
+out_file = path.join(
+  path.dirname(path.realpath(__file__)),
+  'opgen',
+  'onnxops.py')
+
+onnx_ops = {}
+for schema in defs.get_all_schemas_with_history():
+  key = schema.name.lower()
+  if schema.deprecated:
+    continue
+  if key not in onnx_ops or \
+    onnx_ops[key].since_version < schema.since_version:
+    onnx_ops[key] = schema
+
+with open(out_file, 'wt') as fp:
+  def write(s): fp.write(s)
+  def writeline(s = ''): fp.write(s + '\n')
+
+  writeline(f'# AUTO-GENERATED CODE! - DO NOT EDIT!')
+  writeline(f'# $ python {" ".join(argv)}')
+  writeline()
+
+  writeline('from opgen.generator import ONNXAttr, ONNXOp, AttrType')
+  writeline()
+
+  for op_name, schema in sorted(onnx_ops.items()):
+    writeline(f'class {schema.name}(ONNXOp):')
+    writeline(f'  """')
+    doc_str = schema.doc.strip('\r\n')
+    for doc_line in str.splitlines(doc_str, keepends=False):
+      writeline(f'  {doc_line}')
+    writeline(f'  """')
+    writeline()
+    write('  def __init__(self')
+
+    for input in schema.inputs:
+      write(f', {input.name}')
+
+    if len(schema.attributes) > 0:
+      writeline(',')
+      for i, (k, attr) in enumerate(schema.attributes.items()):
+        write(f'    {attr.name}=None')
+        if i < len(schema.attributes) - 1:
+          writeline(', ')
+
+    writeline('):')
+    write(f'    super().__init__(\'{schema.name}\', {len(schema.outputs)}')
+
+    for input in schema.inputs:
+      write(f', {input.name}')
+
+    if len(schema.attributes) > 0:
+      writeline(',')
+      for i, (k, attr) in enumerate(schema.attributes.items()):
+        write(f'      {attr.name}=ONNXAttr({attr.name}, {attr.type})')
+        if i < len(schema.attributes) - 1:
+          writeline(', ')
+
+    writeline(')')
+    writeline()
+
+  writeline('onnx_ops = {')
+  for i, (op_name, schema) in enumerate(onnx_ops.items()):
+    writeline(f'  \'{op_name}\': {schema.name},')
+  write('}')
+
+print(f'File updated: {out_file}')

orttraining/orttraining/eager/opgen/opgen.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+
+import argparse
+
+from pathlib import Path
+
+from copy import deepcopy
+
+from opgen.generator import \
+  ORTGen as ORTGen, \
+  ONNXOp as ONNXOp, \
+  SignatureOnly as SignatureOnly, \
+  MakeFallthrough as MakeFallthrough
+
+from opgen.onnxops import *
+
+kMSDomain = 'onnxruntime::kMSDomain'
+
+parser = argparse.ArgumentParser(description='Generate ORT ATen operations')
+parser.add_argument('--output_file', default=None, type=str, help='Output file [default to std out]')
+parser.add_argument('--use_preinstalled_torch', action='store_true', help='Use pre-installed torch from the python environment')
+
+args = parser.parse_args()
+
+
+class ReluGrad(ONNXOp):
+  def __init__(self, dY, X):
+    super().__init__('ReluGrad', 1, dY, X)
+    self.domain = kMSDomain
+
+class Gelu(ONNXOp):
+  def __init__(self, X):
+    super().__init__('Gelu', 1, X)
+    self.domain = kMSDomain
+
+class GeluGrad(ONNXOp):
+  def __init__(self, dY, X):
+    super().__init__('GeluGrad', 1, dY, X)
+    self.domain = kMSDomain
+
+ops = {
+  # Hand-Implemented Ops
+  'aten::empty.memory_format': SignatureOnly(),
+  'aten::empty_strided': SignatureOnly(),
+  'aten::zero_': SignatureOnly(),
+  'aten::copy_': SignatureOnly(),
+  'aten::reshape': SignatureOnly(),
+  'aten::view': SignatureOnly(),
+
+  'aten::addmm': Gemm('mat1', 'mat2', 'self', alpha='alpha', beta='beta'),
+  'aten::t': Transpose('self'),
+  'aten::mm': MatMul('self', 'mat2'),
+  'aten::zeros_like': ConstantOfShape(Shape('self')), #the default constant is 0, so don't need to speicify attribute
+
+  'aten::sum.dim_IntList': ReduceSum('self', 'dim', keepdims='keepdim'),
+  'aten::threshold_backward': ReluGrad('grad_output', 'self'),
+
+  'aten::fmod.Scalar': Mod('self', 'other', fmod=1),
+  'aten::fmod.Tensor': Mod('self', 'other', fmod=1),
+
+  'aten::softshrink': Shrink('self', bias='lambd', lambd='lambd'), #yes, bias is set to 'lambd'
+  'aten::hardshrink': Shrink('self', bias=0, lambd='lambd'),
+  'aten::gelu' : Gelu('self'),
+  'aten::gelu_backward' : GeluGrad('grad', 'self')
+}
+
+for binary_op, onnx_op in {
+  'add': Add('self', Mul('alpha', 'other')),
+  'sub': Sub('self', Mul('alpha', 'other')),
+  'mul': Mul('self', 'other'),
+  'div': Div('self', 'other')}.items():
+  for dtype in ['Tensor', 'Scalar']:
+    for variant in ['', '_']:
+      ops[f'aten::{binary_op}{variant}.{dtype}'] = deepcopy(onnx_op)
+
+for unary_op in [
+  'abs','acos','acosh', 'asinh', 'atanh', 'asin', 'atan', 'ceil', 'cos',
+  'cosh', 'erf', 'exp', 'floor', 'isnan', 'log', 'reciprocal', 'neg', 'round',
+  'relu', 'selu', 'sigmoid', 'sin', 'sinh', 'sqrt', 'tan', 'tanh', 'nonzero',
+  'sign', 'min', 'max', 'hardsigmoid', 'isinf', 'det']:
+  aten_name = f'aten::{unary_op}'
+  onnx_op = onnx_ops[unary_op]('self')
+  ops[aten_name] = onnx_op
+  # produce the in-place variant as well for ops that support it
+  if unary_op not in ['isnan', 'nonzero', 'min', 'max', 'isinf', 'det']:
+    ops[f'{aten_name}_'] = onnx_op
+
+ortgen = ORTGen(ops)
+
+import os
+import sys
+
+from opgen.parser import cpp_create_from_file as CPPParser
+from opgen.writer import SourceWriter as SourceWriter
+
+if args.use_preinstalled_torch:
+  import torch
+  regdecs_path = Path(torch.__file__).parent.joinpath('include/ATen/RegistrationDeclarations.h')
+else:
+  regdecs_path = os.path.realpath(os.path.join(
+    os.path.dirname(__file__),
+    '..',
+    '..',
+    '..',
+    'external',
+    'pytorch',
+    'build',
+    'aten',
+    'src',
+    'ATen',
+    'RegistrationDeclarations.h'))
+
+print(f"INFO: Using ATen RegistrationDeclations from: {regdecs_path}")
+output = sys.stdout
+if not args.output_file is None:
+  output = open(args.output_file, 'wt')
+
+with CPPParser(regdecs_path) as parser, SourceWriter(output) as writer:
+  ortgen.run(parser, writer)

orttraining/orttraining/eager/opgen/opgen/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.

orttraining/orttraining/eager/opgen/opgen/ast.py

@@ -0,0 +1,347 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import io
+from typing import TextIO, List, Union
+from opgen.lexer import Token
+
+class Node(object):
+  def __init__(self):
+    self.tokens = []
+
+  def write(self, writer: TextIO):
+    raise NotImplementedError(self.write)
+
+  def __str__(self):
+    writer = io.StringIO()
+    self.write(writer)
+    return writer.getvalue()
+
+#region Syntax List
+
+class SyntaxListMember(Node):
+  def __init__(self, member: Node, trailing_separator: Token = None):
+    super().__init__()
+    self.member = member
+    self.trailing_separator = trailing_separator
+
+  def write(self, writer: TextIO):
+    self.member.write(writer)
+    if self.trailing_separator:
+      writer.write(self.trailing_separator.value)
+      writer.write(" ")
+
+class SyntaxList(Node):
+  open_token: Token
+  members: List[SyntaxListMember]
+  close_token: Token
+
+  def __init__(self):
+    super().__init__()
+    self.open_token = None
+    self.members = []
+    self.close_token = None
+
+  def __iter__(self):
+    return self.members.__iter__()
+  
+  def __getitem__(self, key):
+    return self.members.__getitem__(key)
+  
+  def __len__(self):
+    return len(self.members)
+
+  def append(self, member: Node, trailing_separator: Token):
+    self.members.append(SyntaxListMember(member, trailing_separator))
+
+  def write(self, writer: TextIO):
+    if self.open_token:
+      writer.write(self.open_token.value)
+    for member in self.members:
+      member.write(writer)
+    if self.close_token:
+      writer.write(self.close_token.value)
+
+#endregion
+
+#region Expressions
+
+class Expression(Node): pass
+
+class LiteralExpression(Expression):
+  def __init__(self, token: Token):
+    super().__init__()
+    self.token = token
+
+  def write(self, writer: TextIO):
+    writer.write(self.token.value)
+
+class ArrayExpression(Expression):
+  def __init__(self, elements: SyntaxList):
+    self.elements = elements
+
+#endregion
+
+#region Types
+
+class Type(Node):
+  def _desugar_self(self) -> "Type":
+    return self
+
+  def desugar(self) -> "Type":
+    desugared = self
+    while True:
+      _desugared = desugared._desugar_self()
+      if _desugared == desugared:
+        return desugared
+      desugared = _desugared
+
+class ExpressionType(Type):
+  def __init__(self, expression: Expression):
+    super().__init__()
+    self.expression = expression
+
+  def write(self, writer: TextIO):
+    self.expression.write(writer)
+
+class ConcreteType(Type):
+  def __init__(self, identifier_tokens: Union[Token, List[Token]]):
+    super().__init__()
+    if isinstance(identifier_tokens, Token):
+      self.identifier_tokens = [identifier_tokens]
+    else:
+      self.identifier_tokens = identifier_tokens
+
+  def write(self, writer: TextIO):
+    for identifier_token in self.identifier_tokens:
+      writer.write(identifier_token.value)
+
+class ConstType(Type):
+  def __init__(self, const_token: Token, inner_type: Type):
+    super().__init__()
+    self.const_token = const_token
+    self.inner_type = inner_type
+
+  def write(self, writer: TextIO):
+    writer.write(self.const_token.value)
+    writer.write(" ")
+    self.inner_type.write(writer)
+
+  def _desugar_self(self) -> Type:
+    return self.inner_type
+
+class ReferenceType(Type):
+  def __init__(self, inner_type: Type, reference_token: Token):
+    super().__init__()
+    self.inner_type = inner_type
+    self.reference_token = reference_token
+
+  def write(self, writer: TextIO):
+    self.inner_type.write(writer)
+    writer.write(self.reference_token.value)
+
+  def _desugar_self(self) -> Type:
+    return self.inner_type
+
+class ModifiedType(Type):
+  def __init__(self, base_type: Type):
+    super().__init__()
+    self.base_type = base_type
+
+  def _desugar_self(self) -> Type:
+    return self.base_type
+
+class OptionalType(ModifiedType):
+  def __init__(self, base_type: Type, token: Token):
+    super().__init__(base_type)
+    self.token = token
+
+  def write(self, writer: TextIO):
+    self.base_type.write(writer)
+    writer.write(self.token.value)
+
+class ArrayType(ModifiedType):
+  def __init__(
+    self,
+    base_type: Type,
+    open_token: Token,
+    length_token: Token,
+    close_token: Token):
+    super().__init__(base_type)
+    self.open_token = open_token
+    self.length_token = length_token
+    self.close_token = close_token
+
+  def write(self, writer: TextIO):
+    self.base_type.write(writer)
+    writer.write(self.open_token.value)
+    if self.length_token:
+      writer.write(self.length_token.value)
+    writer.write(self.close_token.value)
+
+class TemplateType(Type):
+  def __init__(
+    self,
+    identifier_tokens: Union[Token, List[Token]],
+    type_arguments: SyntaxList):
+    super().__init__()
+    if isinstance(identifier_tokens, Token):
+      self.identifier_tokens = [identifier_tokens]
+    else:
+      self.identifier_tokens = identifier_tokens
+    self.type_arguments = type_arguments
+
+  def write(self, writer: TextIO):
+    for identifier_token in self.identifier_tokens:
+      writer.write(identifier_token.value)
+    self.type_arguments.write(writer)
+
+class TupleMemberType(Type):
+  def __init__(self, element_type: Type, element_name: Token):
+    super().__init__()
+    self.element_type = element_type
+    self.element_name = element_name
+
+  def write(self, writer: TextIO):
+    self.element_type.write(writer)
+
+  def _desugar_self(self) -> Type:
+    return self.element_name
+
+class TupleType(Type):
+  def __init__(self, elements: SyntaxList):
+    super().__init__()
+    self.elements = elements
+
+  def write(self, writer: TextIO):
+    self.elements.write(writer)
+
+class AliasInfo(Node):
+  before_set: List[str]
+  after_set: List[str]
+  contained_types: List[Type]
+  tokens: List[Token]
+
+  def __init__(self):
+    super().__init__()
+    self.before_set = []
+    self.after_set = []
+    self.contained_types = []
+    self.tokens = []
+    self.is_writable = False
+  
+  def __str__(self):
+    buffer = io.StringIO()
+    self.write(buffer)
+    return buffer.getvalue()
+
+  def __eq__(self, obj):
+    return isinstance(obj, AliasInfo) and str(self) == str(obj)
+
+  def __ne__(self, obj):
+    return not self.__eq__(obj)
+
+  def write(self, writer: TextIO):
+    writer.write("(")
+    writer.write("|".join(self.before_set))
+    if self.is_writable:
+      writer.write("!")
+    writer.write(" -> ")
+    writer.write("|".join(self.after_set))
+    writer.write(")")
+
+class AliasInfoType(Type):
+  def __init__(self, inner_type: Type, alias_info: AliasInfo):
+    super().__init__()
+    self.inner_type = inner_type
+    self.alias_info = alias_info
+    self.inner_type.alias_info = alias_info
+
+  def write(self, writer: TextIO):
+    self.inner_type.write(writer)
+    self.alias_info.write(writer)
+
+  def _desugar_self(self) -> Type:
+    return self.inner_type
+
+class KWArgsSentinelType(Type):
+  def __init__(self, token: Token):
+    super().__init__()
+    self.token = token
+
+  def write(self, writer: TextIO):
+    writer.write(self.token.value)
+
+class TensorType(ConcreteType): pass
+class IntType(ConcreteType): pass
+class FloatType(ConcreteType): pass
+class BoolType(ConcreteType): pass
+class StrType(ConcreteType): pass
+class ScalarType(ConcreteType): pass
+class ScalarTypeType(ConcreteType): pass
+class DimnameType(ConcreteType): pass
+class GeneratorType(ConcreteType): pass
+class TensorOptionsType(ConcreteType): pass
+class LayoutType(ConcreteType): pass
+class DeviceType(ConcreteType): pass
+class MemoryFormatType(ConcreteType): pass
+class QSchemeType(ConcreteType): pass
+class StorageType(ConcreteType): pass
+class ConstQuantizerPtrType(ConcreteType): pass
+class StreamType(ConcreteType): pass
+
+#region Decls
+
+class Decl(Node): pass
+
+class ParameterDecl(Decl):
+  def __init__(
+    self,
+    parameter_type: Type,
+    identifier: Token = None,
+    equals: Token = None,
+    default_value: Expression = None):
+    super().__init__()
+    self.parameter_type = parameter_type
+    self.identifier = identifier
+    self.equals = equals
+    self.default_value = default_value
+
+  def write(self, writer: TextIO):
+    self.parameter_type.write(writer)
+    if self.identifier:
+      writer.write(" ")
+      writer.write(self.identifier.value)
+
+class FunctionDecl(Decl):
+  def __init__(
+    self,
+    identifier: Token,
+    parameters: SyntaxList,
+    return_type: Type = None,
+    semicolon: Token = None,
+    arrow: Token = None):
+    super().__init__()
+    self.is_leaf = False
+    self.identifier = identifier
+    self.return_type = return_type
+    self.parameters = parameters
+    self.semicolon = semicolon
+    self.arrow = arrow
+  
+  def get_parameter(self, identifier: str) -> ParameterDecl:
+    for param in self.parameters:
+      id = param.member.identifier
+      if id and id.value == identifier:
+        return param.member
+    return None
+
+class TranslationUnitDecl(Decl):
+  def __init__(self, decls: List[FunctionDecl]):
+    super().__init__()
+    self.decls = decls
+
+  def __iter__(self):
+    return self.decls.__iter__()
+
+#endregion

orttraining/orttraining/eager/opgen/opgen/generator.py

@@ -0,0 +1,466 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+from typing import Optional, Dict, List, Union
+
+import sys
+import json
+
+import opgen.lexer as lexer
+import opgen.parser as parser
+import opgen.ast as ast
+import opgen.writer as writer
+
+class Outputs:
+  def __init__(self, count: int):
+    self.count = count
+    self.name = None
+
+  def __str__(self):
+    return self.name if self.name else f'<unbound output>'
+
+class AttrType:
+  FLOAT = 'at::ScalarType::Float'
+  FLOATS = '<unsupported:FLOATS>'
+  INT = 'at::ScalarType::Int'
+  INTS = '<unsupported:INTS>'
+  STRING = 'const char*'
+  STRINGS = '<unsupported:STRINGS>'
+  TENSOR = 'at::Tensor'
+
+class ONNXAttr:
+  def __init__(self, value, type: AttrType=None):
+    self.value = value
+    self.type = type
+
+class ONNXOpEvalContext:
+  ops: List['ONNXOp']
+
+  def __init__(self):
+    self.ops = []
+
+  def prepare_outputs(self):
+    for i, op in enumerate(self.ops):
+      op.outputs.name = f'ort_outputs_{i}_{op.name}'
+
+class ONNXOp:
+  def __init__(self,
+    name: str,
+    outputs: int,
+    *inputs: Union[str, Outputs],
+    **attributes: Optional[Union[str, Outputs]]):
+    self.name = name
+    self.outputs = Outputs(outputs)
+    self.inputs = inputs
+    self.attributes = attributes
+    self.domain = None
+
+  def eval(self, ctx: ONNXOpEvalContext):
+    evaluated_inputs = []
+
+    for i in self.inputs:
+      if isinstance(i, ONNXOp):
+        i = i.eval(ctx)
+      evaluated_inputs.append(i)
+
+    self.inputs = evaluated_inputs
+
+    ctx.ops.append(self)
+
+    return self.outputs
+
+class SignatureOnly(ONNXOp):
+  def __init__(self): super().__init__(None, 0)
+
+class MakeFallthrough(ONNXOp):
+  def __init__(self): super().__init__(None, 0)
+
+class FunctionGenerationError(NotImplementedError):
+  def __init__(self, cpp_func: ast.FunctionDecl, message: str):
+    super().__init__(f'{message} (torch: {cpp_func.torch_func.torch_schema})')
+
+class MappedOpFunction:
+  def __init__(
+    self,
+    torch_op_namespace: str,
+    torch_op_name: str,
+    onnx_op: ONNXOp,
+    cpp_func: ast.FunctionDecl,
+    torch_func: ast.FunctionDecl,
+    signature_only: bool,
+    make_fallthrough: bool):
+    self.torch_op_namespace = torch_op_namespace
+    self.torch_op_name = torch_op_name
+    self.onnx_op = onnx_op
+    self.cpp_func = cpp_func
+    self.torch_func = torch_func
+    self.signature_only = signature_only
+    self.make_fallthrough = make_fallthrough
+
+class ORTGen:
+  _mapped_ops: Dict[str, ONNXOp]
+
+  def __init__(
+    self,
+    ops: Optional[Dict[str, ONNXOp]] = None):
+    self._mapped_ops = {}
+    if ops:
+      self.register_many(ops)
+
+  def register(self, aten_name: str, onnx_op: ONNXOp):
+    self._mapped_ops[aten_name] = onnx_op
+
+  def register_many(self, ops: Dict[str, ONNXOp]):
+    for k, v in ops.items():
+      self.register(k, v)
+
+  def run(self, cpp_parser: parser.CPPParser, writer: writer.SourceWriter):
+    self._write_file_prelude(writer)
+
+    generated_funcs = []
+    current_ns = None
+
+    for mapped_func in self._parse_mapped_function_decls(cpp_parser):
+      del self._mapped_ops[mapped_func.torch_func.identifier.value]
+      generated_funcs.append(mapped_func)
+
+      if mapped_func.make_fallthrough:
+        continue
+
+      ns = mapped_func.torch_op_namespace
+      if current_ns and current_ns != ns:
+        current_ns = None
+        writer.pop_namespace()
+      if ns != current_ns:
+        current_ns = ns
+        writer.writeline()
+        writer.push_namespace(ns)
+
+      writer.writeline()
+      writer.writeline(f'// {mapped_func.torch_func.torch_schema}')
+
+      self._write_function_signature(writer, mapped_func.cpp_func)
+      if mapped_func.signature_only:
+        writer.writeline(';')
+      else:
+        writer.writeline(' {')
+        writer.push_indent()
+        self._write_function_body(writer, mapped_func)
+        writer.pop_indent()
+        writer.writeline('}')
+
+    if current_ns:
+      current_ns = None
+      writer.pop_namespace()
+
+    self._write_function_registrations(writer, generated_funcs)
+    self._write_file_postlude(writer)
+
+    if len(self._mapped_ops) > 0:
+      raise Exception('Torch operation(s) could not be parsed for mapping: ' + \
+        ', '.join([f'\'{o}\'' for o in self._mapped_ops.keys()]))
+
+  def _write_file_prelude(self, writer: writer.SourceWriter):
+    writer.writeline('// AUTO-GENERATED CODE! - DO NOT EDIT!')
+    writer.writeline(f'// $ python {" ".join(sys.argv)}')
+    writer.writeline()
+    writer.writeline('#include <torch/extension.h>')
+    writer.writeline()
+    writer.writeline('#include <core/providers/dml/OperatorAuthorHelper/Attributes.h>')
+    writer.writeline()
+    writer.writeline('#include "ort_tensor.h"')
+    writer.writeline('#include "ort_aten.h"')
+    writer.writeline('#include "ort_log.h"')
+    writer.writeline()
+    writer.push_namespace('torch_ort')
+    writer.push_namespace('eager')
+    writer.writeline()
+    writer.writeline('using namespace at;')
+    writer.writeline('using NodeAttributes = onnxruntime::NodeAttributes;')
+
+  def _write_file_postlude(self, writer: writer.SourceWriter):
+    writer.pop_namespaces()
+
+  def _write_function_signature(
+    self,
+    writer: writer.SourceWriter,
+    cpp_func: ast.FunctionDecl):
+    cpp_func.return_type.write(writer)
+    writer.write(f' {cpp_func.identifier.value}(')
+    writer.push_indent()
+    for param_list_member in cpp_func.parameters:
+      writer.writeline()
+      if isinstance(
+        param_list_member.member.parameter_type,
+        ast.KWArgsSentinelType):
+        writer.write('// ')
+      param_list_member.write(writer)
+    writer.pop_indent()
+    writer.write(')')
+
+  def _write_function_body(
+    self,
+    writer: writer.SourceWriter,
+    mapped_func: MappedOpFunction):
+    onnx_op, cpp_func = mapped_func.onnx_op, mapped_func.cpp_func
+
+    assert(len(cpp_func.parameters) > 0)
+
+    return_alias_info = self._get_alias_info(cpp_func.torch_func.return_type)
+    if return_alias_info and not return_alias_info.is_writable:
+      return_alias_info = None
+    in_place_param: ast.ParameterDecl = None
+
+    # Eval the outer ONNX op to produce a topologically ordered list of ops
+    ctx = ONNXOpEvalContext()
+    onnx_op.eval(ctx)
+    ctx.prepare_outputs()
+
+    # Debug Logging
+    log_params = ', '.join([p.member.identifier.value for p \
+      in cpp_func.parameters if p.member.identifier])
+    writer.writeline(f'ORT_LOG_FN({log_params});')
+    writer.writeline()
+
+    # Fetch the ORT invoker from an at::Tensor.device()
+    # FIXME: find the first at::Tensor param anywhere in the signature
+    # instead of simply the first parameter?
+    first_torch_param = cpp_func.torch_func.parameters[0].member
+    if not isinstance(
+      first_torch_param.parameter_type.desugar(),
+      ast.TensorType):
+      raise FunctionGenerationError(
+        cpp_func,
+        'First parameter must be an at::Tensor')
+
+    writer.write('auto& invoker = GetORTInvoker(')
+    writer.write(first_torch_param.identifier.value)
+    writer.writeline('.device());')
+    writer.writeline()
+
+    # FIXME: warn if we have not consumed all torch parameters (either as
+    # an ORT input or ORT attribute).
+
+    # Perform kernel fission on the ATen op to yield a chain of ORT Invokes
+    # e.g. aten::add(x, y, α) -> onnx::Add(x, onnx::Mul(α, y))
+    for onnx_op_index, onnx_op in enumerate(ctx.ops):
+      # Torch -> ORT inputs
+      for op_input in onnx_op.inputs:
+        if isinstance(op_input, Outputs):
+          continue
+        # See if this input is aliased as an in-place tensor
+        cpp_param = cpp_func.get_parameter(op_input)
+        if return_alias_info and cpp_param and \
+          len(cpp_param.torch_param) == 1 and \
+          self._get_alias_info(cpp_param.torch_param[0]) == return_alias_info:
+          in_place_param = cpp_param
+
+        writer.write(f'auto ort_input_{op_input} = ')
+        writer.writeline(f'create_ort_value(invoker, {op_input});')
+
+      # Torch kwargs -> ORT attributes
+      attrs = { k:v for k, v in onnx_op.attributes.items() if v and v.value }
+      if len(attrs) > 0:
+        attrs_arg = 'attrs'
+        writer.writeline()
+        writer.writeline(f'NodeAttributes {attrs_arg}({len(attrs)});')
+
+        for attr_name, attr in attrs.items():
+          writer.write(f'{attrs_arg}["{attr_name}"] = ')
+          writer.writeline('create_ort_attribute(')
+          writer.push_indent()
+          writer.write(f'"{attr_name}", {attr.value}')
+          if attr.type.startswith('at::ScalarType::'):
+            writer.write(f', {attr.type}')
+          elif attr.type != AttrType.STRING:
+            raise FunctionGenerationError(
+              cpp_func,
+              f'Unsure how how to map ONNX op "{onnx_op.name}" attribute ' + 
+              f'"{attr_name}" of type "{attr.type}" to a call to ' +
+              'create_ort_attribute. Please teach generator.py.')
+          writer.writeline(');')
+          writer.pop_indent()
+        attrs_arg = f'&{attrs_arg}'
+      else:
+        attrs_arg = 'nullptr'
+
+      # Outputs vector
+      writer.writeline()
+      writer.write(f'std::vector<OrtValue> {onnx_op.outputs}')
+      writer.writeline(f'({onnx_op.outputs.count});')
+
+      # Perform the invocation
+      writer.writeline()
+      if onnx_op_index == 0:
+        writer.write('auto ')
+      writer.writeline(f'status = invoker.Invoke("{onnx_op.name}", {{')
+      writer.push_indent()
+      for op_input in onnx_op.inputs:
+        if isinstance(op_input, Outputs):
+          if op_input.count != 1:
+            raise FunctionGenerationError(
+              cpp_func,
+              'multiple outputs not supported')
+          op_input = f'{op_input}[0]'
+        else:
+          op_input = f'ort_input_{op_input}'
+        writer.writeline(f'std::move({op_input}),')
+      writer.pop_indent()
+      writer.write(f'}}, {onnx_op.outputs}, {attrs_arg}')
+      if onnx_op.domain:
+        writer.write(f', {onnx_op.domain}')
+      writer.writeline(');')
+      writer.writeline()
+
+      # Assert invocation
+      writer.writeline('if (!status.IsOK())')
+      writer.push_indent()
+      writer.writeline('throw std::runtime_error(')
+      writer.push_indent()
+      writer.writeline('"ORT return failure status:" + status.ErrorMessage());')
+      writer.pop_indent()
+      writer.pop_indent()
+      writer.writeline()
+
+      # We'll potentially return back to Torch from this op
+      return_outputs = onnx_op.outputs
+
+    # TODO: Pick the right "out" Torch parameter; do not assume the first one
+    # TODO: Handle mutliple results
+    # TODO: Assert return type
+
+    if not return_alias_info:
+      writer.writeline('return aten_tensor_from_ort(')
+      writer.push_indent()
+      writer.writeline(f'std::move({return_outputs}[0]),')
+      writer.writeline(f'{first_torch_param.identifier.value}.options());')
+      writer.pop_indent()
+      return
+
+    if not in_place_param:
+      raise Exception(f'"{cpp_func.torch_func.torch_schema}" ' +
+        'has alias info on its return type but no associated parameter')
+
+    writer.writeline(f'copy(invoker, {return_outputs}[0], ort_input_{in_place_param.identifier.value});')
+    writer.writeline(f'return {in_place_param.identifier.value};')
+
+  def _write_function_registrations(
+    self,
+    writer: writer.SourceWriter,
+    generated_funcs: List[MappedOpFunction]):
+    writer.writeline()
+    writer.writeline('TORCH_LIBRARY_IMPL(aten, ORT, m) {')
+    writer.push_indent()
+    writer.writeline('ORT_LOG_DEBUG << "ATen init";')
+
+    for mapped_func in generated_funcs:
+      cpp_func, torch_func = mapped_func.cpp_func, mapped_func.torch_func
+
+      if mapped_func.make_fallthrough:
+        reg_function_arg = 'torch::CppFunction::makeFallthrough()'
+      else:
+        if mapped_func.torch_op_namespace:
+          reg_function_arg = f'{mapped_func.torch_op_namespace}::'
+        else:
+          reg_function_arg = ''
+        reg_function_arg += cpp_func.identifier.value
+
+      writer.write('m.impl(')
+      if not mapped_func.make_fallthrough:
+        reg_function_arg = f'TORCH_FN({reg_function_arg})'
+
+      writer.writeline(f'"{torch_func.identifier.value}", {reg_function_arg});')
+
+    writer.pop_indent()
+    writer.writeline('}')
+    writer.writeline()
+
+  def _get_alias_info(self, torch_type_or_param: Union[ast.Type, ast.ParameterDecl]):
+    if isinstance(torch_type_or_param, ast.ParameterDecl):
+      torch_type = torch_type_or_param.parameter_type
+    else:
+      torch_type = torch_type_or_param
+    return getattr(torch_type.desugar(), 'alias_info', None)
+
+  def _parse_mapped_function_decls(self, cpp_parser: parser.CPPParser):
+    for cpp_func, torch_func in self._parse_function_decls(cpp_parser):
+      torch_op_name = torch_func.identifier.value
+      if torch_op_name not in self._mapped_ops:
+        continue
+      onnx_op = self._mapped_ops[torch_op_name]
+      if not onnx_op:
+        continue
+
+      try:
+        torch_op_namespace = torch_op_name[0:torch_op_name.index('::')]
+        torch_op_name = torch_op_name[len(torch_op_namespace) + 2:]
+      except:
+        torch_op_namespace = None
+
+      cpp_func.identifier.value = torch_op_name.replace('.', '__')
+
+      yield MappedOpFunction(
+        torch_op_namespace,
+        torch_op_name,
+        onnx_op,
+        cpp_func,
+        torch_func,
+        isinstance(onnx_op, SignatureOnly),
+        isinstance(onnx_op, MakeFallthrough))
+
+  def _parse_function_decls(self, cpp_parser: parser.CPPParser):
+    # Parse the C++ declarations
+    tu = cpp_parser.parse_translation_unit()
+
+    # Parse the Torch schema from the JSON comment that follows each C++ decl
+    # and link associated Torch and C++ decls (functions, parameters, returns)
+    for cpp_func in tu:
+      if cpp_func.semicolon and cpp_func.semicolon.trailing_trivia:
+        for trivia in cpp_func.semicolon.trailing_trivia:
+          if trivia.kind == lexer.TokenKind.SINGLE_LINE_COMMENT:
+            yield self._parse_and_link_torch_function_decl(cpp_func, trivia)
+            break
+
+  def _parse_and_link_torch_function_decl(
+    self,
+    cpp_func: ast.FunctionDecl,
+    torch_schema_comment_trivia: lexer.Token):
+    metadata = json.loads(torch_schema_comment_trivia.value.lstrip('//'))
+    schema = metadata['schema']
+
+    schema_parser = parser.torch_create_from_string(schema)
+    schema_parser.set_source_location(cpp_func.semicolon.location)
+    torch_func = schema_parser.parse_function()
+
+    torch_func.torch_schema = schema
+    torch_func.torch_dispatch = metadata['dispatch'] == 'True'
+    torch_func.torch_default = metadata['default'] == 'True'
+
+    cpp_func.torch_func = torch_func
+
+    if cpp_func.return_type:
+      cpp_func.return_type.torch_type = torch_func.return_type
+
+    # Synthesize KWArgsSentinelType in the C++ declaration if we have one
+    for i, torch_param in enumerate([p.member for p in torch_func.parameters]):
+      if isinstance(torch_param.parameter_type, ast.KWArgsSentinelType):
+        cpp_func.parameters.members.insert(i, ast.SyntaxListMember(
+          torch_param,
+          lexer.Token(None, lexer.TokenKind.COMMA, ',')))
+        break
+
+    # Link Torch parameters to their C++ counterparts, special casing
+    # TensorOptions parameters
+    for i, cpp_param in enumerate([p.member for p in cpp_func.parameters]):
+      if not getattr(cpp_param, 'torch_param', None):
+        cpp_param.torch_param = []
+
+      torch_param_range = 1
+      if isinstance(cpp_param.parameter_type.desugar(), ast.TensorOptionsType):
+        torch_param_range = 4
+
+      for j in range(torch_param_range):
+        torch_param = torch_func.parameters[i + j].member
+        cpp_param.torch_param.append(torch_param)
+
+    return cpp_func, torch_func

orttraining/orttraining/eager/opgen/opgen/lexer.py

@@ -0,0 +1,365 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+from enum import Enum
+from abc import ABC
+from typing import List, Optional, Union, Tuple
+
+class SourceLocation(object):
+  def __init__(self,
+    offset: int = 0,
+    line: int = 1,
+    column: int = 1):
+    self.offset = offset
+    self.line = line
+    self.column = column
+
+  def increment_line(self):
+    return SourceLocation(self.offset + 1, self.line + 1, 1)
+
+  def increment_column(self, count: int = 1):
+    return SourceLocation(self.offset + count, self.line, self.column + count)
+
+  def __str__(self) -> str:
+    return f"({self.line},{self.column})"
+
+  def __repr__(self) -> str:
+    return f"({self.offset},{self.line},{self.column})"
+
+  def __eq__(self, other) -> bool:
+    return (
+      self.__class__ == other.__class__ and
+      self.offset == other.offset and
+      self.line == other.line and
+      self.column == other.column)
+
+class TokenKind(Enum):
+  UNKNOWN = 1
+  EOF = 2
+  WHITESPACE = 3
+  SINGLE_LINE_COMMENT = 4
+  MULTI_LINE_COMMENT = 5
+  IDENTIFIER = 6
+  NUMBER = 7
+  STRING = 8
+  OPEN_PAREN = 9
+  CLOSE_PAREN = 10
+  OPEN_BRACKET = 11
+  CLOSE_BRACKET = 12
+  LESS_THAN = 13
+  GREATER_THAN = 14
+  COMMA = 15
+  SEMICOLON = 16
+  COLON = 17
+  DOUBLECOLON = 18
+  AND = 19
+  OR = 20
+  DIV = 21
+  MUL = 22
+  MINUS = 23
+  EQUALS = 24
+  QUESTION_MARK = 25
+  EXCLAIMATION_MARK = 26
+  ARROW = 27
+
+class Token(object):
+  def __init__(self,
+    location: Union[SourceLocation, Tuple[int, int, int]],
+    kind: TokenKind,
+    value: str,
+    leading_trivia: Optional[List["Token"]] = None,
+    trailing_trivia: Optional[List["Token"]] = None):
+    if isinstance(location, tuple) or isinstance(location, list):
+      location = SourceLocation(location[0], location[1], location[2])
+
+    self.location = location
+    self.kind = kind
+    self.value = value
+    self.leading_trivia = leading_trivia
+    self.trailing_trivia = trailing_trivia
+
+  def is_trivia(self) -> bool:
+    return (
+      self.kind == TokenKind.WHITESPACE or
+      self.kind == TokenKind.SINGLE_LINE_COMMENT or
+      self.kind == TokenKind.MULTI_LINE_COMMENT)
+
+  def has_trailing_trivia(self, trivia_kind: TokenKind) -> bool:
+    if not self.trailing_trivia:
+      return False
+    for trivia in self.trailing_trivia:
+      if trivia.kind == trivia_kind:
+        return True
+    return False
+
+  def __str__(self) -> str:
+    return f"{self.location}: [{self.kind}] '{self.value}'"
+
+  def __repr__(self) -> str:
+    rep = f"Token({repr(self.location)},{self.kind}"
+    if self.value:
+      rep += ",\"" + self.value + "\""
+    if self.leading_trivia:
+      rep += f",leading_trivia={self.leading_trivia}"
+    if self.trailing_trivia:
+      rep += f",trailing_trivia={self.trailing_trivia}"
+    return rep + ")"
+
+  def __eq__(self, other) -> bool:
+    return (
+      self.__class__ == other.__class__ and
+      self.location == other.location and
+      self.kind == other.kind and
+      self.value == other.value and
+      self.leading_trivia == other.leading_trivia and
+      self.trailing_trivia == other.trailing_trivia)
+
+class Reader(ABC):
+  def open(self):
+    pass
+
+  def close(self):
+    pass
+
+  def read_char(self) -> str:
+    return None
+
+class FileReader(Reader):
+  def __init__(self, path: str):
+    self.path = path
+
+  def open(self):
+    self.fp = open(self.path)
+
+  def close(self):
+    self.fp.close()
+
+  def read_char(self) -> str:
+    return self.fp.read(1)
+
+class StringReader(Reader):
+  def __init__(self, buffer: str):
+    self.buffer = buffer
+    self.position = 0
+
+  def read_char(self) -> str:
+    if self.position < len(self.buffer):
+      c = self.buffer[self.position]
+      self.position += 1
+      return c
+    return None
+
+class Lexer(object):
+  _peek: str
+  _next_token: Token
+  _first_token_leading_trivia: List[Token]
+
+  char_to_token_kind = {
+    '(': TokenKind.OPEN_PAREN,
+    ')': TokenKind.CLOSE_PAREN,
+    '<': TokenKind.LESS_THAN,
+    '>': TokenKind.GREATER_THAN,
+    '[': TokenKind.OPEN_BRACKET,
+    ']': TokenKind.CLOSE_BRACKET,
+    ',': TokenKind.COMMA,
+    ';': TokenKind.SEMICOLON,
+    '&': TokenKind.AND,
+    '|': TokenKind.OR,
+    '=': TokenKind.EQUALS,
+    '?': TokenKind.QUESTION_MARK,
+    '!': TokenKind.EXCLAIMATION_MARK,
+    '*': TokenKind.MUL
+  }
+
+  def __init__(self, reader: Reader):
+    self._reader = reader
+    self._peek = None
+    self._current_token_location = SourceLocation()
+    self._next_token_location = SourceLocation()
+    self._next_token = None
+    self._first_token_leading_trivia = []
+
+  def __enter__(self):
+    self._reader.open()
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    self._reader.close()
+
+  def _make_token(self, kind: TokenKind, value: str) -> Token:
+    return Token(self._next_token_location, kind, value)
+
+  def _peek_char(self) -> str:
+    if self._peek:
+      return self._peek
+    self._peek = self._reader.read_char()
+    return self._peek
+
+  def _read_char(self) -> str:
+    if self._peek:
+      c = self._peek
+      self._peek = None
+    else:
+      c = self._reader.read_char()
+    if c:
+      self._current_token_location = \
+        self._current_token_location.increment_line() \
+          if c == '\n' \
+          else self._current_token_location.increment_column()
+    return c
+
+  def set_source_location(self, origin: SourceLocation):
+    self._current_token_location = origin
+
+  def lex(self) -> Token:
+    """
+    Lex a single semantic token from the source, gathering into it
+    any trailing whitespace or comment trivia that may follow. The
+    first non-trivia token in the buffer may also have leading trivia
+    attached to it.
+    """
+    token: Token
+    leading_trivia: Optional[List[Token]] = None
+    trailing_trivia: Optional[List[Token]] = None
+
+    while True:
+      token = self._lex_core()
+      if token.is_trivia():
+        if not leading_trivia:
+          leading_trivia = [token]
+        else:
+          leading_trivia.append(token)
+      else:
+        break
+
+    while True:
+      trailing = self._lex_core()
+      if trailing.is_trivia():
+        if not trailing_trivia:
+          trailing_trivia = [trailing]
+        else:
+          trailing_trivia.append(trailing)
+      else:
+        self._next_token = trailing
+        break
+
+    token.leading_trivia = leading_trivia
+    token.trailing_trivia = trailing_trivia
+
+    return token
+
+  def _lex_core(self) -> Token:
+    """Lex a single token from the source including comments and whitespace."""
+    if self._next_token:
+      token = self._next_token
+      self._next_token = None
+      return token
+
+    self._next_token_location = self._current_token_location
+
+    c = self._peek_char()
+    if not c:
+      return self._make_token(TokenKind.EOF, None)
+
+    kind = Lexer.char_to_token_kind.get(c)
+    if kind:
+      return self._make_token(kind, self._read_char())
+
+    if c.isspace():
+      return self._lex_sequence(
+        TokenKind.WHITESPACE,
+        lambda c: c.isspace())
+
+    if self._is_identifier_char(c, first_char = True):
+      return self._lex_sequence(
+        TokenKind.IDENTIFIER,
+        lambda c: self._is_identifier_char(c))
+
+    if c == ':':
+      self._read_char()
+      if self._peek_char() == ':':
+        return self._make_token(TokenKind.DOUBLECOLON, c + self._read_char())
+      else:
+        return self._make_token(TokenKind.COLON, c)
+
+    if c == '/':
+      self._read_char()
+      if self._peek_char() == '/':
+        return self._lex_sequence(
+          TokenKind.SINGLE_LINE_COMMENT,
+          lambda c: c != "\n", "/")
+      elif self._peek_char() == '*':
+        raise NotImplementedError("Multi-line comments not supported")
+      else:
+        return self._make_token(TokenKind.DIV, c)
+
+    if c == '-':
+      self._read_char()
+      p = self._peek_char()
+      if p == '>':
+        return self._make_token(TokenKind.ARROW, c + self._read_char())
+      elif p == '.' or p.isnumeric():
+        return self._lex_number(c)
+      else:
+        return self._make_token(TokenKind.MINUS, c)
+
+    if c == '.' or c.isnumeric():
+      return self._lex_number()
+
+    if c == '"' or c == '\'':
+      return self._lex_string()
+
+    return self._make_token(TokenKind.UNKNOWN, c)
+
+  def _lex_number(self, s: str = "") -> Token:
+    s += self._read_char()
+
+    in_exponent = False
+    have_decimal_separator = '.' in s
+
+    while True:
+      p = self._peek_char()
+      if not p:
+        break
+      if p.isnumeric():
+        s += self._read_char()
+      elif not have_decimal_separator and not in_exponent and p == '.':
+        have_decimal_separator = True
+        s += self._read_char()
+      elif not in_exponent and (p == 'e' or p == 'E'):
+        in_exponent = True
+        s += self._read_char()
+        if self._peek_char() == '-':
+          s += self._read_char()
+      else:
+        break
+
+    return self._make_token(TokenKind.NUMBER, s)
+
+  def _lex_string(self) -> Token:
+    term = self._read_char()
+    s = ""
+    while True:
+      p = self._peek_char()
+      if p == '\\':
+        self._read_char()
+        s += self._read_char()
+      elif p == term:
+        self._read_char()
+        return self._make_token(TokenKind.STRING, s)
+      else:
+        s += self._read_char()
+
+  def _is_identifier_char(self, c: str, first_char = False) -> bool:
+    if c == '_' or c.isalpha():
+      return True
+    return c in [':', '.'] or c.isnumeric() if not first_char else False
+
+  def _lex_sequence(self, kind: TokenKind, predicate, s: str = ""):
+    while True:
+      c = self._read_char()
+      if c:
+        s += c
+      p = self._peek_char()
+      if not p or not predicate(p):
+        return self._make_token(kind, s)

orttraining/orttraining/eager/opgen/opgen/parser.py

@@ -0,0 +1,363 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+from opgen.lexer import *
+from opgen.ast import *
+from typing import List, Tuple, Union, Optional
+
+class UnexpectedTokenError(RuntimeError):
+  def __init__(self, expected: TokenKind, actual: Token):
+    self.expected = expected
+    self.actual = actual
+    super().__init__(f"unexpected token {actual}; expected {expected}")
+
+class ExpectedSyntaxError(RuntimeError):
+  def __init__(self, expected: str, actual: Token = None):
+    super().__init__(f"expected {expected}; actual {actual}")
+
+class ParserBase(object):
+  _peek_queue: List[Token]
+
+  def __init__(self, lexer: Union[Lexer, Reader]):
+    self._own_lexer = False
+    if isinstance(lexer, Reader):
+      self._own_lexer = True
+      lexer = Lexer(lexer)
+    elif not isinstance(lexer, Lexer):
+      raise TypeError("lexer must be a Lexer or Reader")
+    self._lexer = lexer
+    self._peek_queue = []
+
+  def __enter__(self):
+    if self._own_lexer:
+      self._lexer.__enter__()
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    if self._own_lexer:
+      self._lexer.__exit__(exc_type, exc_val, exc_tb)
+
+  def set_source_location(self, origin: SourceLocation):
+    self._lexer.set_source_location(origin)
+
+  def _peek_token(
+    self,
+    kinds: Union[TokenKind, List[TokenKind]] = None,
+    value: str = None,
+    look_ahead: int = 1) -> Optional[Token]:
+    if look_ahead < 1:
+      raise IndexError("look_ahead must be at least 1")
+    if look_ahead >= len(self._peek_queue):
+      for _ in range(look_ahead - len(self._peek_queue)):
+        self._peek_queue = [self._lexer.lex()] + self._peek_queue
+    peek = self._peek_queue[-look_ahead]
+    if not kinds:
+      return peek
+    if not isinstance(kinds, list):
+      kinds = [kinds]
+    for kind in kinds:
+      if peek.kind == kind:
+        if value:
+          return peek if peek.value == value else None
+        return peek
+    return None
+
+  def _read_token(self) -> Token:
+    return self._peek_queue.pop() if self._peek_queue else self._lexer.lex()
+
+  def _expect_token(self, kind: TokenKind) -> Token:
+    token = self._read_token()
+    if token.kind != kind:
+      raise UnexpectedTokenError(kind, token)
+    return token
+
+  def _parse_list(
+    self,
+    open_token_kind: TokenKind,
+    separator_token_kind: TokenKind,
+    close_token_kind: TokenKind,
+    member_parser: callable) -> SyntaxList:
+    syntax_list = SyntaxList()
+    if open_token_kind:
+      syntax_list.open_token = self._expect_token(open_token_kind)
+    while True:
+      if close_token_kind and self._peek_token(close_token_kind):
+        break
+      member = member_parser()
+      if not self._peek_token(separator_token_kind):
+        syntax_list.append(member, None)
+        break
+      syntax_list.append(member, self._read_token())
+    if close_token_kind:
+      syntax_list.close_token = self._expect_token(close_token_kind)
+    return syntax_list
+
+  def parse_translation_unit(self) -> TranslationUnitDecl:
+    decls = []
+    while not self._peek_token(TokenKind.EOF):
+      decls.append(self.parse_function())
+    return TranslationUnitDecl(decls)
+
+  def parse_function_parameter_default_value_expression(self) -> Expression:
+    return self.parse_expression()
+
+  def parse_function_parameter(self) -> ParameterDecl:
+    parameter_type = self.parse_type()
+
+    if not self._peek_token(TokenKind.IDENTIFIER):
+      return ParameterDecl(parameter_type)
+
+    parameter_name = self._read_token()
+
+    if not self._peek_token(TokenKind.EQUALS):
+      return ParameterDecl(parameter_type, parameter_name)
+
+    return ParameterDecl(
+      parameter_type,
+      parameter_name,
+      self._read_token(),
+      self.parse_function_parameter_default_value_expression())
+
+  def parse_function_parameters(self) -> SyntaxList:
+    return self._parse_list(
+      TokenKind.OPEN_PAREN,
+      TokenKind.COMMA,
+      TokenKind.CLOSE_PAREN,
+      self.parse_function_parameter)
+
+  def parse_function(self) -> FunctionDecl:
+    raise NotImplementedError()
+
+  def parse_expression(self) -> Expression:
+    raise NotImplementedError()
+
+  def parse_type(self) -> Type:
+    raise NotImplementedError()
+
+class CPPParser(ParserBase):
+  def parse_function(self) -> FunctionDecl:
+    return_type = self.parse_type()
+    return FunctionDecl(
+      self._expect_token(TokenKind.IDENTIFIER),
+      self.parse_function_parameters(),
+      return_type,
+      semicolon = self._expect_token(TokenKind.SEMICOLON))
+
+  def parse_expression(self) -> Expression:
+    if self._peek_token(TokenKind.IDENTIFIER) or \
+      self._peek_token(TokenKind.NUMBER) or \
+      self._peek_token(TokenKind.STRING):
+      return LiteralExpression(self._read_token())
+    else:
+      raise UnexpectedTokenError("expression", self._peek_token())
+
+  def parse_type(self) -> Type:
+    if self._peek_token(TokenKind.IDENTIFIER, "const"):
+      parsed_type = ConstType(
+        self._read_token(),
+        self.parse_type())
+    elif self._peek_token([TokenKind.IDENTIFIER, TokenKind.DOUBLECOLON]):
+      identifiers = []
+      while True:
+        token = self._peek_token([TokenKind.IDENTIFIER, TokenKind.DOUBLECOLON])
+        if not token:
+          break
+        identifiers.append(self._read_token())
+        if token.has_trailing_trivia(TokenKind.WHITESPACE):
+          break
+      if self._peek_token(TokenKind.LESS_THAN):
+        parsed_type = TemplateType(
+          identifiers,
+          self._parse_list(
+            TokenKind.LESS_THAN,
+            TokenKind.COMMA,
+            TokenKind.GREATER_THAN,
+            self._parse_template_type_argument))
+      elif identifiers[-1].value == "TensorOptions":
+        parsed_type = TensorOptionsType(identifiers)
+      else:
+        parsed_type = ConcreteType(identifiers)
+    else:
+      raise ExpectedSyntaxError("type", self._peek_token())
+
+    while True:
+      if self._peek_token(TokenKind.AND):
+        parsed_type = ReferenceType(
+          parsed_type,
+          self._read_token())
+      else:
+        return parsed_type
+
+  def _parse_template_type_argument(self) -> Type:
+    if self._peek_token(TokenKind.NUMBER):
+      return ExpressionType(self.parse_expression())
+    return self.parse_type()
+
+class TorchParser(ParserBase):
+  def __init__(self, lexer: Union[Lexer, Reader]):
+    super().__init__(lexer)
+    self._next_anonymous_alias_id = 0
+
+  def parse_function(self) -> FunctionDecl:
+    return FunctionDecl(
+      self._expect_token(TokenKind.IDENTIFIER),
+      self.parse_function_parameters(),
+      arrow = self._expect_token(TokenKind.ARROW),
+      return_type = self.parse_type())
+
+  def parse_expression(self) -> Expression:
+    if self._peek_token(TokenKind.NUMBER) or \
+      self._peek_token(TokenKind.IDENTIFIER) or \
+      self._peek_token(TokenKind.STRING):
+      return LiteralExpression(self._read_token())
+    elif self._peek_token(TokenKind.OPEN_BRACKET):
+      return ArrayExpression(self._parse_list(
+        TokenKind.OPEN_BRACKET,
+        TokenKind.COMMA,
+        TokenKind.CLOSE_BRACKET,
+        self.parse_expression))
+    else:
+      raise UnexpectedTokenError("expression", self._peek_token())
+
+  def parse_type(self) -> Type:
+    parsed_type, alias_info = self._parse_type_and_alias()
+    if not alias_info:
+      return parsed_type
+    if isinstance(parsed_type, ModifiedType):
+      parsed_type.base_type = AliasInfoType(parsed_type.base_type, alias_info)
+    else:
+      parsed_type = AliasInfoType(parsed_type, alias_info)
+    return parsed_type
+
+  def _parse_type_and_alias(self) -> Tuple[Type, AliasInfo]:
+    parsed_type: Type = None
+    alias_info: AliasInfo = None
+
+    if self._peek_token(TokenKind.MUL):
+      return (KWArgsSentinelType(self._read_token()), None)
+
+    if self._peek_token(TokenKind.OPEN_PAREN):
+      def parse_tuple_element():
+        element_type, element_alias_info = self._parse_type_and_alias()
+        if alias_info and element_alias_info:
+          alias_info.add_contained_type(element_alias_info)
+        return TupleMemberType(
+          element_type,
+          self._read_token() \
+            if self._peek_token(TokenKind.IDENTIFIER) \
+              else None)
+
+      parsed_type = TupleType(self._parse_list(
+        TokenKind.OPEN_PAREN,
+        TokenKind.COMMA,
+        TokenKind.CLOSE_PAREN,
+        parse_tuple_element))
+    elif self._peek_token(TokenKind.IDENTIFIER, "Tensor"):
+      parsed_type = TensorType(self._read_token())
+      alias_info = self._parse_torch_alias_info()
+    else:
+      parsed_type = self._parse_torch_base_type()
+      alias_info = self._parse_torch_alias_info()
+
+    while True:
+      if self._peek_token(TokenKind.OPEN_BRACKET):
+        parsed_type = ArrayType(
+          parsed_type,
+          self._read_token(),
+          self._read_token() \
+            if self._peek_token(TokenKind.NUMBER) \
+            else None,
+          self._expect_token(TokenKind.CLOSE_BRACKET))
+      elif self._peek_token(TokenKind.QUESTION_MARK):
+        parsed_type = OptionalType(
+          parsed_type,
+          self._read_token())
+      else:
+        return (parsed_type, alias_info)
+
+  def _parse_torch_base_type(self) -> Type:
+    base_type_parsers = {
+      "int": IntType,
+      "float": FloatType,
+      "bool": BoolType,
+      "str": StrType,
+      "Scalar": ScalarType,
+      "ScalarType": ScalarTypeType,
+      "Dimname": DimnameType,
+      "Layout": LayoutType,
+      "Device": DeviceType,
+      "Generator": GeneratorType,
+      "MemoryFormat": MemoryFormatType,
+      "QScheme": QSchemeType,
+      "Storage": StorageType,
+      "ConstQuantizerPtr": ConstQuantizerPtrType,
+      "Stream": StreamType
+    }
+    identifier = self._expect_token(TokenKind.IDENTIFIER)
+    base_type_parser = base_type_parsers.get(identifier.value)
+    if not base_type_parser:
+      raise ExpectedSyntaxError(
+        "|".join(base_type_parsers.keys()),
+        identifier)
+    base_type = base_type_parser(identifier)
+    return base_type
+
+  def _parse_torch_alias_info(self) -> AliasInfo:
+    alias_info = AliasInfo()
+
+    def parse_set(alias_set: List[str]):
+      while True:
+        if self._peek_token(TokenKind.MUL):
+          alias_info.tokens.append(self._read_token())
+          alias_set.append('*')
+        elif '*' not in alias_set:
+          identifier = self._expect_token(TokenKind.IDENTIFIER)
+          alias_info.tokens.append(identifier)
+          alias_set.append(identifier.value)
+        else:
+          raise ExpectedSyntaxError(
+            "alias wildcard * or alias identifier",
+            self._peek_token())
+
+        if self._peek_token(TokenKind.OR):
+          alias_info.tokens.append(self._read_token())
+        else:
+          return
+
+    if self._peek_token(TokenKind.OPEN_PAREN):
+      alias_info.tokens.append(self._read_token())
+
+      parse_set(alias_info.before_set)
+
+      if self._peek_token(TokenKind.EXCLAIMATION_MARK):
+        alias_info.tokens.append(self._read_token())
+        alias_info.is_writable = True
+
+      if self._peek_token(TokenKind.ARROW):
+        alias_info.tokens.append(self._read_token())
+        parse_set(alias_info.after_set)
+      else:
+        # no '->' so assume before and after are identical
+        alias_info.after_set = alias_info.before_set
+
+      alias_info.tokens.append(self._expect_token(TokenKind.CLOSE_PAREN))
+    elif self._peek_token(TokenKind.EXCLAIMATION_MARK):
+      alias_info.is_writable = True
+      alias_info.before_set.append(str(self._next_anonymous_alias_id))
+      self._next_anonymous_alias_id += 1
+    else:
+      return None
+
+    return alias_info
+
+def cpp_create_from_file(path: str) -> CPPParser:
+  return CPPParser(FileReader(path))
+
+def cpp_create_from_string(buffer: str) -> CPPParser:
+  return CPPParser(StringReader(buffer))
+
+def torch_create_from_file(path: str) -> TorchParser:
+  return TorchParser(FileReader(path))
+
+def torch_create_from_string(buffer: str) -> TorchParser:
+  return TorchParser(StringReader(buffer))

orttraining/orttraining/eager/opgen/opgen/writer.py

@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+from typing import TextIO, List
+
+class SourceWriter:
+  _writer: TextIO
+  _indent_str: str
+  _indent_depth: int
+  _needs_indent: bool
+  _namespaces: List[str]
+
+  def __init__(self, base_writer: TextIO, indent_str: str = '  '):
+    self._writer = base_writer
+    self._indent_str = indent_str
+    self._indent_depth = 0
+    self._needs_indent = False
+    self._namespaces = []
+
+  def __enter__(self):
+    self._writer.__enter__()
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    self._writer.__exit__(exc_type, exc_val, exc_tb)
+
+  def write(self, str: str):
+    if not str or len(str) <=0:
+      return
+
+    for c in str:
+      if self._needs_indent:
+        self._needs_indent = False
+        if self._indent_depth > 0:
+          self._writer.write(self._indent_str * self._indent_depth)
+
+      if c == '\n':
+        self._needs_indent = True
+      
+      self._writer.write(c)
+
+  def writeline(self, str: str = None):
+    if str:
+      self.write(str)
+    self.write('\n')
+
+  def push_indent(self):
+    self._indent_depth += 1
+  
+  def pop_indent(self):
+    self._indent_depth -= 1
+
+  def push_namespace(self, namespace: str):
+    self._namespaces.append(namespace)
+    self.writeline(f"namespace {namespace} {{")
+  
+  def pop_namespace(self):
+    self.writeline(f"}} // namespace {self._namespaces.pop()}")
+
+  def pop_namespaces(self):
+    while len(self._namespaces) > 0:
+      self.pop_namespace()

orttraining/orttraining/eager/opgen/opgen_test/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.

orttraining/orttraining/eager/opgen/opgen_test/lexer_test.py

@@ -0,0 +1,87 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import unittest
+
+from opgen.lexer import StringReader, Lexer, Token, TokenKind, SourceLocation
+
+class LexerTestCase(unittest.TestCase):
+  def create_lexer(self, buffer: str) -> Lexer:
+    return Lexer(StringReader(buffer))
+
+  def lex_single(self, buffer: str, expected_kind: TokenKind):
+    lexer = self.create_lexer(buffer)
+    token = lexer.lex()
+    self.assertEqual(expected_kind, token.kind)
+    self.assertIsNone(token.leading_trivia)
+    self.assertIsNone(token.trailing_trivia)
+    eof = lexer.lex()
+    self.assertEqual(TokenKind.EOF, eof.kind)
+    self.assertIsNone(eof.value)
+    self.assertIsNone(eof.leading_trivia)
+    self.assertIsNone(eof.trailing_trivia)
+    return token
+
+  def test_empty(self):
+    lexer = self.create_lexer("")
+    self.assertEqual(lexer.lex().kind, TokenKind.EOF)
+
+  def test_trivia(self):
+    lexer = self.create_lexer("    // hello\nid\r\n// world")
+    self.assertEqual(
+      Token(
+        (13, 2, 1),
+        TokenKind.IDENTIFIER,
+        "id",
+        leading_trivia = [
+          Token((0, 1, 1), TokenKind.WHITESPACE, "    "),
+          Token((4, 1, 5), TokenKind.SINGLE_LINE_COMMENT, "// hello"),
+          Token((12, 1, 13), TokenKind.WHITESPACE, "\n")
+        ],
+        trailing_trivia = [
+          Token((15, 2, 3), TokenKind.WHITESPACE, "\r\n"),
+          Token((17, 3, 1), TokenKind.SINGLE_LINE_COMMENT, "// world")
+        ]),
+      lexer.lex())
+    self.assertEqual(
+      Token(
+        (25, 3, 9),
+        TokenKind.EOF,
+        None),
+      lexer.lex())
+
+  def test_number(self):
+    def assert_number(number):
+      self.assertEqual(
+        number,
+        self.lex_single(
+          number,
+          TokenKind.NUMBER).value)
+
+    for number in [
+      "0", "01", "-1", "-.5", ".5", "0.5", "0e1", "1E10", "-0.5e6",
+      "-0.123e-456", "1234567891011231314151617181920", "-12345.6789E-123456"]:
+      assert_number(number)
+
+    for number in [
+      "1.2.3", "e1", "-e1", "123e0.5"]:
+      self.assertRaises(
+        BaseException,
+        lambda: assert_number(number))
+
+    lexer = self.create_lexer("1.2.3.4e5.6")
+    self.assertEqual(
+      Token((0, 1, 1), TokenKind.NUMBER, "1.2"),
+      lexer.lex())
+    self.assertEqual(
+      Token((3, 1, 4), TokenKind.NUMBER, ".3"),
+      lexer.lex())
+    self.assertEqual(
+      Token((5, 1, 6), TokenKind.NUMBER, ".4e5"),
+      lexer.lex())
+    self.assertEqual(
+      Token((9, 1, 10), TokenKind.NUMBER, ".6"),
+      lexer.lex())
+    self.assertEqual(
+      Token((11, 1, 12), TokenKind.EOF, None),
+      lexer.lex())

orttraining/orttraining/eager/opgen/opgen_test/parser_test.py

@@ -0,0 +1,68 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import unittest
+
+from opgen.lexer import TokenKind
+from opgen.parser import create_from_string as Parser
+
+class ParserLookaheadTests(unittest.TestCase):
+  def test_no_peeks(self):
+    parser = Parser("1 2 3 4 5")
+    self.assertEqual("1", parser._read_token().value)
+    self.assertEqual("2", parser._read_token().value)
+    self.assertEqual("3", parser._read_token().value)
+    self.assertEqual("4", parser._read_token().value)
+    self.assertEqual("5", parser._read_token().value)
+
+  def test_peek_0(self):
+    parser = Parser("1 2 3 4 5")
+    self.assertRaises(IndexError, lambda: parser._peek_token(look_ahead = 0))
+
+  def test_backward_peek_no_reads(self):
+    parser = Parser("1 2 3 4 5")
+    self.assertEqual("1", parser._peek_token(look_ahead = 1).value)
+    self.assertEqual("2", parser._peek_token(look_ahead = 2).value)
+    self.assertEqual("3", parser._peek_token(look_ahead = 3).value)
+    self.assertEqual("4", parser._peek_token(look_ahead = 4).value)
+    self.assertEqual("5", parser._peek_token(look_ahead = 5).value)
+
+  def test_forward_peek_no_reads(self):
+    parser = Parser("1 2 3 4 5")
+    self.assertEqual("5", parser._peek_token(look_ahead = 5).value)
+    self.assertEqual("4", parser._peek_token(look_ahead = 4).value)
+    self.assertEqual("3", parser._peek_token(look_ahead = 3).value)
+    self.assertEqual("2", parser._peek_token(look_ahead = 2).value)
+    self.assertEqual("1", parser._peek_token(look_ahead = 1).value)
+
+  def test_peek_and_read(self):
+    parser = Parser("1 2 3 4 5 6 7 8")
+    self.assertEqual("1", parser._read_token().value)
+    self.assertEqual("2", parser._read_token().value)
+    self.assertEqual("4", parser._peek_token(look_ahead = 2).value)
+    self.assertEqual("3", parser._peek_token(look_ahead = 1).value)
+    self.assertEqual("6", parser._peek_token(look_ahead = 4).value)
+    self.assertEqual("3", parser._read_token().value)
+    self.assertEqual("4", parser._read_token().value)
+    self.assertEqual("5", parser._peek_token(look_ahead = 1).value)
+    self.assertEqual("8", parser._peek_token(look_ahead = 4).value)
+    self.assertEqual("5", parser._read_token().value)
+    self.assertEqual("6", parser._read_token().value)
+    self.assertEqual("7", parser._read_token().value)
+    self.assertEqual("8", parser._peek_token(look_ahead = 1).value)
+
+  def test_peek_value_and_read(self):
+    parser = Parser("1 2 3 4 5 6 7 8")
+    self.assertEqual("1", parser._read_token().value)
+    self.assertEqual("2", parser._read_token().value)
+    self.assertIsNotNone(parser._peek_token(TokenKind.NUMBER, "4", look_ahead = 2))
+    self.assertIsNotNone(parser._peek_token(TokenKind.NUMBER, "3", look_ahead = 1))
+    self.assertIsNotNone(parser._peek_token(TokenKind.NUMBER, "6", look_ahead = 4))
+    self.assertEqual("3", parser._read_token().value)
+    self.assertEqual("4", parser._read_token().value)
+    self.assertIsNotNone(parser._peek_token(TokenKind.NUMBER, "5", look_ahead = 1))
+    self.assertIsNotNone(parser._peek_token(TokenKind.NUMBER, "8", look_ahead = 4))
+    self.assertEqual("5", parser._read_token().value)
+    self.assertEqual("6", parser._read_token().value)
+    self.assertEqual("7", parser._read_token().value)
+    self.assertIsNotNone(parser._peek_token(TokenKind.NUMBER, "8", look_ahead = 1))

orttraining/orttraining/eager/ort_aten.cpp

@@ -0,0 +1,291 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "ort_aten.h"
+#include "ort_tensor.h"
+
+namespace torch_ort {
+namespace eager {
+
+//#pragma region Helpers
+
+namespace {
+  inline bool is_device_supported(at::DeviceType type) {
+    return type == at::kORT || type == at::kCPU;
+  }
+
+  inline void assert_tensor_supported(const at::Tensor& tensor) {
+    if (tensor.is_sparse()) {
+      throw std::runtime_error("ORT copy: sparse not supported");
+    }
+
+    if (tensor.is_quantized()) {
+      throw std::runtime_error("ORT copy: quantized not supported");
+    }
+
+    if (!is_device_supported(tensor.device().type())) {
+      throw std::runtime_error("ORT copy: device not supported");
+    }
+  }
+}
+
+at::Tensor aten_tensor_from_ort(
+  OrtValue&& ot,
+  const at::TensorOptions& options) {
+  return at::Tensor(c10::make_intrusive<ORTTensorImpl>(
+    std::move(ot),
+    options));
+}
+
+onnxruntime::MLDataType ort_scalar_type_from_aten(
+  at::ScalarType dtype) {
+  switch (dtype){
+    case at::kFloat:
+      return onnxruntime::DataTypeImpl::GetType<float>();
+    case at::kDouble:
+      return onnxruntime::DataTypeImpl::GetType<double>();
+    case at::kHalf:
+      return onnxruntime::DataTypeImpl::GetType<onnxruntime::MLFloat16>();
+    case at::kBFloat16:
+      return onnxruntime::DataTypeImpl::GetType<onnxruntime::BFloat16>();
+    case at::kInt:
+      return onnxruntime::DataTypeImpl::GetType<int>();
+    case at::kShort:
+      return onnxruntime::DataTypeImpl::GetType<int16_t>();
+    case at::kLong:
+      return onnxruntime::DataTypeImpl::GetType<int64_t>();
+    default:
+      ORT_THROW("Unsupport aten scalar type: ", dtype);
+  }
+}
+
+OrtValue create_ort_value(
+  onnxruntime::ORTInvoker& invoker,
+  const at::Scalar& scalar) {
+  // TODO: support more types
+  float val = scalar.toFloat();
+  OrtValue ort_val;
+  CreateMLValue(
+    invoker.GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault),
+    ort_scalar_type_from_aten(at::kFloat),
+    {},
+    &ort_val);
+  auto* ort_tensor = ort_val.GetMutable<onnxruntime::Tensor>();
+  CopyVectorToTensor<float>(invoker, {val}, *ort_tensor);
+  return ort_val;
+}
+
+OrtValue create_ort_value(
+  onnxruntime::ORTInvoker& invoker,
+  const at::Tensor& tensor) {
+  assert_tensor_supported(tensor);
+
+  auto* impl = dynamic_cast<ORTTensorImpl*>(tensor.unsafeGetTensorImpl());
+  if (impl) {
+    return impl->tensor();
+  }
+
+  OrtValue ort_tensor;
+  CreateMLValue(
+    tensor.data_ptr(),
+    ort_scalar_type_from_aten(tensor.scalar_type()),
+    tensor.sizes().vec(),
+    &ort_tensor);
+  return ort_tensor;
+}
+
+OrtValue create_ort_value(const at::Tensor& tensor){
+  auto& invoker = GetORTInvoker(tensor.device());
+  return create_ort_value(invoker, tensor);
+}
+
+onnx::AttributeProto create_ort_attribute(
+  const char* name,
+  at::Scalar value) {
+  return create_ort_attribute(name, value, value.type());
+}
+
+onnx::AttributeProto create_ort_attribute(
+  const char* name,
+  at::Scalar value,
+  at::ScalarType type) {
+  onnx::AttributeProto attr;
+  attr.set_name(name);
+  switch (type) {
+    case at::ScalarType::Float:
+    case at::ScalarType::Double:
+      attr.set_type(onnx::AttributeProto_AttributeType::AttributeProto_AttributeType_FLOAT);
+      attr.set_f(value.to<double>());
+      break;
+    case at::ScalarType::Bool:
+    case at::ScalarType::Int:
+    case at::ScalarType::Long:
+      attr.set_type(onnx::AttributeProto_AttributeType::AttributeProto_AttributeType_INT);
+      attr.set_i(value.to<int64_t>());
+      break;
+    default:
+      // For most at::ScalarType, it should be safe to just call value.to<>
+      // on it, but for now we want to explicitly know when we've encountered
+      // a new scalar type while bringing up ORT eager mode.
+      ORT_THROW("Unsupported: at::ScalarType::", value.type());
+  }
+
+  return attr;
+}
+
+onnx::AttributeProto create_ort_attribute(
+  const char* name,
+  const char* value) {
+  onnx::AttributeProto attr;
+  attr.set_name(name);
+  attr.set_type(onnx::AttributeProto_AttributeType::AttributeProto_AttributeType_STRING);
+  attr.set_s(value);
+  return attr;
+}
+
+//#pragma endregion
+
+//#pragma region Hand-Implemented ATen Ops
+
+namespace aten {
+
+at::Tensor empty__memory_format(
+  at::IntArrayRef size,
+  // *,
+  c10::optional<at::ScalarType> dtype_opt,
+  c10::optional<at::Layout> layout_opt,
+  c10::optional<at::Device> device_opt,
+  c10::optional<bool> pin_memory,
+  c10::optional<at::MemoryFormat> memory_format) {
+  ORT_LOG_FN(size, dtype_opt, layout_opt, device_opt, pin_memory, memory_format);
+
+  assert(dtype_opt.has_value());
+  assert(device_opt.has_value());
+  assert(!layout_opt.has_value());
+
+  // TODO: validate options and memory format
+  // TODO: figure out how to get the correct element type.
+  OrtValue ot;
+  auto& invoker = GetORTInvoker(*device_opt);
+  CreateMLValue(
+    invoker.GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault),
+    ort_scalar_type_from_aten(at::kFloat),
+    size.vec(),
+    &ot);
+
+  return aten_tensor_from_ort(
+    std::move(ot),
+    at::TensorOptions()
+      .device(*device_opt)
+      .dtype(*dtype_opt));
+}
+
+at::Tensor empty_strided(
+  at::IntArrayRef size,
+  at::IntArrayRef stride,
+  // *
+  c10::optional<at::ScalarType> dtype_opt,
+  c10::optional<at::Layout> layout_opt,
+  c10::optional<at::Device> device_opt,
+  c10::optional<bool> pin_memory_opt) {
+  ORT_LOG_FN(size, stride, dtype_opt, layout_opt, device_opt, pin_memory_opt);
+
+  // TODO: handle stride
+  // TODO: how to handle type conversion
+  OrtValue ot;
+  assert(device_opt.has_value());
+  // TODO: how to support layout
+  //assert(!layout_opt.has_value());
+  at::ScalarType dtype = c10::dtype_or_default(dtype_opt);
+  auto& invoker = GetORTInvoker(*device_opt);
+  CreateMLValue(
+    invoker.GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault),
+    ort_scalar_type_from_aten(dtype),
+    size.vec(),
+    &ot);
+  return aten_tensor_from_ort(
+    std::move(ot),
+    at::TensorOptions()
+      .device(*device_opt)
+      .dtype(dtype));
+}
+
+at::Tensor reshape(at::Tensor const& self, at::IntArrayRef shape) {
+  ORT_LOG_FN(self, shape);
+
+  auto& invoker = GetORTInvoker(self.device());
+  return aten_tensor_from_ort(
+    reshape_copy(
+      invoker,
+      create_ort_value(invoker, self),
+      shape.vec()),
+    self.options());
+}
+
+at::Tensor view(const at::Tensor& self, at::IntArrayRef size) {
+  ORT_LOG_FN(self, size);
+
+  auto& invoker = GetORTInvoker(self.device());
+  return aten_tensor_from_ort(
+    reshape_copy(
+      invoker,
+      create_ort_value(invoker, self),
+      at::infer_size(
+        size,
+        self.numel())),
+    self.options());
+}
+
+at::Tensor& copy_(
+  at::Tensor& self,
+  const at::Tensor& src,
+  bool non_blocking) {
+  ORT_LOG_FN(self, src, non_blocking);
+
+  assert_tensor_supported(self);
+  assert_tensor_supported(src);
+
+  auto& invoker = GetORTInvoker(self.device().type() == at::kORT
+    ? self.device()
+    : src.device());
+  const auto ort_src = create_ort_value(invoker, src);
+  auto ort_self = create_ort_value(invoker, self);
+
+  copy(invoker, ort_src, ort_self);
+
+  return self;
+}
+
+at::Tensor& zero_(at::Tensor& self){
+  auto& invoker = GetORTInvoker(self.device());
+  auto ort_in_self = create_ort_value(invoker, self);
+  OrtValue flag_val;
+  //construct a constant tensor
+  auto element_type = onnxruntime::DataTypeImpl::GetType<int64_t>();
+  CreateMLValue(invoker.GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault),
+                element_type, {}, &flag_val);
+  auto* ort_flag_tensor = flag_val.GetMutable<onnxruntime::Tensor>();
+  CopyVectorToTensor<int64_t>(invoker, {1}, *ort_flag_tensor);
+
+  std::vector<OrtValue> ort_out(1);
+
+  auto status = invoker.Invoke(
+    "ZeroGradient", {
+      std::move(ort_in_self),
+      std::move(flag_val)
+    }, ort_out, nullptr, onnxruntime::kMSDomain);
+
+  if (!status.IsOK())
+    throw std::runtime_error(
+      "ORT return failure status:" + status.ErrorMessage());
+
+  copy(invoker, ort_out[0], ort_in_self);
+  return self;
+}
+
+} // namespace aten
+
+//#pragma endregion
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/ort_aten.h

@@ -0,0 +1,73 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <torch/extension.h>
+#include <core/framework/ml_value.h>
+
+#include "ort_util.h"
+#include "ort_ops.h"
+#include "ort_log.h"
+
+namespace torch_ort {
+namespace eager {
+
+at::Tensor aten_tensor_from_ort(
+  OrtValue&& ot,
+  const at::TensorOptions& options);
+
+onnxruntime::MLDataType ort_scalar_type_from_aten(
+  at::ScalarType dtype);
+
+OrtValue create_ort_value(
+  onnxruntime::ORTInvoker& invoker,
+  const at::Scalar& scalar);
+
+OrtValue create_ort_value(
+  onnxruntime::ORTInvoker& invoker,
+  const at::Tensor& tensor);
+
+OrtValue create_ort_value(const at::Tensor& tensor);
+
+template<typename T>
+OrtValue create_ort_value(
+  onnxruntime::ORTInvoker& invoker, 
+  const std::vector<T> values) {
+  OrtValue ort_value;
+  CreateMLValue(
+    invoker.GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault),
+    onnxruntime::DataTypeImpl::GetType<T>(),
+    {(int64_t)values.size(),},
+    &ort_value);
+  CopyVectorToTensor<T>(
+    invoker,
+    values,
+    *ort_value.GetMutable<onnxruntime::Tensor>());
+  return ort_value;
+}
+
+template<typename T>
+OrtValue create_ort_value(
+  onnxruntime::ORTInvoker& invoker, 
+  const at::ArrayRef<T> values) {
+  std::vector<T> values_vector;
+  values_vector.assign(values.begin(), values.end());
+  return create_ort_value(invoker, values_vector);
+}
+
+onnx::AttributeProto create_ort_attribute(
+  const char* name,
+  at::Scalar value);
+
+onnx::AttributeProto create_ort_attribute(
+  const char* name,
+  at::Scalar value,
+  at::ScalarType type);
+
+onnx::AttributeProto create_ort_attribute(
+  const char* name,
+  const char* value);
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/ort_backends.cpp

@@ -0,0 +1,131 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include <core/providers/cpu/cpu_execution_provider.h>
+#include <core/providers/cpu/cpu_provider_factory_creator.h>
+#include <core/common/logging/sinks/clog_sink.h>
+#include <core/providers/get_execution_providers.h>
+#include "ort_backends.h"
+#include "ort_log.h"
+#include "core/platform/env.h"
+#include "core/providers/shared_library/provider_host_api.h"
+
+
+#ifdef USE_MSNPU
+  namespace onnxruntime {
+    std::unique_ptr<onnxruntime::IExecutionProvider> CreateMSNPU_ExecutionProvider();
+  }
+#endif
+
+//use the environment from python module
+namespace onnxruntime{
+namespace python{
+  onnxruntime::Environment& GetEnv();
+}
+}
+
+namespace torch_ort {
+namespace eager {
+
+using namespace onnxruntime;
+
+constexpr const char* kMSNPUExecutionProvider = "MSNPUExecutionProvider";
+
+ORTBackendsManager& GetORTBackendsManager() {
+  auto& env = onnxruntime::python::GetEnv();
+  static ORTBackendsManager instance {env.GetLoggingManager()->DefaultLogger()};
+  return instance;
+}
+
+onnxruntime::ORTInvoker& GetORTInvoker(const at::Device device) {
+  return GetORTBackendsManager().GetInvoker(device);
+}
+
+ORTBackendsManager::ORTBackendsManager(const onnxruntime::logging::Logger& logger): logger_(logger){
+  // set device index 0 to cpu EP as default backend.
+  auto status = set_device(0, kCpuExecutionProvider, {});
+  if (!status.IsOK()){
+    throw std::runtime_error("Init CPU device failed: " + status.ErrorMessage());
+  }
+}
+
+void ORTBackendsManager::RegisterProviderLib(const std::string& provider_type, const std::string& lib_path){
+  additional_provider_libs_.insert({provider_type, lib_path});
+}
+
+onnxruntime::Status ORTBackendsManager::set_device(size_t device_index, const std::string& provider_type,
+                                 const ProviderOptions& provider_options){
+  // query avalible device
+  auto& available_providers = GetAvailableExecutionProviderNames();
+  std::unique_ptr<IExecutionProvider> provider_p;
+  if (std::find(available_providers.begin(), available_providers.end(), provider_type) != available_providers.end()){
+    if (provider_type == kCpuExecutionProvider){
+      provider_p = onnxruntime::CreateExecutionProviderFactory_CPU(0)->CreateProvider();
+    }else if (provider_type == kMSNPUExecutionProvider){
+#ifdef USE_MSNPU
+      provider_p = onnxruntime::CreateMSNPU_ExecutionProvider();
+#else
+      return onnxruntime::Status(common::StatusCategory::ONNXRUNTIME,
+                          common::StatusCode::INVALID_ARGUMENT, 
+                          "Execution provider: " + std::string(kMSNPUExecutionProvider) + " is not supported.");
+#endif
+    }
+  }
+  else{
+    auto shared_lib_path_it = additional_provider_libs_.find(provider_type);
+    if (shared_lib_path_it == provider_options.end()){
+      return onnxruntime::Status(common::StatusCategory::ONNXRUNTIME,
+                          common::StatusCode::INVALID_ARGUMENT, 
+                          "Execution provider: " + provider_type + " is not supported.");
+    }
+
+    void* handle;
+    auto error = Env::Default().LoadDynamicLibrary(shared_lib_path_it->second, false, &handle);
+    if (!error.IsOK()) {
+      return onnxruntime::Status(common::StatusCategory::ONNXRUNTIME,
+                                 common::StatusCode::INVALID_ARGUMENT, 
+                                 "Load shared execution provider: " + provider_type + " failed: "
+                                 + error.ErrorMessage());
+    }
+
+    Provider* (*PGetProvider)();
+    Env::Default().GetSymbolFromLibrary(handle, "GetProvider", (void**)&PGetProvider);
+
+    Provider* provider = PGetProvider();
+    std::shared_ptr<IExecutionProviderFactory> ep_factory = provider->CreateExecutionProviderFactory(&provider_options);
+    provider_p = ep_factory->CreateProvider();
+  }
+
+
+  auto invoker = 
+  std::make_unique<onnxruntime::ORTInvoker>(
+    std::move(provider_p),
+    logger_,
+    custom_op_schema_);
+
+  backends_[device_index] = std::move(invoker);
+  return onnxruntime::Status::OK();
+}
+
+onnxruntime::ORTInvoker& ORTBackendsManager::GetInvoker(const at::Device device) {
+  ORT_LOG_FN(device);
+
+  auto device_index = 0;
+  if (device.has_index()) {
+    device_index = device.index();
+  }
+
+  TORCH_CHECK(device.type() == at::DeviceType::ORT, "must be an ORT device");
+  TORCH_CHECK(device_index >= 0, "must have a valid index");
+
+  auto lookup = backends_.find(device_index);
+  if (lookup != backends_.end()) {
+    return *lookup->second;
+  }else{
+    throw std::runtime_error("ORT device index: " + std::to_string(device_index) + " not initialized, \
+                              please use 'torch_ort.set_device' to initialize it first.");
+  }
+}
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/ort_backends.h

@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <torch/extension.h>
+#include <core/framework/ml_value.h>
+#include "core/common/status.h"
+#include <core/framework/provider_options.h>
+#include <core/eager/ort_kernel_invoker.h>
+#include <core/graph/schema_registry.h>
+#include "onnx/defs/schema.h"
+#include <core/graph/model.h>
+
+namespace torch_ort {
+namespace eager {
+
+class ORTBackendsManager {
+public:
+  ORTBackendsManager(const onnxruntime::logging::Logger& logger);
+
+  void RegisterProviderLib(const std::string& provider_type, const std::string& lib_path);
+
+  onnxruntime::Status set_device(size_t device_index, const std::string& provider_type,
+                                 const onnxruntime::ProviderOptions& provider_options);
+
+  onnxruntime::ORTInvoker& GetInvoker(const at::Device device);
+
+private:
+  std::map<at::DeviceIndex, std::unique_ptr<onnxruntime::ORTInvoker>> backends_;
+  const onnxruntime::logging::Logger& logger_;
+  //custom op schema registry
+  //TODO: we might want to support load custom op schema on the fly
+  onnxruntime::IOnnxRuntimeOpSchemaRegistryList custom_op_schema_ = {};
+  std::unordered_map<std::string, std::string> additional_provider_libs_ = {};
+};
+
+ORTBackendsManager& GetORTBackendsManager();
+
+onnxruntime::ORTInvoker& GetORTInvoker(const at::Device device);
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/ort_eager.cpp

@@ -0,0 +1,68 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include <torch/extension.h>
+#include "torch/csrc/autograd/python_variable.h"
+#include "ort_backends.h"
+#include "ort_log.h"
+#include "ort_aten.h"
+#include "ort_backends.h"
+#include "orttraining/core/framework/ortmodule_graph_builder.h"
+#include "python/onnxruntime_pybind_state_common.h"
+#include "orttraining/core/framework/torch/dlpack_python.h"
+#include <core/session/provider_bridge_ort.h>
+
+namespace onnxruntime{
+namespace python{
+
+using namespace onnxruntime::training;
+using namespace torch_ort::eager;
+
+py::object ORTTensor_toDLPack(const at::Tensor& data)
+{
+  OrtValue ort_value = torch_ort::eager::create_ort_value(data);
+  return py::reinterpret_steal<py::object>(onnxruntime::training::framework::torch::ToDlpack(ort_value));
+}
+
+at::Tensor ORTTensor_FromDLPack(const py::object& dlpack_tensor)
+{
+  OrtValue ort_value = onnxruntime::training::framework::torch::FromDlpack(dlpack_tensor.ptr(), false);
+  return torch_ort::eager::aten_tensor_from_ort(
+    std::move(ort_value),
+    at::TensorOptions()
+      .device(at::Device(at::DeviceType::ORT, 0)));
+}
+
+void addObjectMethodsForEager(py::module& m){
+  ORT_LOG_DEBUG << "pybind11 module init";
+
+  m.def(
+    "device",
+    [](int device_index) {
+      return py::cast<py::object>(
+        THPDevice_New(at::Device(at::DeviceType::ORT, device_index)));
+    },
+    py::arg("device_index") = 0);
+  
+  m.def("ort_to_dlpack", [](at::Tensor data) {
+    return ORTTensor_toDLPack(data);
+  });
+  m.def("ort_from_dlpack", [](py::object dlpack_tensor) {
+    return ORTTensor_FromDLPack(dlpack_tensor);
+  });
+
+  m.def("_register_provider_lib", [](const std::string& name, const std::string& provider_shared_lib_path ) {
+    torch_ort::eager::GetORTBackendsManager().RegisterProviderLib(name, provider_shared_lib_path);
+  });
+
+  m.def("set_device", [](size_t device_index, 
+                                          const std::string& provider_type,
+                                          const std::unordered_map<std::string, std::string>& arguments){
+      auto status = torch_ort::eager::GetORTBackendsManager().set_device(device_index, provider_type, arguments);
+      if (!status.IsOK())
+        throw std::runtime_error(status.ErrorMessage());
+    });
+}
+
+}
+}

orttraining/orttraining/eager/ort_guard.cpp

@@ -0,0 +1,112 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include <torch/extension.h>
+
+#include "ort_backends.h"
+#include "ort_log.h"
+
+namespace torch_ort {
+namespace eager {
+
+struct ORTGuardImpl final : public c10::impl::DeviceGuardImplInterface {
+  ORTGuardImpl() {
+    ORT_LOG_FN();
+  }
+
+  explicit ORTGuardImpl(at::DeviceType t) {
+    ORT_LOG_FN(t);
+    AT_ASSERT(t == at::DeviceType::ORT);
+  }
+
+  at::DeviceType type() const override {
+    ORT_LOG_FN();
+    return at::DeviceType::ORT;
+  }
+
+  at::Device exchangeDevice(at::Device d) const override {
+    ORT_LOG_FN(d);
+    AT_ASSERT(d.type() == at::DeviceType::ORT);
+    auto old_device = getDevice();
+    if (old_device.index() != d.index()) {
+      current_device_ = d.index();
+    }
+    return old_device;
+  }
+
+  at::Device getDevice() const override {
+    ORT_LOG_FN();
+    return at::Device(at::DeviceType::ORT, current_device_);
+  }
+
+  void setDevice(at::Device d) const override {
+    ORT_LOG_FN(d);
+    AT_ASSERT(d.type() == at::DeviceType::ORT);
+    AT_ASSERT(d.index() >= 0);
+    current_device_ = d.index();
+  }
+
+  void uncheckedSetDevice(at::Device d) const noexcept override {
+    ORT_LOG_FN(d);
+    current_device_ = d.index();
+  }
+
+  at::Stream getStream(at::Device d) const noexcept override {
+    ORT_LOG_FN(d);
+    return at::Stream(at::Stream::UNSAFE, d, current_streams_[d.index()]);
+  }
+
+  at::Stream exchangeStream(at::Stream s) const noexcept override {
+    ORT_LOG_FN(s);
+    auto old_id = current_streams_[s.device_index()];
+    current_streams_[s.device_index()] = s.id();
+    return at::Stream(at::Stream::UNSAFE, s.device(), old_id);
+  }
+
+  at::DeviceIndex deviceCount() const noexcept override {
+    ORT_LOG_FN();
+    return 1;
+  }
+
+//  #pragma region events
+
+  #define EVENTS_NIEX TORCH_CHECK(false, "ORT backend doesn't support events.")
+
+  void record(void** event,
+    const at::Stream& stream,
+    const at::DeviceIndex device_index,
+    const at::EventFlag flag) const override {
+    EVENTS_NIEX;
+  }
+
+  void block(
+    void* event,
+    const at::Stream& stream) const override {
+    EVENTS_NIEX;
+  }
+
+  bool queryEvent(void* event) const override {
+    EVENTS_NIEX;
+  }
+
+  void destroyEvent(
+    void* event,
+    const at::DeviceIndex device_index) const noexcept override {
+  }
+
+  #undef EVENTS_NIEX
+
+  //#pragma endregion events
+
+ private:
+  thread_local static at::DeviceIndex current_device_;
+  thread_local static std::map<at::DeviceIndex, at::StreamId> current_streams_;
+};
+
+thread_local at::DeviceIndex ORTGuardImpl::current_device_ = 0;
+thread_local std::map<at::DeviceIndex, at::StreamId> ORTGuardImpl::current_streams_ = {};
+
+C10_REGISTER_GUARD_IMPL(ORT, ORTGuardImpl);
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/ort_hooks.cpp

@@ -0,0 +1,21 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include <ATen/detail/ORTHooksInterface.h>
+
+namespace torch_ort {
+namespace eager {
+
+struct ORTHooks : public at::ORTHooksInterface {
+  ORTHooks(at::ORTHooksArgs) {}
+  std::string showConfig() const override {
+    return "  - ORT is enabled\n";
+  }
+};
+
+using at::ORTHooksRegistry;
+using at::RegistererORTHooksRegistry;
+REGISTER_ORT_HOOKS(ORTHooks);
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/ort_log.cpp

@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include <iostream>
+#include <thread>
+
+#include <c10/util/StringUtil.h>
+
+#include "ort_log.h"
+
+namespace torch_ort {
+namespace eager {
+
+ORTLog::ORTLog(const char* file, int line, ORTLogLevel log_level) {
+  file_ = c10::detail::StripBasename(std::string(file));
+  line_ = line;
+  log_level_ = log_level;
+}
+
+ORTLog::~ORTLog() {
+  static const char* const LOG_PREFIX = "FEWIDVT";
+  static std::mutex mutex;
+
+  mutex.lock();
+
+  auto& out = std::cerr;
+
+  out << "[";
+
+  if (log_level_ < ORTLogLevel::MIN || log_level_ > ORTLogLevel::MAX) {
+    out << "(INVALID_LOG_LEVEL: " << (int)log_level_ << ")";
+  } else {
+    out << LOG_PREFIX[(int)log_level_];
+  }
+
+  out << " ORT " << file_ << ":" << line_ << "] ";
+  out << buffer_.str() << "\n" << std::flush;
+
+  mutex.unlock();
+}
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/ort_log.h

@@ -0,0 +1,125 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <sstream>
+
+#include <c10/core/Scalar.h>
+#include <c10/util/Optional.h>
+#include <ATen/core/Tensor.h>
+
+namespace torch_ort {
+namespace eager {
+
+enum class ORTLogLevel : int {
+  FATAL = 0,
+  ERROR = 1,
+  WARNING = 2,
+  INFO = 3,
+  DEBUG = 4,
+  VERBOSE = 5,
+  TRACE = 6,
+
+  MIN = FATAL,
+  MAX = TRACE 
+};
+
+class ORTLog {
+ public:
+  ORTLog(const char* file, int line, ORTLogLevel log_level);
+  ~ORTLog();
+
+  template<typename T>
+  ORTLog& operator<<(const T value) {
+    buffer_ << value;
+    return *this;
+  }
+
+  ORTLog& operator<<(const at::Device device) {
+    *this << "Device:" << c10::DeviceTypeName(device.type(), true);
+    *this << ":" << (int)device.index();
+    return *this;
+  }
+
+  template<typename T>
+  ORTLog& operator<<(const c10::optional<T> optional) {
+    if (optional.has_value()) {
+      *this << *optional;
+    } else {
+      *this << "None";
+    }
+    return *this;
+  }
+
+  ORTLog& operator<<(const c10::Scalar scalar) {
+    *this << "Scalar:" << scalar.type();
+    switch (scalar.type()) {
+      case c10::ScalarType::Double:
+        *this << "=" << scalar.to<double>();
+        break;
+      case c10::ScalarType::Long:
+        *this << "=" << scalar.to<int64_t>();
+        break;
+      case c10::ScalarType::Bool:
+        *this << "=" << scalar.to<bool>();
+        break;
+      default:
+        break;
+    }
+    return *this;
+  }
+
+  ORTLog& operator<<(const at::Tensor /*tensor*/) {
+    *this << "Tensor";
+    return *this;
+  }
+
+  template<typename F, typename...Ts>
+  static inline void foreach(F f, const Ts&... args) {
+    (void)std::initializer_list<int> {
+      ((void)f(args), 0)...
+    };
+  }
+
+  template<typename...Ts>
+  ORTLog& func(const char* function_name, const Ts&... args) {
+    *this << function_name << "(";
+    
+    unsigned i = 0;
+    foreach([&](const auto& arg) {
+      *this << arg;
+      if (i++ < sizeof...(Ts) - 1)
+        *this << ", ";
+    }, args...);
+
+    *this << ")";
+
+    return *this;
+  }
+
+ private:
+  std::string file_;
+  int line_;
+  ORTLogLevel log_level_;
+  std::stringstream buffer_;
+};
+
+#define ORT_LOG(LEVEL) ORTLog(__FILE__, __LINE__, LEVEL)
+
+#define ORT_LOG_FATAL ORT_LOG(ORTLogLevel::FATAL)
+#define ORT_LOG_ERROR ORT_LOG(ORTLogLevel::ERROR)
+#define ORT_LOG_WARNING ORT_LOG(ORTLogLevel::WARNING)
+#define ORT_LOG_INFO ORT_LOG(ORTLogLevel::INFO)
+#define ORT_LOG_DEBUG ORT_LOG(ORTLogLevel::DEBUG)
+#define ORT_LOG_VERBOSE ORT_LOG(ORTLogLevel::VERBOSE)
+#define ORT_LOG_TRACE ORT_LOG(ORTLogLevel::TRACE)
+
+#ifdef __clang__
+#  define ORT_LOG_FN(...) ORT_LOG_VERBOSE.func(__PRETTY_FUNCTION__,##__VA_ARGS__)
+#else
+#  define ORT_LOG_FN(...) ORT_LOG_VERBOSE << __PRETTY_FUNCTION__
+#endif
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/ort_ops.cpp

@@ -0,0 +1,45 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "ort_ops.h"
+#include "ort_util.h"
+#include "ort_log.h"
+
+namespace torch_ort {
+namespace eager {
+
+OrtValue reshape_copy(
+  onnxruntime::ORTInvoker& invoker,
+  const OrtValue& input,
+  std::vector<int64_t> shape) {
+  
+  // TODO: actual reshape on buffer
+  const onnxruntime::Tensor& input_tensor = input.Get<onnxruntime::Tensor>();
+  auto new_shape = at::infer_size(shape, input_tensor.Shape().Size());
+  OrtValue shape_tensor;
+  //todo: avoid the copy on this small shape vector;
+  auto element_type = onnxruntime::DataTypeImpl::GetType<int64_t>();
+  CreateMLValue(invoker.GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault),
+                element_type, {(int64_t)new_shape.size(),}, &shape_tensor);
+  auto* ort_shape_tensor = shape_tensor.GetMutable<onnxruntime::Tensor>();
+  CopyVectorToTensor<int64_t>(invoker, new_shape, *ort_shape_tensor);
+  std::vector<OrtValue> result(1);
+  auto status = invoker.Invoke("Reshape", {input, shape_tensor}, result, nullptr);
+  if (!status.IsOK())
+    throw std::runtime_error("ORT return failure status: " + status.ErrorMessage());
+  return result[0];
+}
+
+void copy(onnxruntime::ORTInvoker& invoker, 
+          const OrtValue& src, OrtValue& dst){
+  auto& ort_ep = invoker.GetCurrentExecutionProvider();
+  
+  const auto& src_tensor = src.Get<onnxruntime::Tensor>();
+  auto* dst_tensor = dst.GetMutable<onnxruntime::Tensor>();
+  if (!dst_tensor)
+    throw std::runtime_error("ORT copy: dst is not a tensor");
+  ort_ep.GetDataTransfer()->CopyTensor(src_tensor, *dst_tensor);
+}
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/ort_ops.h

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <core/framework/ml_value.h>
+#include <core/eager/ort_kernel_invoker.h>
+
+namespace torch_ort {
+namespace eager {
+
+OrtValue reshape_copy(
+  onnxruntime::ORTInvoker& invoker,
+  const OrtValue& input,
+  std::vector<int64_t> shape);
+
+OrtValue add(onnxruntime::ORTInvoker& invoker,
+             const OrtValue& A,
+             const OrtValue& B);
+
+void copy(onnxruntime::ORTInvoker& invoker, 
+          const OrtValue& src, OrtValue& dst);
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/ort_tensor.cpp

@@ -0,0 +1,109 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "ort_tensor.h"
+#include "ort_util.h"
+
+namespace torch_ort {
+namespace eager {
+
+c10::intrusive_ptr<c10::TensorImpl> ORTTensorImpl::shallow_copy_and_detach(
+  const c10::VariableVersion& version_counter,
+  bool allow_tensor_metadata_change) const {
+  auto impl = c10::make_intrusive<ORTTensorImpl>(
+    tensor_,
+    at::TensorOptions()
+      .dtype(this->dtype())
+      .device(this->device()));
+
+  copy_tensor_metadata(
+    this,
+    impl.get(),
+    version_counter,
+    allow_tensor_metadata_change);
+
+  return impl;
+}
+
+c10::intrusive_ptr<c10::TensorImpl> ORTTensorImpl::shallow_copy_and_detach(
+  c10::VariableVersion&& version_counter,
+  bool allow_tensor_metadata_change) const {
+  auto impl = c10::make_intrusive<ORTTensorImpl>(
+    tensor_,
+    at::TensorOptions()
+      .dtype(this->dtype())
+      .device(this->device()));
+
+  copy_tensor_metadata(
+    this,
+    impl.get(),
+    std::move(version_counter),
+    allow_tensor_metadata_change);
+
+  return impl;
+}
+
+void ORTTensorImpl::shallow_copy_from(
+  const c10::intrusive_ptr<TensorImpl>& impl) {
+  auto* src_impl = dynamic_cast<ORTTensorImpl*>(impl.get());
+  copy_tensor_metadata(
+    src_impl,
+    this,
+    version_counter(),
+    allow_tensor_metadata_change());
+}
+
+at::IntArrayRef ORTTensorImpl::sizes() const {
+  const_cast<ORTTensorImpl*>(this)->cacheSizeMetadata();
+  return c10::TensorImpl::sizes();
+}
+
+int64_t ORTTensorImpl::dim() const {
+  const_cast<ORTTensorImpl*>(this)->cacheSizeMetadata();
+  return c10::TensorImpl::dim();
+}
+
+int64_t ORTTensorImpl::numel() const {
+  const_cast<ORTTensorImpl*>(this)->cacheSizeMetadata();
+  return c10::TensorImpl::numel();
+}
+
+bool ORTTensorImpl::is_contiguous(at::MemoryFormat memory_format) const {
+  return true;
+}
+
+int64_t ORTTensorImpl::size(int64_t d) const {
+  const_cast<ORTTensorImpl*>(this)->cacheSizeMetadata();
+  return c10::TensorImpl::size(d);
+}
+
+void ORTTensorImpl::cacheSizeMetadata() {
+  // TODO: wrap with change generation guard
+  auto& tensor = tensor_.Get<onnxruntime::Tensor>();
+  auto shape = tensor.Shape();
+  auto strides = GetStrides(shape.GetDims());
+
+  numel_ = shape.Size();
+
+  sizes_and_strides_.set_sizes(shape.GetDims());
+
+  for (std::size_t i = 0; i < strides.size(); i++) {
+    sizes_and_strides_.stride_at_unchecked(i) = strides[i];
+  }
+}
+
+const at::Storage& ORTTensorImpl::storage() const {
+  throw std::runtime_error("ORT Tensors do not have storage");
+}
+
+bool ORTTensorImpl::has_storage() const {
+  return false;
+}
+
+at::IntArrayRef ORTTensorImpl::strides() const {
+  const_cast<ORTTensorImpl*>(this)->cacheSizeMetadata();
+  return sizes_and_strides_.strides_arrayref(); 
+}
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/ort_tensor.h

@@ -0,0 +1,65 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <c10/core/TensorImpl.h>
+#include <core/framework/ml_value.h>
+
+namespace torch_ort {
+namespace eager {
+
+class ORTTensorImpl final : public c10::TensorImpl {
+ public:
+  explicit ORTTensorImpl(OrtValue tensor, const at::TensorOptions& options)
+    : c10::TensorImpl(
+        c10::DispatchKeySet {
+          c10::DispatchKey::ORT,
+          c10::DispatchKey::AutogradORT
+        },
+        options.dtype(),
+        options.device()) {
+    set_tensor(tensor);
+  }
+
+  OrtValue& tensor() {
+    return tensor_;
+  }
+
+  void set_tensor(OrtValue tensor) {
+    tensor_ = std::move(tensor);
+  }
+
+  c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach(
+    const c10::VariableVersion& version_counter,
+    bool allow_tensor_metadata_change) const override;
+
+  c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach(
+    c10::VariableVersion&& version_counter,
+    bool allow_tensor_metadata_change) const override;
+
+  void shallow_copy_from(const c10::intrusive_ptr<TensorImpl>& impl) override;
+
+  at::IntArrayRef sizes() const override;
+
+  int64_t dim() const override;
+
+  int64_t numel() const override;
+
+  bool is_contiguous(at::MemoryFormat memory_format) const override;
+
+  int64_t size(int64_t d) const override;
+
+  const at::Storage& storage() const override;
+
+  bool has_storage() const override;
+
+  at::IntArrayRef strides() const override;  
+
+ private:
+  void cacheSizeMetadata();
+  OrtValue tensor_;
+};
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/ort_util.cpp

@@ -0,0 +1,51 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include <core/providers/cpu/cpu_execution_provider.h>
+
+#include "ort_util.h"
+#include "ort_backends.h"
+
+namespace torch_ort {
+namespace eager {
+
+using namespace onnxruntime;
+
+
+void CreateMLValue(onnxruntime::AllocatorPtr alloc, 
+                   onnxruntime::MLDataType element_type, 
+                   const std::vector<int64_t>& dims, 
+                   OrtValue* p_mlvalue) {
+  onnxruntime::TensorShape shape(dims);
+  std::unique_ptr<onnxruntime::Tensor> p_tensor = std::make_unique<onnxruntime::Tensor>(element_type,
+                                                                      shape,
+                                                                      alloc);
+  p_mlvalue->Init(p_tensor.release(),
+                  onnxruntime::DataTypeImpl::GetType<onnxruntime::Tensor>(),
+                  onnxruntime::DataTypeImpl::GetType<onnxruntime::Tensor>()->GetDeleteFunc());
+}
+
+void CreateMLValue(void* data_ptr, onnxruntime::MLDataType element_type, const std::vector<int64_t>& dims, OrtValue* p_mlvalue) {
+  onnxruntime::TensorShape shape(dims);
+  OrtMemoryInfo *cpu_info;
+  Ort::ThrowOnError(Ort::GetApi().CreateCpuMemoryInfo(OrtArenaAllocator, OrtMemTypeDefault, &cpu_info));
+  std::unique_ptr<onnxruntime::Tensor> p_tensor = std::make_unique<onnxruntime::Tensor>(element_type,
+                                                                      shape,
+                                                                      data_ptr,
+                                                                      *cpu_info);
+  
+  p_mlvalue->Init(p_tensor.release(),
+                  onnxruntime::DataTypeImpl::GetType<onnxruntime::Tensor>(),
+                  onnxruntime::DataTypeImpl::GetType<onnxruntime::Tensor>()->GetDeleteFunc());
+}
+
+std::vector<int64_t> GetStrides(const std::vector<int64_t>& shape) {
+  std::vector<int64_t> strides(shape.size(), 1);
+  for (auto i = shape.size(); i > 1; --i) {
+    strides[i - 2] = strides[i - 1] * shape[i - 1];
+  }
+  return strides;
+}
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/ort_util.h

@@ -0,0 +1,58 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <core/session/onnxruntime_cxx_api.h>
+
+#include "ort_backends.h"
+
+namespace torch_ort {
+namespace eager {
+
+void CreateMLValue(onnxruntime::AllocatorPtr alloc, 
+                   onnxruntime::MLDataType element_type, 
+                   const std::vector<int64_t>& dims, 
+                   OrtValue* p_mlvalue);
+
+void CreateMLValue(void* data_ptr, onnxruntime::MLDataType element_type, const std::vector<int64_t>& dims, OrtValue* p_mlvalue);
+
+template <typename T>
+inline void CopyVectorToTensor(onnxruntime::ORTInvoker& invoker,
+                               const std::vector<T>& value,
+                               onnxruntime::Tensor& tensor) {
+  const auto& execution_provider = invoker.GetCurrentExecutionProvider();
+
+  OrtValue* ort_value;
+  int64_t shape = value.size();
+  OrtMemoryInfo cpuMemoryInfo;
+
+  Ort::ThrowOnError(Ort::GetApi().CreateTensorWithDataAsOrtValue(
+    &cpuMemoryInfo,
+    const_cast<void*>(reinterpret_cast<const void*>(value.data())),
+    value.size() * sizeof(T),
+    &shape,
+    1,
+    Ort::TypeToTensorType<T>::type,
+    &ort_value));
+
+  execution_provider.GetDataTransfer()->CopyTensor(
+    ort_value->Get<onnxruntime::Tensor>(),
+    tensor);
+}
+
+// vector<bool> is specialized so we need to handle it separately
+template <>
+inline void CopyVectorToTensor<bool>(onnxruntime::ORTInvoker& /*invoker*/,
+                                     const std::vector<bool>& value,
+                                     onnxruntime::Tensor& tensor) {
+  auto output_span = tensor.MutableDataAsSpan<bool>();
+  for (size_t i = 0, end = value.size(); i < end; ++i) {
+    output_span[i] = value[i];
+  }
+}
+
+std::vector<int64_t> GetStrides(const std::vector<int64_t>& shape);
+
+} // namespace eager
+} // namespace torch_ort

orttraining/orttraining/eager/test/__main__.py

@@ -0,0 +1,15 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import glob
+import os
+import sys
+import subprocess
+
+selfdir = os.path.dirname(os.path.realpath(__file__))
+
+for testpath in glob.glob(os.path.join(selfdir, '*')):
+  if not os.path.basename(testpath).startswith('_'):
+    print(f'Running tests for {testpath} ...')
+    subprocess.check_call([sys.executable, testpath])
+    print()

orttraining/orttraining/eager/test/ort_eps_test.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import unittest
+import torch
+import onnxruntime_pybind11_state as torch_ort
+import os
+
+class OrtEPTests(unittest.TestCase):
+  def get_test_execution_provider_path(self):
+      return os.path.join('.', 'libtest_execution_provider.so')
+
+  def test_import_custom_eps(self):
+    torch_ort.set_device(0, 'CPUExecutionProvider', {})
+
+    torch_ort._register_provider_lib('TestExecutionProvider', self.get_test_execution_provider_path())
+    torch_ort.set_device(1, 'TestExecutionProvider', {'device_id':'0', 'some_config':'val'})
+    ort_device = torch_ort.device(1)
+
+if __name__ == '__main__':
+  unittest.main()

orttraining/orttraining/eager/test/ort_init.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+#
+# This test must run as a separate process since we assert ORT things
+# are unavailable before we import torch_ort, then test that they do
+# in fact become available after importing it. The act of importing
+# torch_ort makes it available implicitly to any tests that may run
+# after the import, hence this test is isolated from the others.
+
+import unittest
+import torch
+
+class OrtInitTests(unittest.TestCase):
+  def test_ort_init(self):
+    config_match = 'ORT is enabled'
+
+    def ort_alloc():
+      torch.zeros(5, 5, device='ort')
+
+    self.assertNotIn(config_match, torch._C._show_config())
+    with self.assertRaises(BaseException):
+      ort_alloc()
+
+    import onnxruntime_pybind11_state as torch_ort
+    ort_alloc()
+    self.assertIn(config_match, torch._C._show_config())
+
+if __name__ == '__main__':
+  unittest.main()

orttraining/orttraining/eager/test/ort_ops.py

@@ -0,0 +1,65 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import unittest
+import torch
+import onnxruntime_pybind11_state as torch_ort
+import numpy as np
+
+class OrtOpTests(unittest.TestCase):
+  def get_device(self):
+    return torch_ort.device()
+
+  def test_add(self):
+    device = self.get_device()
+    cpu_ones = torch.Tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
+    ort_ones = cpu_ones.to(device)
+    cpu_twos = cpu_ones + cpu_ones
+    ort_twos = ort_ones + ort_ones
+    assert torch.allclose(cpu_twos, ort_twos.cpu())
+
+  def test_add_alpha(self):
+    device = self.get_device()
+    cpu_ones = torch.Tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
+    ort_ones = cpu_ones.to(device)
+    assert torch.allclose(
+      torch.add(cpu_ones, cpu_ones, alpha=2.5),
+      torch.add(ort_ones, ort_ones, alpha=2.5).cpu())
+
+  def test_add_(self):
+    device = self.get_device()
+    cpu_ones = torch.Tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
+    ort_ones = cpu_ones.to(device)
+    cpu_twos = cpu_ones
+    cpu_twos += cpu_ones
+    ort_twos = ort_ones
+    ort_twos += ort_ones
+    assert torch.allclose(cpu_twos, ort_twos.cpu())
+
+  def test_sin_(self):
+    device = self.get_device()
+    cpu_sin_pi_ = torch.Tensor([np.pi])
+    torch.sin_(cpu_sin_pi_)
+    ort_sin_pi_ = torch.Tensor([np.pi]).to(device)
+    torch.sin_(ort_sin_pi_)
+    cpu_sin_pi = torch.sin(torch.Tensor([np.pi]))
+    ort_sin_pi = torch.sin(torch.Tensor([np.pi]).to(device))
+    assert torch.allclose(cpu_sin_pi, ort_sin_pi.cpu())
+    assert torch.allclose(cpu_sin_pi_, ort_sin_pi_.cpu())
+    assert torch.allclose(ort_sin_pi.cpu(), ort_sin_pi_.cpu())
+
+  def test_sin(self):
+    device = self.get_device()
+    cpu_sin_pi = torch.sin(torch.Tensor([np.pi]))
+    ort_sin_pi = torch.sin(torch.Tensor([np.pi]).to(device))
+    assert torch.allclose(cpu_sin_pi, ort_sin_pi.cpu())
+  
+  def test_zero_like(self):
+    device = self.get_device()
+    ones = torch.ones((10, 10), dtype=torch.float32)
+    cpu_zeros = torch.zeros_like(ones)
+    ort_zeros = torch.zeros_like(ones.to(device))
+    assert torch.allclose(cpu_zeros, ort_zeros.cpu())
+
+if __name__ == '__main__':
+  unittest.main()

orttraining/orttraining/eager/test/ort_tensor.py

@@ -0,0 +1,24 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import unittest
+import torch
+import onnxruntime_pybind11_state as torch_ort
+
+class OrtTensorTests(unittest.TestCase):
+  def test_is_ort_via_alloc(self):
+    cpu_ones = torch.zeros(10, 10)
+    assert not cpu_ones.is_ort
+    ort_ones = torch.zeros(10, 10, device='ort')
+    assert ort_ones.is_ort
+    assert torch.allclose(cpu_ones, ort_ones.cpu())
+
+  def test_is_ort_via_to(self):
+    cpu_ones = torch.ones(10, 10)
+    assert not cpu_ones.is_ort
+    ort_ones = cpu_ones.to('ort')
+    assert ort_ones.is_ort
+    assert torch.allclose(cpu_ones, ort_ones.cpu())
+
+if __name__ == '__main__':
+  unittest.main()

orttraining/orttraining/eager/test_models/mnist_fc.py

@@ -0,0 +1,45 @@
+from __future__ import print_function
+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import os
+import onnxruntime_pybind11_state as torch_ort
+
+class NeuralNet(nn.Module):
+    def __init__(self, input_size, hidden_size, num_classes):
+        super(NeuralNet, self).__init__()
+        self.fc1 = nn.Linear(input_size, hidden_size)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(hidden_size, num_classes)
+
+    def forward(self, x):
+        out = self.fc1(x)
+        out = self.relu(out)
+        out = self.fc2(out)
+        return out
+
+input_size = 784
+hidden_size = 500
+num_classes = 10
+batch_size = 128
+
+batch = torch.rand((batch_size, input_size))
+device = torch_ort.device()
+
+with torch.no_grad():
+
+    model = NeuralNet(input_size, hidden_size, num_classes)
+    pred = model(batch)
+    print("inference result is: ")
+    print(pred)
+
+    model.to(device)
+
+    ort_batch = batch.to(device)
+    ort_pred = model(ort_batch)
+    print("ORT inference result is:")
+    print(ort_pred.cpu())
+    print("Compare result:")
+    print(torch.allclose(pred, ort_pred.cpu(), atol=1e-6))

orttraining/orttraining/eager/test_models/mnist_fc_training.py

@@ -0,0 +1,109 @@
+## This code is from https://github.com/pytorch/examples/blob/master/mnist/main.py
+## with modification to do training using onnxruntime as backend on cuda device.
+## A private PyTorch build from https://aiinfra.visualstudio.com/Lotus/_git/pytorch (ORTTraining branch) is needed to run the demo.
+
+## Model testing is not complete.
+
+from __future__ import print_function
+import argparse
+import torch
+import onnxruntime_pybind11_state as torch_ort
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+import numpy as np
+import os
+
+dataset_root_dir = os.path.join(
+    os.path.dirname(os.path.realpath(__file__)),
+    'data')
+
+class NeuralNet(nn.Module):
+    def __init__(self, input_size, hidden_size, num_classes):
+        super(NeuralNet, self).__init__()
+        self.fc1 = nn.Linear(input_size, hidden_size)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(hidden_size, num_classes)
+
+    def forward(self, x):
+        out = self.fc1(x)
+        out = self.relu(out)
+        out = self.fc2(out)
+        return out
+
+def my_loss(x, target):
+    return F.nll_loss(F.log_softmax(x, dim=1), target)
+
+def train_with_eager(args, model, optimizer, device, train_loader, epoch):
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data_cpu = data.reshape(data.shape[0], -1)
+        data = data_cpu.to(device)
+
+        x = model(data)
+        loss = my_loss(x.cpu(), target)
+
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+
+
+        # Since the output corresponds to [loss_desc, probability_desc], the first value is taken as loss.
+        if batch_idx % args.log_interval == 0:
+            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
+                epoch, batch_idx * len(data_cpu), len(train_loader.dataset),
+                100. * batch_idx / len(train_loader), loss))
+
+def main():
+#Training settings
+    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
+    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
+                        help='input batch size for training (default: 64)')
+    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
+                        help='input batch size for testing (default: 1000)')
+    parser.add_argument('--epochs', type=int, default=10, metavar='N',
+                        help='number of epochs to train (default: 10)')
+    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
+                        help='learning rate (default: 0.01)')
+    parser.add_argument('--no-cuda', action='store_true', default=False,
+                        help='disables CUDA training')
+    parser.add_argument('--seed', type=int, default=1, metavar='S',
+                        help='random seed (default: 1)')
+    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
+                        help='how many batches to wait before logging training status')
+
+
+    args = parser.parse_args()
+    use_cuda = not args.no_cuda and torch.cuda.is_available()
+
+    torch.manual_seed(args.seed)
+
+    kwargs = {'num_workers': 0, 'pin_memory': True}
+    train_loader = torch.utils.data.DataLoader(
+        datasets.MNIST(dataset_root_dir, train=True, download=True,
+                       transform=transforms.Compose([
+                           transforms.ToTensor(),
+                           transforms.Normalize((0.1307,), (0.3081,))
+                       ])),
+        batch_size=args.batch_size, shuffle=True, **kwargs)
+    test_loader = torch.utils.data.DataLoader(
+        datasets.MNIST(dataset_root_dir, train=False, transform=transforms.Compose([
+            transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])),
+        batch_size=args.test_batch_size, shuffle=True, **kwargs)
+
+    device = torch.device('ort')
+    input_size = 784
+    hidden_size = 500
+    num_classes = 10
+    model = NeuralNet(input_size, hidden_size, num_classes)
+    model.to(device)
+    optimizer = optim.SGD(model.parameters(), lr=0.01)
+
+    print('\nStart Training.')
+
+    for epoch in range(1, args.epochs + 1):
+        train_with_eager(args, model, optimizer, device, train_loader, epoch)
+
+
+if __name__ == '__main__':
+    main()

orttraining/orttraining/eager/test_models/scratchpad.py

@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import torch
+import onnxruntime_pybind11_state as torch_ort
+
+device = torch_ort.device()
+
+ones = torch.ones(2, 3).to(device)
+print(ones.cpu())
+
+twos = ones + ones
+print(twos.cpu())
+
+threes = twos + ones
+print(threes.cpu())
+
+fours = twos * twos
+print(fours.cpu())
+
+fenced_ten = torch.tensor(
+  [[-1, -1, -1],
+   [-1, 10, -1],
+   [-1, -1, -1]],
+  device = device, dtype=torch.float)
+
+print(fenced_ten.numel())
+print(fenced_ten.size())
+print(fenced_ten.cpu())
+print(fenced_ten.relu().cpu())
+
+a = torch.ones(3, 3).to(device)
+b = torch.ones(3, 3)
+c = a + b
+d = torch.sin (c)
+e = torch.tan (c)
+torch.sin_(c)
+print ("sin-in-place:")
+print(c.cpu())
+print ("sin explicit:")
+print (d.cpu ())
+
+a = torch.tensor([[10, 10]], dtype=torch.float).to(device)
+b = torch.tensor([[3.3, 3.3]]).to(device)
+c = torch.fmod(a, b)
+
+print(c.cpu())
+
+a = torch.tensor([[5, 3, -5]], dtype=torch.float).to(device)
+b = torch.hardshrink(a, 3) #should be [5, 0, -5]
+c = torch.nn.functional.softshrink(a, 3) #should be [2, 0, -2]
+print(b.cpu())
+print(c.cpu())

tools/ci_build/build.py

@@ -578,7 +578,7 @@ def get_config_build_dir(build_dir, config):
 
 
 def run_subprocess(args, cwd=None, capture_stdout=False, dll_path=None,
-                   shell=False, env={}):
+                   shell=False, env={}, python_path=None):
     if isinstance(args, str):
         raise ValueError("args should be a sequence of strings, not a string")
 
@@ -591,6 +591,14 @@ def run_subprocess(args, cwd=None, capture_stdout=False, dll_path=None,
                 my_env["LD_LIBRARY_PATH"] += os.pathsep + dll_path
             else:
                 my_env["LD_LIBRARY_PATH"] = dll_path
+    if python_path:
+        if is_windows():
+            my_env["PYTHONPATH"] = python_path + os.pathsep + my_env["PYTHONPATH"]
+        else:
+            if "PYTHONPATH" in my_env:
+                my_env["PYTHONPATH"] += os.pathsep + python_path
+            else:
+                my_env["PYTHONPATH"] = python_path
 
     my_env.update(env)
 
@@ -1010,6 +1018,10 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
     else:
         add_cmake_define_without_override(cmake_extra_defines, "onnxruntime_PYBIND_EXPORT_OPSCHEMA", "OFF")
 
+    if args.build_eager_mode:
+        import torch
+        cmake_args += ["-Donnxruntime_PREBUILT_PYTORCH_PATH=%s" % os.path.dirname(torch.__file__)]
+
     cmake_args += ["-D{}".format(define) for define in cmake_extra_defines]
 
     cmake_args += cmake_extra_args
@@ -1541,6 +1553,11 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs):
                 # run basic frontend tests
                 run_training_python_frontend_tests(cwd=cwd)
 
+            if args.build_eager_mode:
+                # run eager mode test
+                args_list = [sys.executable, os.path.join(cwd, 'eager_test')]
+                run_subprocess(args_list, cwd=cwd, dll_path=dll_path, python_path=cwd)
+
             try:
                 import onnx  # noqa
                 onnx_test = True
@@ -2127,6 +2144,23 @@ def main():
         if args.use_rocm and args.rocm_version is None:
             args.rocm_version = ""
 
+        if args.build_eager_mode:
+            # generate the ort aten backend code
+            def gen_ort_aten_ops(eager_root_dir):
+                gen_cpp_name = os.path.join(eager_root_dir, "ort_aten.g.cpp")
+                if os.path.exists(gen_cpp_name):
+                    os.remove(gen_cpp_name)
+                subprocess.check_call([
+                    sys.executable,
+                    os.path.join(eager_root_dir, 'opgen', 'opgen.py'),
+                    "--output_file",
+                    gen_cpp_name,
+                    "--use_preinstalled_torch"
+                ])
+
+            eager_root_dir = os.path.join(source_dir, "orttraining", "orttraining", "eager")
+            gen_ort_aten_ops(eager_root_dir)
+
         generate_build_tree(
             cmake_path, source_dir, build_dir, cuda_home, cudnn_home, rocm_home, mpi_home, nccl_home,
             tensorrt_home, migraphx_home, acl_home, acl_libs, armnn_home, armnn_libs,