LLVM / Torch-Mlir Upgrade (#8)

- Remove StableHLO dependency
- Apply patches to build torch-mlir without stablehlo and stablehlo lit
tests
- Update APIs realted to` isa<>`, `cast<>`, and `dyn_cast<>`
- Add a new EliminateUnusedTorchOpsPass to remove unused torch ops in
torch-to-tcp (see example below):


```
module {
  func.func @func_main(%arg0: !torch.vtensor<[?,3],f32>, %arg1: !torch.vtensor<[?,3],f32>) -> !torch.vtensor<[?,3],f32> {
    %int2 = torch.constant.int 2
    %int0 = torch.constant.int 0
    %0 = torch.symbolic_int "s35" {min_val = 0, max_val = 9223372036854775807} : !torch.int
    torch.bind_symbolic_shape %arg0, [%0], affine_map<()[s0] -> (s0, 3)> : !torch.vtensor<[?,3],f32>
    torch.bind_symbolic_shape %arg1, [%0], affine_map<()[s0] -> (s0, 3)> : !torch.vtensor<[?,3],f32>
    %1 = torch.aten.size.int %arg0, %int0 : !torch.vtensor<[?,3],f32>, !torch.int -> !torch.int
    %2 = torch.aten.size.int %arg1, %int0 : !torch.vtensor<[?,3],f32>, !torch.int -> !torch.int
    %3 = torch.aten.sub.Tensor %arg0, %arg1, %int2 : !torch.vtensor<[?,3],f32>, !torch.vtensor<[?,3],f32>, !torch.int -> !torch.vtensor<[?,3],f32>
    torch.bind_symbolic_shape %3, [%0], affine_map<()[s0] -> (s0, 3)> : !torch.vtensor<[?,3],f32>
    %4 = torch.aten.eq.int %1, %2 : !torch.int, !torch.int -> !torch.bool
    %5 = torch.aten.Int.bool %4 : !torch.bool -> !torch.int
    %6 = torch.aten.Bool.int %5 : !torch.int -> !torch.bool
    torch.runtime.assert %6, "Runtime assertion failed for expression Eq(s35, s58) on node 'eq_2'"
    return %3 : !torch.vtensor<[?,3],f32>
  }
}
```

where the `torch.runtime.assert` and related checking ops can be
removed.

Author: zezhang

.bazelignore

@@ -3,7 +3,6 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 # Also available under a BSD-style license. See LICENSE.
 
-# ignore local_repos of llvm-project, torch-mlir, stablehlo
+# ignore local_repos of llvm-project, torch-mlir
 third_party/llvm-project
 third_party/torch-mlir
-third_party/stablehlo

.github/workflows/bazelBuildAndTestStablehlo.yml

@@ -9,10 +9,9 @@ bazel-out
 bazel-mlir-tcp
 bazel-testlogs
 
-# ignore local_repos of llvm, torch-mlir, stablehlo
+# ignore local_repos of llvm, torch-mlir
 third_party/llvm-project
 third_party/torch-mlir
-third_party/stablehlo
 
 # clangd related
 .cache

.gitignore

@@ -150,6 +150,7 @@ cc_library(
     name = "TcpDialectPasses",
     srcs = [
         "lib/Dialect/Transforms/DropSymbolicShapeOpsPass.cpp",
+        "lib/Dialect/Transforms/EliminateUnusedTorchOpsPass.cpp",
         "lib/Dialect/Transforms/FuseTcpOpsPass.cpp",
         "lib/Dialect/Transforms/FusionPatterns.cpp",
         "lib/Dialect/Transforms/IsolateGroupOpsPass.cpp",
@@ -160,6 +161,7 @@ cc_library(
     ],
     hdrs = [
         "include/mlir-tcp/Dialect/Transforms/DropSymbolicShapeOpsPass.h",
+        "include/mlir-tcp/Dialect/Transforms/EliminateUnusedTorchOpsPass.h",
         "include/mlir-tcp/Dialect/Transforms/FuseTcpOpsPass.h",
         "include/mlir-tcp/Dialect/Transforms/FusionPatterns.h",
         "include/mlir-tcp/Dialect/Transforms/IsolateGroupOpsPass.h",
@@ -175,6 +177,7 @@ cc_library(
         "@llvm-project//mlir:TensorDialect",
         "@llvm-project//mlir:TensorTransforms",
         "@llvm-project//mlir:Transforms",
+        "@torch-mlir//:TorchMLIRTorchDialect",
     ],
 )
 
@@ -198,7 +201,6 @@ cc_library(
     hdrs = ["include/mlir-tcp/Conversion/Passes.h"],
     strip_include_prefix = "include",
     deps = [
-        ":StablehloToTcp",
         ":TcpToArith",
         ":TcpToLinalg",
         ":TcpToTensor",
@@ -237,25 +239,6 @@ cc_library(
     ],
 )
 
-cc_library(
-    name = "StablehloToTcp",
-    srcs = [
-        "lib/Conversion/PassDetail.h",
-        "lib/Conversion/StablehloToTcp/StablehloToTcp.cpp",
-    ],
-    hdrs = ["include/mlir-tcp/Conversion/StablehloToTcp/StablehloToTcp.h"],
-    strip_include_prefix = "include",
-    deps = [
-        ":TcpConversionPassesIncGen",
-        ":TcpDialect",
-        "@llvm-project//mlir:Dialect",
-        "@llvm-project//mlir:LinalgDialect",
-        "@llvm-project//mlir:Pass",
-        "@llvm-project//mlir:Transforms",
-        "@stablehlo//:stablehlo_ops",
-    ],
-)
-
 cc_library(
     name = "TcpToLinalg",
     srcs = [
@@ -364,6 +347,5 @@ cc_binary(
         "@llvm-project//mlir:AllPassesAndDialects",
         "@llvm-project//mlir:MlirOptLib",
         "@llvm-project//mlir:QuantOps",
-        "@stablehlo//:register",
     ],
 )

BUILD

@@ -51,17 +51,15 @@ bazel run //tools/clangd:refresh_compile_commands
 ```
 When run successfully, a `compile_commands.json` is generated at the workspace root (and refreshed upon re-runs). If you're using VSCode, just hit CMD+SHIFT+P and select `clangd: Restart language server` to start clangd. Note that this only works for non-docker builds at the moment.
 
-When bumping upstream dependencies (LLVM, Torch-MLIR, StableHLO), you may validate the set of "green commits" by running the corresponding third-party tests:
+When bumping upstream dependencies (LLVM, Torch-MLIR), you may validate the set of "green commits" by running the corresponding third-party tests:
 ```shell
 bazel test @llvm-project//mlir/...
 bazel test @torch-mlir//...
-bazel test @stablehlo//...
 ```
 
 The following CI workflows are automatically triggered anytime upstream dependencies (`deps.bzl`) are updated:
 - [![Bazel Build and Test (llvm-project)](https://github.com/llvm/mlir-tcp/actions/workflows/bazelBuildAndTestLlvm.yml/badge.svg)](https://github.com/llvm/mlir-tcp/actions/workflows/bazelBuildAndTestLlvm.yml)
 - [![Bazel Build and Test (torch-mlir)](https://github.com/llvm/mlir-tcp/actions/workflows/bazelBuildAndTestTorchmlir.yml/badge.svg)](https://github.com/llvm/mlir-tcp/actions/workflows/bazelBuildAndTestTorchmlir.yml)
-- [![Bazel Build and Test (stablehlo)](https://github.com/llvm/mlir-tcp/actions/workflows/bazelBuildAndTestStablehlo.yml/badge.svg)](https://github.com/llvm/mlir-tcp/actions/workflows/bazelBuildAndTestStablehlo.yml)
 
 To use newer `torch-mlir` and/or `torch` python packages in our hermetic python sandbox, just regenerate `requirements_lock.txt` as follows:
 ```shell

README.md

@@ -7,10 +7,8 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 load(
     ":local_repos.bzl",
     "local_llvm_repo_path",
-    "local_stablehlo_repo_path",
     "local_torch_mlir_repo_path",
     "use_local_llvm_repo",
-    "use_local_stablehlo_repo",
     "use_local_torch_mlir_repo",
 )
 
@@ -22,8 +20,8 @@ def third_party_deps():
             path = local_llvm_repo_path(),
         )
     else:
-        LLVM_COMMIT = "72144d119a7291f8b6b8e022a2947fbe31e66afc"
-        LLVM_SHA256 = "2caacb6925a13cb5886a5d7f225fa408b80ca8e1efe0736186954b2abc4ee1c3"
+        LLVM_COMMIT = "b231e5ff504295641b0f580ceefa2e1048011614"
+        LLVM_SHA256 = "88dfa59052730710cb48fa20b00a4344144edd1c3cb524c06d983899835e491a"
         http_archive(
             name = "llvm-raw",
             build_file_content = "# empty",
@@ -39,32 +37,15 @@ def third_party_deps():
             path = local_torch_mlir_repo_path(),
         )
     else:
-        TORCH_MLIR_COMMIT = "9f2ba5abaa85cefd95cc85579fafd0c53c1101e8"
-        TORCH_MLIR_SHA256 = "09444281839eeae4aff42c029d87b1728f307fa26511b896ff448d51aaa98049"
+        TORCH_MLIR_COMMIT = "1ad9702d2a290b693c4f6f17921d0e0a8d14a999"
+        TORCH_MLIR_SHA256 = "8843399168c34ca3ca16d2417703fe4e1440ca7240d9e04844b3deedf256f0ab"
         http_archive(
             name = "torch-mlir-raw",
             build_file_content = "# empty",
+            patches = ["//third_party/patches:torch-mlir-bazel-build.1.patch", "//third_party/patches:torch-mlir-bazel-build.2.patch"],
             sha256 = TORCH_MLIR_SHA256,
             strip_prefix = "torch-mlir-" + TORCH_MLIR_COMMIT,
             urls = ["https://github.com/llvm/torch-mlir/archive/{commit}.tar.gz".format(commit = TORCH_MLIR_COMMIT)],
-            patches = [
-                "//third_party/patches:torch-mlir.1.patch",
-            ],
-        )
-
-    if use_local_stablehlo_repo():
-        native.local_repository(
-            name = "stablehlo",
-            path = local_stablehlo_repo_path(),
-        )
-    else:
-        STABLEHLO_COMMIT = "a54938f0651d3b4b7be9771848eda2463c92a8e7"
-        STABLEHLO_SHA256 = "edab2288f0b19e3efbf08815d17d4efb106984aa6fe02fed0cb2165284e6a5b7"
-        http_archive(
-            name = "stablehlo",
-            sha256 = STABLEHLO_SHA256,
-            strip_prefix = "stablehlo-" + STABLEHLO_COMMIT,
-            urls = ["https://github.com/openxla/stablehlo/archive/{commit}.tar.gz".format(commit = STABLEHLO_COMMIT)],
         )
 
     SKYLIB_VERSION = "1.3.0"

deps.bzl

@@ -46,23 +46,6 @@ def ConvertTorchToTcpCustomOp : Pass<"convert-torch-to-tcp-custom-op", "func::Fu
   ];
 }
 
-//===----------------------------------------------------------------------===//
-// StablehloToTcp
-//===----------------------------------------------------------------------===//
-
-def ConvertStablehloToTcp
-      : Pass<"convert-stablehlo-to-tcp", "func::FuncOp"> {
-  let summary = "Lower StableHLO to TCP";
-  let description = [{
-    Pass that converts StableHLO operations to equivalent operations in TCP.
-  }];
-
-  let constructor = "mlir::tcp::createConvertStablehloToTcpPass()";
-  let dependentDialects = [
-    "mlir::tcp::TcpDialect",
-  ];
-}
-
 //===----------------------------------------------------------------------===//
 // TcpToLinalg
 //===----------------------------------------------------------------------===//

include/mlir-tcp/Conversion/Passes.td

@@ -24,8 +24,8 @@ include "mlir-tcp/Dialect/IR/TcpBase.td"
 // Where low and high ends are 0,255 when unsigned, -128,127 when signed, for
 // the 8-bit case.
 class Tcp_QuantizedType<string n, list<int> params, bit signed>
-  : Type<And<[CPred<"$_self.isa<mlir::quant::QuantizedType>()">,
-              CPred<"$_self.cast<mlir::quant::QuantizedType>()" #
+  : Type<And<[CPred<"::llvm::isa<mlir::quant::QuantizedType>($_self)">,
+              CPred<"::llvm::cast<mlir::quant::QuantizedType>($_self)" #
                     ".getStorageTypeIntegralWidth() == " # !head(params)>]>,
     "Q" # !if (signed, "int", "uint") # !head(params) # " type"> {
   string name = n;

include/mlir-tcp/Dialect/IR/TcpTypes.td

@@ -9,17 +9,13 @@
 
 #pragma once
 
-#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/Pass.h"
+#include <memory>
 
-namespace mlir {
+namespace mlir::tcp {
 
-#define GEN_PASS_DECL_CONVERTSTABLEHLOTOTCP
-#include "mlir-tcp/Conversion/Passes.h.inc"
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+createEliminateUnusedTorchOpsPass();
 
-namespace tcp {
-
-std::unique_ptr<OperationPass<func::FuncOp>> createConvertStablehloToTcpPass();
-
-} // namespace tcp
-} // namespace mlir
+} // namespace mlir::tcp

include/mlir-tcp/Conversion/StablehloToTcp/StablehloToTcp.h renamed to include/mlir-tcp/Dialect/Transforms/EliminateUnusedTorchOpsPass.h

@@ -45,4 +45,10 @@ def DropSymbolicShapeOps : Pass<"drop-symbolic-shape-ops", "func::FuncOp"> {
   let constructor = "mlir::tcp::createDropSymbolicShapeOpsPass()";
 }
 
+// \brief This pass removes unused torch ops.
+def EliminateUnusedTorchOps : Pass<"eliminate-unused-torch-ops", "ModuleOp"> {
+  let summary = "Removes unused/unnecessary torch ops";
+  let constructor = "mlir::tcp::createEliminateUnusedTorchOpsPass()";
+}
+
 #endif // TCP_PASSES

include/mlir-tcp/Dialect/Transforms/Passes.td

@@ -9,7 +9,6 @@
 
 #include "mlir-tcp/Conversion/Passes.h"
 
-#include "mlir-tcp/Conversion/StablehloToTcp/StablehloToTcp.h"
 #include "mlir-tcp/Conversion/TcpToArith/TcpToArith.h"
 #include "mlir-tcp/Conversion/TcpToLinalg/TcpToLinalg.h"
 #include "mlir-tcp/Conversion/TcpToTensor/TcpToTensor.h"