Merge branch 'main' into sambhav/upgrade_sep_14

.github/workflows/merge-rollpytorch.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     if: |
       github.repository == 'llvm/torch-mlir' &&
-      github.event.workflow_run.actor.login == 'silvasean' &&
+      github.event.workflow_run.actor.login == 'stellaraccident' &&
       github.event.workflow_run.conclusion == 'success'
 
     steps:

.gitmodules

@@ -1,6 +1,6 @@
 [submodule "externals/llvm-project"]
 	path = externals/llvm-project
 	url = https://github.com/llvm/llvm-project.git
-[submodule "externals/mlir-hlo"]
-	path = externals/mlir-hlo
-	url = https://github.com/tensorflow/mlir-hlo.git
+[submodule "externals/stablehlo"]
+	path = externals/stablehlo
+	url = https://github.com/openxla/stablehlo.git

CMakeLists.txt

@@ -123,14 +123,7 @@ else()
 endif()
 
 if (TORCH_MLIR_ENABLE_STABLEHLO)
-  set(STABLEHLO_BUILD_EMBEDDED ON)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/externals/mlir-hlo
-    ${CMAKE_CURRENT_BINARY_DIR}/mlir-hlo
-    EXCLUDE_FROM_ALL)
-  include_directories(${CMAKE_CURRENT_SOURCE_DIR}/externals/mlir-hlo/include)
-  include_directories(${CMAKE_CURRENT_SOURCE_DIR}/externals/mlir-hlo)
-  include_directories(${CMAKE_CURRENT_BINARY_DIR}/mlir-hlo/include)
-  include_directories(${CMAKE_CURRENT_BINARY_DIR}/mlir-hlo)
+  include_directories(${CMAKE_CURRENT_SOURCE_DIR}/externals/stablehlo)
 endif()
 
 set(TORCH_MLIR_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
@@ -234,3 +227,18 @@ if (NOT LLVM_INSTALL_TOOLCHAIN_ONLY)
             COMPONENT torch-mlir-headers)
   endif()
 endif()
+
+# Important: If loading StableHLO in this fashion, it must come last,
+# after all of our libraries and test targets have been defined.
+# It seems that they both abuse upstream CMake macros that accumulate
+# properties.
+# Getting this wrong results in building large parts of the stablehlo
+# project that we don't actually depend on. Further some of those parts
+# do not even compile on all platforms.
+if (TORCH_MLIR_ENABLE_STABLEHLO)
+  set(STABLEHLO_BUILD_EMBEDDED ON)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/externals/stablehlo
+    ${CMAKE_CURRENT_BINARY_DIR}/stablehlo
+    EXCLUDE_FROM_ALL)
+  include_directories(${CMAKE_CURRENT_SOURCE_DIR}/externals/stablehlo)
+endif()

build_tools/autogen_ltc_backend.yaml

@@ -1,14 +1,11 @@
 blacklist:
-# List of unsupported ops in LTC autogen because of some error
-- _index_put_impl_  # Error: TODO not sure if there are other valid types to handle here
-- _index_put_impl # Error: TODO not sure if there are other valid types to handle here
-- empty_like  # Error: TODO add support for type BaseType(name=<BaseTy.MemoryFormat: 12>)
-- index.Tensor  # Error: TODO not sure if there are other valid types to handle here
-- index_put  # Error: TODO not sure if there are other valid types to handle here
-- index_put_  # Error: TODO not sure if there are other valid types to handle here
+# Disabled in favour of `aten::index_put` which supports optional indices via `hacked_twin` JIT hack.
+# It also doesn't have confusing `unsafe` argument.
+- _index_put_impl
 
 # Ops with list of tensors output
 - split.Tensor
+- split_with_sizes
 - unbind.int
 - chunk
 
@@ -61,6 +58,8 @@ supported:
 - unbind_copy.int
 - split_copy.Tensor
 - split_with_sizes_copy
+- index.Tensor
+- index_put
 
 # ops required for functionalization
 - lift

build_tools/python_deploy/build_linux_packages.sh

@@ -300,6 +300,11 @@ function test_in_tree() {
 
       echo ":::: Run Linalg e2e integration tests"
       python -m e2e_testing.main --config=linalg -v
+
+      # Dynamo is changing a lot in nightly versions, and thus the implementation
+      # tends to become incompatible to the stable version.
+      echo ":::: Run TorchDynamo e2e integration tests"
+      python -m e2e_testing.main --config=torchdynamo -v
       ;;
     stable)
       echo ":::: Test with stable torch"
@@ -317,12 +322,6 @@ function test_in_tree() {
   echo ":::: Run make_fx + TOSA e2e integration tests"
   python -m e2e_testing.main --config=make_fx_tosa -v
 
-  echo ":::: Run TorchDynamo e2e integration tests"
-  python -m e2e_testing.main --config=torchdynamo -v
-
-  echo ":::: Run StableHLO e2e integration tests"
-  python -m e2e_testing.main --config=stablehlo -v
-
   echo ":::: Run TOSA e2e integration tests"
   python -m e2e_testing.main --config=tosa -v
 

docs/development.md

@@ -408,13 +408,18 @@ Torch-MLIR by default builds with the latest nightly PyTorch version. This can b
 # Updating the LLVM and MLIR-HLO submodules
 
 Torch-MLIR depends on `llvm-project` (which contains, among other things,
-upstream MLIR) and `mlir-hlo`, both of which are submodules in the `externals/`
+upstream MLIR) and `stablehlo`, both of which are submodules in the `externals/`
 directory. We aim to update these at least weekly to bring in the latest
 features and spread out over time the effort of updating our code for MLIR API
 breakages.
 
 ## Which LLVM commit should I pick?
 
+NOTE: This section is in flux. Specifically, the `mlir-hlo` dep has been
+dropped and the project is running off of a `stablehlo` fork which can be
+patched for certain OS combinations. As of 2023-09-12, stellaraccident@
+is massaging this situation. Please reach out for advice updating.
+
 Since downstream projects may want to build Torch-MLIR (and thus LLVM and
 MLIR-HLO) in various configurations (Release versus Debug builds; on Linux,
 Windows, or macOS; possibly with Clang, LLD, and LLDB enabled), it is crucial to

e2e_testing/main.py

@@ -25,7 +25,6 @@
 )
 
 from torch_mlir_e2e_test.linalg_on_tensors_backends.refbackend import RefBackendLinalgOnTensorsBackend
-from torch_mlir_e2e_test.stablehlo_backends.linalg_on_tensors import LinalgOnTensorsStablehloBackend
 from torch_mlir_e2e_test.tcp_backends.linalg_on_tensors import LinalgOnTensorsTcpBackend
 from torch_mlir_e2e_test.tosa_backends.linalg_on_tensors import LinalgOnTensorsTosaBackend
 
@@ -37,6 +36,7 @@
     TCP_PASS_SET,
     TOSA_PASS_SET,
     LTC_XFAIL_SET,
+    LTC_CRASHING_SET,
     TORCHDYNAMO_XFAIL_SET,
     TORCHDYNAMO_CRASHING_SET
 )
@@ -46,15 +46,14 @@
 register_all_tests()
 
 def _get_argparse():
-    config_choices = ["native_torch", "torchscript", "linalg", "stablehlo", "tcp", "make_fx_tosa", "tosa", "lazy_tensor_core", "torchdynamo"]
+    config_choices = ["native_torch", "torchscript", "linalg", "tcp", "make_fx_tosa", "tosa", "lazy_tensor_core", "torchdynamo"]
     parser = argparse.ArgumentParser(description="Run torchscript e2e tests.")
     parser.add_argument("-c", "--config",
         choices=config_choices,
         default="linalg",
         help=f"""
 Meaning of options:
 "linalg": run through torch-mlir"s default Linalg-on-Tensors backend.
-"stablehlo": run through torch-mlir"s default StableHLO backend.
 "tcp": run through torch-mlir's default TCP backend.
 "tosa": run through torch-mlir"s default TOSA backend.
 "native_torch": run the torch.nn.Module as-is without compiling (useful for verifying model is deterministic; ALL tests should pass in this configuration).
@@ -107,10 +106,6 @@ def main():
         config = TosaBackendTestConfig(LinalgOnTensorsTosaBackend(), use_make_fx=True)
         xfail_set = all_test_unique_names - MAKE_FX_TOSA_PASS_SET
         crashing_set = set()
-    elif args.config == "stablehlo":
-        config = StablehloBackendTestConfig(LinalgOnTensorsStablehloBackend())
-        xfail_set = all_test_unique_names - STABLEHLO_PASS_SET
-        crashing_set = STABLEHLO_CRASHING_SET
     elif args.config == "native_torch":
         config = NativeTorchTestConfig()
         xfail_set = set()
@@ -122,7 +117,7 @@ def main():
     elif args.config == "lazy_tensor_core":
         config = LazyTensorCoreTestConfig()
         xfail_set = LTC_XFAIL_SET
-        crashing_set = set()
+        crashing_set = LTC_CRASHING_SET
     elif args.config == "torchdynamo":
         config = TorchDynamoTestConfig(RefBackendLinalgOnTensorsBackend())
         xfail_set = TORCHDYNAMO_XFAIL_SET