Fix torch/LAPIS integration (#47)

* Create variables to test if torch-mlir/mpact enabled

* Get rid of test translations in lapis-translate
This just duplicates what mlir-translate already does

* Add lower_and_emit_kokkos CAPI/py function

Given an MLIR source string for a module, it
parses, lowers and emits Kokkos in one shot

Works for matadd when called directly from Python

* Add lapis-emit CLI utility
Does end-to-end lowering and emitting.
Takes input as MLIR source, as file or stdin (default=stdin)

matadd.py and torchscript_resnet18.py work again when called from
Python. Dense inputs and outputs are in device space so they're suitable
for benchmarking from C++ driver.

* Fixes for MPACT version of spmv
- handle casted integers in loop mapping pass (CSR pattern detection)
- support memref.subview in emitter

README
- use https, not ssh URLs for repositories 
- add torch, torchvision SHAs and instruction link to build from source.

* Modify util functions to make params/returns DualView

* WIP: updates to emitter
- multiple results
- making all memref returns of non-extern funcs DualView

* More detailed IR dumps from lapis-emit
* Add KokkosBackend.compile_mpact function

Runs new sparse pipeline and emits C++.
Python wrapper still broken but C++ is good.

* Minimal emitter fixes to get sparseMHA to compile

* Update examples (from Dec review)

- make a copy of resnet18 that uses dynamic batch size
  (resnet18_dynamic.py), since E3SM needs this to work.
- Rename torchscript_resnet18.py to resnet18_static.py since
  it has a compile-time batch dimension (1).
- Add MPACT based spmv example

* Sync results to host before returning to numpy
---------
Signed-off-by: Brian Kelley <[email protected]>

Author: brian-kelley

CMakeLists.txt

@@ -87,6 +87,20 @@ include_directories(${STANDALONE_BINARY_DIR}/include)
 set(LAPIS_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
 set(LAPIS_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}")
 
+if("torch-mlir" IN_LIST LLVM_EXTERNAL_PROJECTS)
+  message(STATUS "LAPIS has torch-mlir support")
+  set(LAPIS_HAS_TORCH_MLIR ON)
+else()
+  set(LAPIS_HAS_TORCH_MLIR OFF)
+endif()
+
+if("mpact" IN_LIST LLVM_EXTERNAL_PROJECTS)
+  message(STATUS "LAPIS has MPACT support")
+  set(LAPIS_HAS_MPACT ON)
+else()
+  set(LAPIS_HAS_MPACT OFF)
+endif()
+
 link_directories(${LLVM_BUILD_LIBRARY_DIR})
 add_definitions(${LLVM_DEFINITIONS})
 

README.md

@@ -29,19 +29,19 @@ are set to the paths of these repositories.
 
 The following commands will clone the correct versions of all the repositories and set these environment variables.
 ```
-git clone git@github.com:MPACT-ORG/mpact-compiler
+git clone https://github.com/MPACT-ORG/mpact-compiler
 cd mpact-compiler
 git checkout 556009cd
 git submodule update --init --recursive
 export MPACT_SRC=`pwd`
 export TORCH_MLIR_SRC="$MPACT_SRC/externals/torch-mlir"
 export LLVM_SRC="$TORCH_MLIR_SRC/externals/llvm-project"
 cd ..
-git clone git@github.com:sandialabs/LAPIS
+git clone https://github.com/sandialabs/LAPIS
 cd LAPIS
 export LAPIS_SRC=`pwd`
 cd ..
-git clone -b master git@github.com:kokkos/kokkos
+git clone -b master https://github.com/kokkos/kokkos
 ```
 
 Building with ninja is not required but useful as it automatically uses all cores for parallel compilation. Pass ``-Gninja`` to
@@ -92,15 +92,35 @@ cd ..
 This recipe builds LAPIS as an external project with LLVM.
 torch-mlir and mpact require this recipe, but torch-mlir and mpact are still optional.
 mpact requires torch-mlir, however.
-**This requires ninja due to an issue in torch-mlir. make will not work.**
+**ninja is required due to an issue in torch-mlir. make will not work.**
 ```
 # If enabling torch-mlir, need to install Python dependencies first.
 # This can be done inside a python virtual env.
 
 cd $TORCH_MLIR_SRC
-pip install -r requirements.txt
-pip install -r torchvision-requirements.txt
+pip install -r build-requirements.txt
+pip install -r test-requirements.txt
+```
+
+Then install the torch and torchvision Python packages from source.
+It needs to be from source because torch-mlir only works with specific nightly versions
+of these packages, but the package repository only stores nightly binaries for a limited amount of time.
 
+A CPU-only build is sufficient for LAPIS.
+* [torch repository](https://github.com/pytorch/pytorch), version 995ec16c)
+* [torchvision repository](https://github.com/pytorch/vision), version c7ea645b)
+[Instructions for building from source can be found here](https://github.com/pytorch/pytorch#from-source)
+
+These Git SHAs correspond to the nightly versions 2.5.0.dev20240909 and 0.20.0.dev20240909 respectively.
+*Note for developers:* from installed nightly versions, the exact Git versions can be found with:
+```
+import torch
+import torchvision
+print(torch.version.git_version)
+print(torchvision.version.git_version)
+```
+
+```
 cd $WORKSPACE
 mkdir build
 cd build

examples/MTTKRP.py

@@ -35,9 +35,10 @@
     torch.ones(1, 3, 224, 224),
     output_type="linalg-on-tensors",
     func_name=resnet18.__class__.__name__,
+    experimental_support_mutation=True
 )
-backend = KokkosBackend.KokkosBackend()
-compiled = backend.compile_sparse(module)
+backend = KokkosBackend.KokkosBackend(dump_mlir=True)
+compiled = backend.compile(module)
 #fx_module = backend.load(compiled)
 
 params = {

examples/fx_resnet18.py

@@ -22,7 +22,7 @@ def main():
 
     mlir_module = torchscript.compile(m, (a, b), output_type='linalg-on-tensors')
 
-    backend = KokkosBackend(dump_mlir=True)
+    backend = KokkosBackend.KokkosBackend(dump_mlir=True)
     k_backend = backend.compile(mlir_module)
 
     c = k_backend.forward(a, b)
@@ -34,3 +34,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+

examples/matadd.py

@@ -1,4 +1,5 @@
-#define PYTACO_CPP_DRIVER
+#define LAPIS_CPP_DRIVER
+
 #include <iostream>
 #include <fstream>
 #include "mlir/Dialect/SparseTensor/IR/Enums.h"

examples/matadd_pytaco.py

@@ -0,0 +1,56 @@
+# 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.
+
+import sys
+from pathlib import Path
+
+import torch
+import torchvision.models as models
+from torch_mlir import torchscript
+from torch_mlir.compiler_utils import TensorPlaceholder
+from torch_mlir_e2e_test.linalg_on_tensors_backends import refbackend
+from lapis import KokkosBackend
+
+sys.path.append(str(Path(__file__).absolute().parent))
+from _example_utils import (
+    top3_possibilities,
+    load_and_preprocess_image,
+    load_labels,
+    DEFAULT_IMAGE_URL,
+)
+
+def predictions(torch_func, kokkos_func, img, labels):
+    golden_prediction = top3_possibilities(torch_func(img), labels)
+    print("PyTorch prediction")
+    print(golden_prediction)
+    prediction = top3_possibilities(torch.from_numpy(kokkos_func(img.numpy())), labels)
+    print("LAPIS prediction")
+    print(prediction)
+
+print("load image from " + DEFAULT_IMAGE_URL, file=sys.stderr)
+img = load_and_preprocess_image(DEFAULT_IMAGE_URL)
+labels = load_labels()
+#print("Dumping preprocessed dog image to dog.bin")
+#img.numpy().tofile('dog.bin')
+
+resnet18 = models.resnet18(pretrained=True).eval()
+#print(help(torchscript.compile))
+
+module = torchscript.compile(
+    resnet18, torch.ones(1, 3, 224, 224), output_type="linalg-on-tensors"
+)
+backend = refbackend.RefBackendLinalgOnTensorsBackend()
+compiled = backend.compile(module)
+jit_module = backend.load(compiled)
+
+imgPH = TensorPlaceholder([-1, 3, 224, 224], torch.float32)
+#imgPH = TensorPlaceholder([1, 3, 224, 224], torch.float32)
+
+kModule = torchscript.compile(resnet18, imgPH, output_type="linalg-on-tensors")
+kBackend = KokkosBackend.KokkosBackend(dump_mlir=True)
+kCompiledModule = kBackend.compile(kModule)
+
+predictions(resnet18.forward, kCompiledModule.forward, img, labels)
+

examples/pytaco_cpp_driver/main.cpp

@@ -0,0 +1,55 @@
+# 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.
+
+import sys
+from pathlib import Path
+
+import torch
+import torchvision.models as models
+from torch_mlir import torchscript
+from torch_mlir.compiler_utils import TensorPlaceholder
+from torch_mlir_e2e_test.linalg_on_tensors_backends import refbackend
+from lapis import KokkosBackend
+
+sys.path.append(str(Path(__file__).absolute().parent))
+from _example_utils import (
+    top3_possibilities,
+    load_and_preprocess_image,
+    load_labels,
+    DEFAULT_IMAGE_URL,
+)
+
+def predictions(torch_func, kokkos_func, img, labels):
+    golden_prediction = top3_possibilities(torch_func(img), labels)
+    print("PyTorch prediction")
+    print(golden_prediction)
+    prediction = top3_possibilities(torch.from_numpy(kokkos_func(img.numpy())), labels)
+    print("LAPIS prediction")
+    print(prediction)
+
+print("load image from " + DEFAULT_IMAGE_URL, file=sys.stderr)
+img = load_and_preprocess_image(DEFAULT_IMAGE_URL)
+labels = load_labels()
+#print("Dumping preprocessed dog image to dog.bin")
+#img.numpy().tofile('dog.bin')
+
+resnet18 = models.resnet18(pretrained=True).eval()
+#print(help(torchscript.compile))
+
+module = torchscript.compile(
+    resnet18, torch.ones(1, 3, 224, 224), output_type="linalg-on-tensors"
+)
+backend = refbackend.RefBackendLinalgOnTensorsBackend()
+compiled = backend.compile(module)
+jit_module = backend.load(compiled)
+
+imgPH = TensorPlaceholder([1, 3, 224, 224], torch.float32)
+
+kModule = torchscript.compile(resnet18, imgPH, output_type="linalg-on-tensors")
+kBackend = KokkosBackend.KokkosBackend(dump_mlir=True)
+kCompiledModule = kBackend.compile(kModule)
+
+predictions(resnet18.forward, kCompiledModule.forward, img, labels)
+