Replace tm_tensor with torch-mlir proper

.gitmodules

@@ -44,3 +44,6 @@
 [submodule "third_party/stablehlo"]
 	path = third_party/stablehlo
 	url = https://github.com/iree-org/stablehlo.git
+[submodule "third_party/torch-mlir"]
+	path = third_party/torch-mlir
+	url = https://github.com/llvm/torch-mlir.git

CMakeLists.txt

@@ -812,10 +812,7 @@ else()
   # Add default external projects.
   iree_llvm_add_external_project(mlir-iree-dialects ${CMAKE_CURRENT_SOURCE_DIR}/llvm-external-projects/iree-dialects)
   iree_llvm_add_external_project(stablehlo ${CMAKE_CURRENT_SOURCE_DIR}/third_party/stablehlo)
-  if(IREE_INPUT_TORCH)
-    iree_llvm_add_external_project(torch-mlir-dialects ${CMAKE_CURRENT_SOURCE_DIR}/third_party/torch-mlir-dialects)
-  endif()
-
+
   # Ensure that LLVM-based dependencies needed for testing are included.
   add_dependencies(iree-test-deps FileCheck)
   if(IREE_LLD_TARGET)

build_tools/bazel/workspace.bzl

@@ -152,13 +152,6 @@ def configure_iree_submodule_deps(iree_repo_alias = "@", iree_path = "./"):
         path = paths.join(iree_path, "third_party/spirv_cross"),
     )
 
-    maybe(
-        native.new_local_repository,
-        name = "torch-mlir-dialects",
-        build_file = iree_repo_alias + "//:build_tools/third_party/torch-mlir-dialects/BUILD.overlay",
-        path = paths.join(iree_path, "third_party/torch-mlir-dialects"),
-    )
-
     maybe(
         native.new_local_repository,
         name = "tracy_client",

build_tools/bazel_to_cmake/bazel_to_cmake_targets.py

@@ -86,10 +86,6 @@ def __init__(self, repo_map: Dict[str, str]):
                 "@nccl//:headers": [
                     "nccl::headers",
                 ],
-                # Torch-MLIR.
-                "@torch-mlir-dialects//:TorchMLIRTMTensorDialect": [
-                    "TorchMLIRTMTensorDialect"
-                ],
                 # Tracy.
                 "@tracy_client//:runtime": ["tracy_client::runtime"],
                 # Vulkan

compiler/plugins/input/Torch/CMakeLists.txt

@@ -0,0 +1,25 @@
+# Copyright 2023 Nod Labs, Inc
+#
+# Licensed 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
+
+set(TORCH_MLIR_ROOT_DIR "${IREE_SOURCE_DIR}/third_party/torch-mlir")
+set(IREE_PACKAGE_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}")
+set(IREE_PACKAGE_ROOT_PREFIX "")
+set(IREE_COMPILER_TABLEGEN_INCLUDE_DIRS "${CMAKE_CURRENT_SOURCE_DIR}")
+
+add_library(torch-iree_compiler_defs INTERFACE)
+target_include_directories(torch-iree_compiler_defs
+  INTERFACE
+  ${CMAKE_CURRENT_SOURCE_DIR}
+  ${CMAKE_CURRENT_BINARY_DIR}
+)
+
+# Configures all iree_cc_* targets to take this implicit dep,
+# which provides common includes and copts for the tree.
+set(IREE_IMPLICIT_DEFS_CC_DEPS torch-iree_compiler_defs)
+
+add_subdirectory(torch-iree)
+add_subdirectory(torch-mlir)
+add_subdirectory(torch-mlir-dialects)

compiler/plugins/input/Torch/torch-iree/CMakeLists.txt

@@ -0,0 +1,32 @@
+# Copyright 2023 Nod Labs, Inc
+#
+# Licensed 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
+
+add_subdirectory(InputConversion)
+
+iree_cc_library(
+  NAME
+    registration
+  SRCS
+    "PluginRegistration.cpp"
+  DEPS
+    MLIRIR
+    MLIRPass
+    iree::compiler::PluginAPI
+    torch-iree::InputConversion
+    torch-mlir::TorchConversionDialectIR
+    torch-mlir::TorchDialectIR
+    torch-mlir::TorchDialectPasses
+    torch-mlir::ConversionPasses
+    torch-mlir-dialects::TMTensorDialectIR
+  PUBLIC
+)
+
+iree_compiler_register_plugin(
+  PLUGIN_ID
+    torch_iree
+  TARGET
+    ::registration
+)

compiler/plugins/input/Torch/torch-iree/InputConversion/CMakeLists.txt

@@ -0,0 +1,51 @@
+# Copyright 2023 Nod Labs, Inc
+#
+# Licensed 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
+
+add_subdirectory(test)
+
+iree_tablegen_library(
+  NAME
+    PassesIncGen
+  TD_FILE
+    "Passes.td"
+  OUTS
+    --gen-pass-decls Passes.h.inc
+)
+
+iree_cc_library(
+  NAME
+    PassHeaders
+  HDRS
+    "PassDetail.h"
+    "Passes.h"
+    "Passes.h.inc"
+  DEPS
+    ::PassesIncGen
+    MLIRPass
+    MLIRTransforms
+  PUBLIC
+)
+
+iree_cc_library(
+  NAME
+    InputConversion
+  HDRS
+    "Passes.h"
+  SRCS
+    "ConvertTMTensorToLinalgExt.cpp"
+    "Passes.cpp"
+  DEPS
+    ::PassHeaders
+    ::PassesIncGen
+    IREELinalgExtDialect
+    MLIRFuncDialect
+    MLIRIR
+    MLIRPass
+    MLIRTensorDialect
+    MLIRTransforms
+    torch-mlir-dialects::TMTensorDialectIR
+  PUBLIC
+)

compiler/src/iree/compiler/InputConversion/TMTensor/ConvertTMTensorToLinalgExt.cpp → compiler/plugins/input/Torch/torch-iree/InputConversion/ConvertTMTensorToLinalgExt.cpp

@@ -9,16 +9,16 @@
 
 #include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtDialect.h"
 #include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
-#include "iree/compiler/InputConversion/TMTensor/PassDetail.h"
-#include "iree/compiler/InputConversion/TMTensor/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "torch-iree/InputConversion/PassDetail.h"
+#include "torch-iree/InputConversion/Passes.h"
 #include "torch-mlir-dialects/Dialect/TMTensor/IR/TMTensorOps.h"
 
 namespace mlir {
 namespace iree_compiler {
-namespace TMTensor {
+namespace TorchInput {
 
 namespace {
 
@@ -205,6 +205,6 @@ createConvertTMTensorToLinalgExtPass() {
   return std::make_unique<ConvertTMTensorToLinalgExtPass>();
 }
 
-} // namespace TMTensor
+} // namespace TorchInput
 } // namespace iree_compiler
 } // namespace mlir

compiler/src/iree/compiler/InputConversion/TMTensor/PassDetail.h → compiler/plugins/input/Torch/torch-iree/InputConversion/PassDetail.h

@@ -4,21 +4,21 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#ifndef IREE_COMPILER_INPUTCONVERSION_TMTENSOR_PASSDETAIL_H_
-#define IREE_COMPILER_INPUTCONVERSION_TMTENSOR_PASSDETAIL_H_
+#ifndef TORCH_IREE_INPUTCONVERSION_PASSDETAIL_H_
+#define TORCH_IREE_INPUTCONVERSION_PASSDETAIL_H_
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 namespace iree_compiler {
-namespace TMTensor {
+namespace TorchInput {
 
 #define GEN_PASS_CLASSES
-#include "iree/compiler/InputConversion/TMTensor/Passes.h.inc"
+#include "torch-iree/InputConversion/Passes.h.inc"
 
-} // namespace TMTensor
+} // namespace TorchInput
 } // namespace iree_compiler
 } // namespace mlir
 
-#endif // IREE_COMPILER_INPUTCONVERSION_TMTENSOR_PASSDETAIL_H_
+#endif // TORCH_IREE_INPUTCONVERSION_PASSDETAIL_H_

compiler/src/iree/compiler/InputConversion/TMTensor/Passes.cpp → compiler/plugins/input/Torch/torch-iree/InputConversion/Passes.cpp

@@ -4,22 +4,22 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include "iree/compiler/InputConversion/TMTensor/Passes.h"
+#include "torch-iree/InputConversion/Passes.h"
 
 namespace mlir {
 namespace iree_compiler {
-namespace TMTensor {
+namespace TorchInput {
 
 namespace {
 #define GEN_PASS_REGISTRATION
-#include "iree/compiler/InputConversion/TMTensor/Passes.h.inc" // IWYU pragma: export
+#include "torch-iree/InputConversion/Passes.h.inc" // IWYU pragma: export
 } // namespace
 
 void registerTMTensorConversionPasses() {
   // Generated.
   registerPasses();
 }
 
-} // namespace TMTensor
+} // namespace TorchInput
 } // namespace iree_compiler
 } // namespace mlir