Torch-MLIR LTC Backend Lowering Codegen (#621)

* Codegen and build LTC lowering

* Add LazyShapeInference header

.gitignore

@@ -20,3 +20,16 @@ __pycache__
 
 # Bazel
 bazel-*
+
+# Autogenerated files
+/generated_native_functions.yaml
+/generated_backend.hash
+/python/torch_mlir/csrc/backend/LazyLazyIr.h
+/python/torch_mlir/csrc/backend/LazyNativeFunctions.cpp
+/python/torch_mlir/csrc/backend/LazyNativeFunctions.h
+/python/torch_mlir/csrc/backend/LazyShapeInference.cpp
+/python/torch_mlir/csrc/backend/RegisterLazy.cpp
+
+# Libraries
+*.so
+*.a

python/torch_mlir/csrc/CMakeLists.txt

@@ -20,9 +20,14 @@ link_directories("${TORCH_INSTALL_PREFIX}/lib")
 
 
 add_library(torch_mlir_ltc_backend SHARED
-  backend/backend_impl.cc
-  backend/mlir_lowering_context.cc
-  backend/mlir_node.cc
+  backend/aten_eager_fallback.cpp
+  backend/aten_ltc_mlir_type.cpp
+  backend/backend_impl.cpp
+  backend/LazyNativeFunctions.cpp
+  backend/LazyShapeInference.cpp
+  backend/mlir_lowering_context.cpp
+  backend/mlir_node.cpp
+  backend/RegisterLazy.cpp
 )
 
 target_link_libraries(torch_mlir_ltc_backend
@@ -32,12 +37,13 @@ target_link_libraries(torch_mlir_ltc_backend
   torch_python
 )
 
-message(STATUS "TORCH_CXXFLAGS=${TORCH_CXXFLAGS} -Wpedantic")
+message(STATUS "TORCH_CXXFLAGS=${TORCH_CXXFLAGS} -Wno-pedantic")
 set_target_properties(torch_mlir_ltc_backend PROPERTIES
   LIBRARY_OUTPUT_DIRECTORY "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/"
   OUTPUT_NAME _MLIR_LTC
   PREFIX "${PYTHON_MODULE_PREFIX}"
   SUFFIX "${PYTHON_MODULE_EXTENSION}"
   CXX_VISIBILITY_PRESET "hidden"
-  COMPILE_FLAGS "${TORCH_CXXFLAGS} -Wpedantic"
+  COMPILE_FLAGS "${TORCH_CXXFLAGS} -Wno-pedantic"
 )
+

python/torch_mlir/csrc/backend/backend_impl.h

@@ -1,3 +1,18 @@
+//===- backend_impl.h -----------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// Also available under a BSD-style license. See LICENSE.
+//
+//===----------------------------------------------------------------------===//
+// The Torch-MLIR backend class API that handles lowering LTC ATen ops to MLIR
+// using the Torch-MLIR ATen dialect
+//
+// This file is adapted from pytorch/pytorch
+// https://github.com/pytorch/pytorch/blob/lazy_tensor_staging/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/backend_impl.h
+//===----------------------------------------------------------------------===//
+
 #pragma once
 
 #include <torch/csrc/lazy/backend/backend_data.h>

python/torch_mlir/csrc/backend/mlir_lowering_context.h

@@ -1,9 +1,23 @@
+//===- mlir_lowering_context.h --------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// Also available under a BSD-style license. See LICENSE.
+//
+//===----------------------------------------------------------------------===//
+// This file is adapted from pytorch/pytorch
+// https://github.com/pytorch/pytorch/blob/lazy_tensor_staging/torch/csrc/lazy/ts_backend/ts_lowering_context.h
+//===----------------------------------------------------------------------===//
+
+
 #pragma once
 
 #include <vector>
 
 #include <torch/csrc/lazy/backend/lowering_context.h>
 
+
 namespace torch {
 namespace lazy {
 
@@ -38,9 +52,9 @@ class MlirLoweringContext : public torch::lazy::LoweringContext {
     // shape. Only used for the operator-by-operator execution, mostly for
     // debugging purposes.
     virtual void AddParameter(const torch::lazy::Output& output,
-                      size_t index,
-                      const torch::lazy::Shape& shape,
-                      const std::string& name) override;
+                              size_t index,
+                              const torch::lazy::Shape& shape,
+                              const std::string& name) override;
 
     // Build the computation capturing all the operations created with the
     // embedded builder (returned by the builder() API).