diff --git a/apps/stable_diffusion/profiling_with_iree.md b/apps/stable_diffusion/profiling_with_iree.md index 424fd0dd..0ae3844d 100644 --- a/apps/stable_diffusion/profiling_with_iree.md +++ b/apps/stable_diffusion/profiling_with_iree.md @@ -10,7 +10,7 @@ Vulkan AMD: iree-compile --iree-input-type=none --iree-hal-target-backends=vulkan --iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 /path/to/input/mlir -o /path/to/output/vmfb # add --mlir-print-debuginfo --mlir-print-op-on-diagnostic=true for debug -# use –iree-input-type=mhlo for tf models +# use –iree-input-type=auto or "mhlo_legacy" or "stablehlo" for TF models CUDA NVIDIA: iree-compile --iree-input-type=none --iree-hal-target-backends=cuda --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 /path/to/input/mlir -o /path/to/output/vmfb diff --git a/cpp/README.md b/cpp/README.md index 60f75328..59b89305 100644 --- a/cpp/README.md +++ b/cpp/README.md @@ -40,7 +40,7 @@ cmake --build build/ *Prepare the model* ```bash wget https://storage.googleapis.com/shark_tank/latest/resnet50_tf/resnet50_tf.mlir -iree-compile --iree-input-type=mhlo --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --iree-llvmcpu-embedded-linker-path=`python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])'`/iree/compiler/tools/../_mlir_libs/iree-lld --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --mlir-pass-pipeline-crash-reproducer=ist/core-reproducer.mlir --iree-llvmcpu-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 resnet50_tf.mlir -o resnet50_tf.vmfb +iree-compile --iree-input-type=auto --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --iree-llvmcpu-embedded-linker-path=`python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])'`/iree/compiler/tools/../_mlir_libs/iree-lld --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --mlir-pass-pipeline-crash-reproducer=ist/core-reproducer.mlir --iree-llvmcpu-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 resnet50_tf.mlir -o resnet50_tf.vmfb ``` *Prepare the input* @@ -65,18 +65,18 @@ A tool for benchmarking other models is built and can be invoked with a command see `./build/vulkan_gui/iree-vulkan-gui --help` for an explanation on the function input. For example, stable diffusion unet can be tested with the following commands: ```bash wget https://storage.googleapis.com/shark_tank/quinn/stable_diff_tf/stable_diff_tf.mlir -iree-compile --iree-input-type=mhlo --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvmcpu-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 stable_diff_tf.mlir -o stable_diff_tf.vmfb +iree-compile --iree-input-type=auto --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvmcpu-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 stable_diff_tf.mlir -o stable_diff_tf.vmfb ./build/vulkan_gui/iree-vulkan-gui --module-file=stable_diff_tf.vmfb --function_input=2x4x64x64xf32 --function_input=1xf32 --function_input=2x77x768xf32 ``` VAE and Autoencoder are also available ```bash # VAE wget https://storage.googleapis.com/shark_tank/quinn/stable_diff_tf/vae_tf/vae.mlir -iree-compile --iree-input-type=mhlo --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvmcpu-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 vae.mlir -o vae.vmfb +iree-compile --iree-input-type=auto --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvmcpu-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 vae.mlir -o vae.vmfb ./build/vulkan_gui/iree-vulkan-gui --module-file=stable_diff_tf.vmfb --function_input=1x4x64x64xf32 # CLIP Autoencoder wget https://storage.googleapis.com/shark_tank/quinn/stable_diff_tf/clip_tf/clip_autoencoder.mlir -iree-compile --iree-input-type=mhlo --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvmcpu-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 clip_autoencoder.mlir -o clip_autoencoder.vmfb +iree-compile --iree-input-type=auto --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvmcpu-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 clip_autoencoder.mlir -o clip_autoencoder.vmfb ./build/vulkan_gui/iree-vulkan-gui --module-file=stable_diff_tf.vmfb --function_input=1x77xi32 --function_input=1x77xi32 ``` diff --git a/cpp/vision_inference/CMakeLists.txt b/cpp/vision_inference/CMakeLists.txt index 15922392..03c6a58b 100644 --- a/cpp/vision_inference/CMakeLists.txt +++ b/cpp/vision_inference/CMakeLists.txt @@ -21,7 +21,7 @@ endif() # Compile mnist.mlir to mnist.vmfb. set(_COMPILE_TOOL_EXECUTABLE $) set(_COMPILE_ARGS) -list(APPEND _COMPILE_ARGS "--iree-input-type=mhlo") +list(APPEND _COMPILE_ARGS "--iree-input-type=auto") list(APPEND _COMPILE_ARGS "--iree-hal-target-backends=llvm-cpu") list(APPEND _COMPILE_ARGS "${IREE_SOURCE_DIR}/samples/models/mnist.mlir") list(APPEND _COMPILE_ARGS "-o") diff --git a/shark/iree_utils/compile_utils.py b/shark/iree_utils/compile_utils.py index f0df8648..a2b22224 100644 --- a/shark/iree_utils/compile_utils.py +++ b/shark/iree_utils/compile_utils.py @@ -60,7 +60,7 @@ def get_iree_device_args(device, extra_args=[]): def get_iree_frontend_args(frontend): if frontend in ["torch", "pytorch", "linalg", "tm_tensor"]: return ["--iree-llvmcpu-target-cpu-features=host"] - elif frontend in ["tensorflow", "tf", "mhlo"]: + elif frontend in ["tensorflow", "tf", "mhlo", "stablehlo"]: return [ "--iree-llvmcpu-target-cpu-features=host", "--iree-mhlo-demote-i64-to-i32=false", @@ -265,8 +265,8 @@ def compile_module_to_flatbuffer( args += extra_args if frontend in ["tensorflow", "tf"]: - input_type = "mhlo" - elif frontend in ["mhlo", "tosa"]: + input_type = "auto" + elif frontend in ["stablehlo", "tosa"]: input_type = frontend elif frontend in ["tflite", "tflite-tosa"]: input_type = "tosa" @@ -367,7 +367,7 @@ def export_iree_module_to_vmfb( def export_module_to_mlir_file(module, frontend, directory: str): # TODO: write proper documentation. mlir_str = module - if frontend in ["tensorflow", "tf", "mhlo", "tflite"]: + if frontend in ["tensorflow", "tf", "mhlo", "stablehlo", "tflite"]: mlir_str = module.decode("utf-8") elif frontend in ["pytorch", "torch"]: mlir_str = module.operation.get_asm() diff --git a/shark/shark_runner.py b/shark/shark_runner.py index bc8bbc79..679a7073 100644 --- a/shark/shark_runner.py +++ b/shark/shark_runner.py @@ -25,7 +25,14 @@ import sys # supported dialects by the shark-runtime. -supported_dialects = {"linalg", "mhlo", "tosa", "tf-lite", "tm_tensor"} +supported_dialects = { + "linalg", + "mhlo", + "stablehlo", + "tosa", + "tf-lite", + "tm_tensor", +} class SharkRunner: diff --git a/shark/shark_trainer.py b/shark/shark_trainer.py index 27788e1f..abf242b7 100644 --- a/shark/shark_trainer.py +++ b/shark/shark_trainer.py @@ -59,6 +59,7 @@ class SharkTrainer: "torch", "tensorflow", "tf", + "stablehlo", "mhlo", "linalg", "tosa", @@ -84,7 +85,7 @@ class SharkTrainer: "tm_tensor", extra_args=extra_args, ) - elif self.frontend in ["tensorflow", "tf", "mhlo"]: + elif self.frontend in ["tensorflow", "tf", "mhlo", "stablehlo"]: self.shark_runner = SharkRunner( self.model, self.input, diff --git a/tank/examples/MiniLM_tf/huggingface_MiniLM_run.py b/tank/examples/MiniLM_tf/huggingface_MiniLM_run.py index 9caa19e2..40e97604 100644 --- a/tank/examples/MiniLM_tf/huggingface_MiniLM_run.py +++ b/tank/examples/MiniLM_tf/huggingface_MiniLM_run.py @@ -75,7 +75,7 @@ if __name__ == "__main__": compiler_module, target_backends=[backend], extra_args=args, - input_type="mhlo", + input_type="auto", ) # flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"]) diff --git a/tank/examples/bert_fine_tuning/bert_fine_tune_tf.py b/tank/examples/bert_fine_tuning/bert_fine_tune_tf.py index 649854b1..3dd0167c 100644 --- a/tank/examples/bert_fine_tuning/bert_fine_tune_tf.py +++ b/tank/examples/bert_fine_tuning/bert_fine_tune_tf.py @@ -153,7 +153,7 @@ if __name__ == "__main__": compiler_module, target_backends=[backend], extra_args=args, - input_type="mhlo", + input_type="auto", ) # Save module as MLIR file in a directory diff --git a/tank/examples/bert_tf/bert_large_run.py b/tank/examples/bert_tf/bert_large_run.py index 3d9f20c5..a99e4f70 100644 --- a/tank/examples/bert_tf/bert_large_run.py +++ b/tank/examples/bert_tf/bert_large_run.py @@ -96,7 +96,7 @@ if __name__ == "__main__": compiler_module, target_backends=[backend], extra_args=args, - input_type="mhlo", + input_type="auto", ) # flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"]) diff --git a/tank/examples/bert_tf/bert_small_run.py b/tank/examples/bert_tf/bert_small_run.py index a2322a81..5668ccf3 100644 --- a/tank/examples/bert_tf/bert_small_run.py +++ b/tank/examples/bert_tf/bert_small_run.py @@ -91,7 +91,7 @@ if __name__ == "__main__": compiler_module, target_backends=[backend], extra_args=args, - input_type="mhlo", + input_type="auto", ) # flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"])