Synchronize submodules with LLVM at llvm/llvm-project@6a7a2ee8161d Updates LLVM dependencies to match [6a7a2ee8161d](https://github.com/llvm/llvm-project/commit/6a7a2ee8161d). - TensorFlow to [4f8db85aa9ab](https://github.com/tensorflow/tensorflow/commit/4f8db85aa9ab) - MLIR-HLO to [c65dc7a455dc](https://github.com/tensorflow/mlir-hlo/commit/c65dc7a455dc) `./scripts/git/update_to_llvm_syncpoint.py ` PiperOrigin-RevId: 387433820
diff --git a/.gitignore b/.gitignore index 4f52c93..6fe9bf2 100644 --- a/.gitignore +++ b/.gitignore
@@ -58,3 +58,6 @@ # Generated documentation files mkdocs/site/ docs/website/site/ + +# Temporary files +iree/compiler/Dialect/HAL/Target/LLVM/librt/bin/librt.ll
diff --git a/CMakeLists.txt b/CMakeLists.txt index c466088..e4df9d5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt
@@ -322,6 +322,23 @@ endif() #------------------------------------------------------------------------------- +# Check if git submodules have been initialized. +# This will only run if python3 is available. +#------------------------------------------------------------------------------- + +find_package(Python3 COMPONENTS Interpreter QUIET) +if(Python3_FOUND) + execute_process( + COMMAND ${Python3_EXECUTABLE} scripts/git/check_submodule_init.py + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} + RESULT_VARIABLE ret + ) + if(NOT ret EQUAL "0") + message(FATAL_ERROR "check_submodule_init.py failed, see the logs above") + endif() +endif() + +#------------------------------------------------------------------------------- # MLIR/LLVM Dependency # We treat the LLVM dependency specially because we support several different # ways to use it:
diff --git a/bindings/python/iree/compiler/xla.py b/bindings/python/iree/compiler/xla.py index eb04231..a0b0578 100644 --- a/bindings/python/iree/compiler/xla.py +++ b/bindings/python/iree/compiler/xla.py
@@ -74,6 +74,7 @@ import_format: Union[ImportFormat, str] = ImportFormat.BINARY_PROTO, import_extra_args: Sequence[str] = (), + save_temp_mhlo_input: Optional[str] = None, save_temp_iree_input: Optional[str] = None, **kwargs): """Initialize options from keywords. @@ -87,6 +88,7 @@ self.import_only = import_only self.import_format = ImportFormat.parse(import_format) self.import_extra_args = import_extra_args + self.save_temp_mhlo_input = save_temp_mhlo_input self.save_temp_iree_input = save_temp_iree_input @@ -121,6 +123,10 @@ cl.append("--mlir-print-op-generic") # Save temps flags. + save_mhlo_input = tfs.alloc_optional("tf-mhlo.mlir", + export_as=options.save_temp_mhlo_input) + if save_mhlo_input: + cl.append(f"--save-temp-mhlo-input={save_mhlo_input}") iree_input = tfs.alloc_optional("xla-iree-input.mlir", export_as=options.save_temp_iree_input) if iree_input:
diff --git a/build_tools/buildkite/cmake/build_configurations.yml b/build_tools/buildkite/cmake/build_configurations.yml index fb79ee9..426e5fb 100644 --- a/build_tools/buildkite/cmake/build_configurations.yml +++ b/build_tools/buildkite/cmake/build_configurations.yml
@@ -5,7 +5,7 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception steps: - - label: "Build with tracing enabled" + - label: ":zap: Build with tracing enabled" commands: - "./scripts/git/submodule_versions.py init" - "docker run --user=$(id -u):$(id -g) --volume=\\$PWD:\\$IREE_DOCKER_WORKDIR --workdir=\\$IREE_DOCKER_WORKDIR --rm gcr.io/iree-oss/cmake@sha256:9d9953acf5ca0cf1ff3e8de32f10f24dfab1c4e8ec5d1fc047f556024ee4bed6 ./build_tools/cmake/build_tracing.sh" @@ -14,7 +14,7 @@ agents: - "queue=build" - - label: "Build the runtime only" + - label: ":hammer_and_wrench: Build the runtime only" commands: - "./scripts/git/submodule_versions.py init" - "docker run --user=$(id -u):$(id -g) --volume=\\$PWD:\\$IREE_DOCKER_WORKDIR --workdir=\\$IREE_DOCKER_WORKDIR --rm gcr.io/iree-oss/cmake@sha256:9d9953acf5ca0cf1ff3e8de32f10f24dfab1c4e8ec5d1fc047f556024ee4bed6 ./build_tools/cmake/build_runtime.sh" @@ -22,3 +22,27 @@ IREE_DOCKER_WORKDIR: "/usr/src/github/iree" agents: - "queue=build" + + - label: ":linux: Build host install" + key: "build-host-install" + commands: + - "./scripts/git/submodule_versions.py init" + - "docker run --user=$(id -u):$(id -g) --volume=\\$PWD:\\$IREE_DOCKER_WORKDIR --workdir=\\$IREE_DOCKER_WORKDIR --rm gcr.io/iree-oss/cmake@sha256:9d9953acf5ca0cf1ff3e8de32f10f24dfab1c4e8ec5d1fc047f556024ee4bed6 ./build_tools/cmake/build_host_install.sh" + - "tar -czvf build-artifacts.tgz build-host/install" + artifact_paths: "build-artifacts.tgz" + env: + IREE_DOCKER_WORKDIR: "/usr/src/github/iree" + agents: + - "queue=build" + + - label: ":webassembly: Build WebAssembly runtime with Emscripten" + depends_on: "build-host-install" + commands: + - "buildkite-agent artifact download --step build-host-install build-artifacts.tgz ./" + - "tar xzf build-artifacts.tgz" + - "./scripts/git/submodule_versions.py init" + - "docker run --user=$(id -u):$(id -g) --volume=\\$PWD:\\$IREE_DOCKER_WORKDIR --workdir=\\$IREE_DOCKER_WORKDIR --rm gcr.io/iree-oss/cmake-emscripten@sha256:8acad361d23cb586187c2ea29df3a1ab301b5283c3648beb328681d69ecd0ab0 ./build_tools/cmake/build_runtime_emscripten.sh" + env: + IREE_DOCKER_WORKDIR: "/usr/src/github/iree" + agents: + - "queue=build"
diff --git a/build_tools/cmake/build_host_install.sh b/build_tools/cmake/build_host_install.sh new file mode 100755 index 0000000..b5f8130 --- /dev/null +++ b/build_tools/cmake/build_host_install.sh
@@ -0,0 +1,36 @@ +#!/bin/bash +# Copyright 2021 The IREE Authors +# +# 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 + +# Builds host binaries (compiler tools and other utilities) and installs them +# in build-host/install. Designed for CI, but can be run manually. This uses +# previously cached build results and does not clear build directories. + +set -x +set -e + +CMAKE_BIN=${CMAKE_BIN:-$(which cmake)} +"${CMAKE_BIN?}" --version +ninja --version + +ROOT_DIR=$(git rev-parse --show-toplevel) +cd ${ROOT_DIR?} + +if [ -d "build-host" ] +then + echo "build-host directory already exists. Will use cached results there." +else + echo "build-host directory does not already exist. Creating a new one." + mkdir build-host +fi +cd build-host + +# Configure, build, install. +"${CMAKE_BIN?}" -G Ninja .. \ + -DCMAKE_INSTALL_PREFIX=./install \ + -DIREE_BUILD_TESTS=OFF \ + -DIREE_BUILD_SAMPLES=OFF +"${CMAKE_BIN?}" --build . --target install
diff --git a/build_tools/cmake/build_runtime_emscripten.sh b/build_tools/cmake/build_runtime_emscripten.sh new file mode 100755 index 0000000..ca2abc3 --- /dev/null +++ b/build_tools/cmake/build_runtime_emscripten.sh
@@ -0,0 +1,51 @@ +#!/bin/bash +# Copyright 2021 The IREE Authors +# +# 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 + +# Cross-compile IREE's runtime through Emscripten to WebAssembly with CMake. +# Designed for CI, but can be run manually. This uses previously cached build +# results and does not clear build directories. +# +# Host binaries (e.g. compiler tools) should already be built at +# ./build-host/install. Emscripten binaries (e.g. .wasm and .js files) will be +# built in ./build-emscripten/. + +set -x +set -e + +if ! command -v emcmake &> /dev/null +then + echo "'emcmake' not found, setup environment according to https://emscripten.org/docs/getting_started/downloads.html" + exit +fi + +CMAKE_BIN=${CMAKE_BIN:-$(which cmake)} +"${CMAKE_BIN?}" --version +ninja --version + +ROOT_DIR=$(git rev-parse --show-toplevel) +cd ${ROOT_DIR?} + +if [ -d "build-emscripten" ] +then + echo "build-emscripten directory already exists. Will use cached results there." +else + echo "build-emscripten directory does not already exist. Creating a new one." + mkdir build-emscripten +fi +cd build-emscripten + +# Configure using Emscripten's CMake wrapper, then build. +emcmake "${CMAKE_BIN?}" -G Ninja .. \ + -DIREE_HOST_BINARY_ROOT=$PWD/../build-host/install \ + -DIREE_HAL_DRIVERS_TO_BUILD=VMVX\;DyLib \ + -DIREE_BUILD_COMPILER=OFF \ + -DIREE_ENABLE_MLIR=OFF \ + -DIREE_BUILD_TESTS=OFF \ + -DIREE_BUILD_SAMPLES=ON + +# TODO(scotttodd): expand this list of targets +"${CMAKE_BIN?}" --build . --target iree_samples_simple_embedding_simple_embedding_vmvx_sync
diff --git a/build_tools/docker/cmake-emscripten/Dockerfile b/build_tools/docker/cmake-emscripten/Dockerfile new file mode 100644 index 0000000..4f3f600 --- /dev/null +++ b/build_tools/docker/cmake-emscripten/Dockerfile
@@ -0,0 +1,44 @@ +# Copyright 2021 The IREE Authors +# +# 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 + +# An image for building IREE through Emscripten using CMake. + +FROM gcr.io/iree-oss/cmake@sha256:9d9953acf5ca0cf1ff3e8de32f10f24dfab1c4e8ec5d1fc047f556024ee4bed6 AS final + +# See also +# * https://github.com/emscripten-core/emsdk/blob/main/docker/Dockerfile +# * https://hub.docker.com/r/emscripten/emsdk + +ARG EMSDK_COMMIT=5c0e31a03a72136ccaa761c408baf5a640f942ec +ARG SDK_VERSION=2.0.25 + +# Follow https://emscripten.org/docs/getting_started/downloads.html. +RUN git clone https://github.com/emscripten-core/emsdk +RUN cd emsdk && git checkout "${EMSDK_COMMIT?}" && \ + ./emsdk install ${SDK_VERSION?} && \ + ./emsdk activate ${SDK_VERSION?} + +# Set some environment variables for Emscripten to use. +ENV EMSDK=/emsdk +ENV EM_DATA=${EMSDK}/.data +ENV EM_CONFIG=${EMSDK}/.emscripten +ENV EM_CACHE=${EM_DATA}/cache +ENV EM_PORTS=${EM_DATA}/ports +# Emscripten writes into its cache location (outside of the CMake build +# directory). +# We can either +# (A) Grant broad write permissions to the cache directory to be able to run +# our scripts under different users. +# (B) Mount a user home directory when using the image. +# Since (A) requires less configuration, we'll do that. If multiple tools would +# want a user directory (like Bazel), we should switch to (B). +# See https://github.com/emscripten-core/emsdk/issues/535 +RUN mkdir -p ${EM_CACHE} && chmod -R 777 ${EM_CACHE} + +# Normally we'd run `source emsdk_env.sh`, but that doesn't integrate with +# Docker's environment properties model. Instead, we directly extend the path +# to include the directories suggested by `emsdk activate`. +ENV PATH="$PWD/emsdk:$PWD/emsdk/node/14.15.5_64bit/bin:$PWD/emsdk/upstream/emscripten:$PATH"
diff --git a/build_tools/docker/manage_images.py b/build_tools/docker/manage_images.py index 9e9bfd8..ce61e2f 100755 --- a/build_tools/docker/manage_images.py +++ b/build_tools/docker/manage_images.py
@@ -46,6 +46,7 @@ 'bazel': ['base', 'util'], 'cmake': ['base', 'util'], 'cmake-android': ['cmake-python', 'util'], + 'cmake-emscripten': ['cmake'], 'cmake-python': ['cmake'], 'cmake-python-vulkan': ['cmake-python', 'vulkan'], 'cmake-python-swiftshader': ['cmake-python-vulkan', 'swiftshader'],
diff --git a/build_tools/docker/prod_digests.txt b/build_tools/docker/prod_digests.txt index f5e1f59..dd0a5b7 100644 --- a/build_tools/docker/prod_digests.txt +++ b/build_tools/docker/prod_digests.txt
@@ -17,3 +17,4 @@ gcr.io/iree-oss/cmake-riscv@sha256:95489593bc9b0cd325ce9c1a32b47389c01b174a5b8190a16d937d2e8828d384 gcr.io/iree-oss/cmake-bazel-frontends-android@sha256:1392e3a27cddbdc597817168fb61e125bbdcbfd9076eff9d70bd8012b0a0c5ba gcr.io/iree-oss/samples@sha256:be5465585706b620d6c722caa6237eafdfaa8dd11ce20db0981b979f2d3387b3 +gcr.io/iree-oss/cmake-emscripten@sha256:8acad361d23cb586187c2ea29df3a1ab301b5283c3648beb328681d69ecd0ab0
diff --git a/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader-asan/build.sh b/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader-asan/build.sh index bd48c3f..1f5b793 100755 --- a/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader-asan/build.sh +++ b/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader-asan/build.sh
@@ -88,6 +88,7 @@ fi if [[ "${IREE_CUDA_DISABLE?}" == 1 ]]; then label_exclude_args+=("^driver=cuda$") + label_exclude_args+=("^uses_cuda_runtime$") fi if [[ "${IREE_VULKAN_F16_DISABLE?}" == 1 ]]; then label_exclude_args+=("^vulkan_uses_vk_khr_shader_float16_int8$")
diff --git a/docs/developers/best_practices.md b/docs/developers/best_practices.md new file mode 100644 index 0000000..d3d9c91 --- /dev/null +++ b/docs/developers/best_practices.md
@@ -0,0 +1,63 @@ +# IREE Best Practices + +This page contains a list of best practices for getting the most out of IREE, +spanning model authoring, ahead-of-time compilation, and runtime use. Treat +these as a collection of ideas to consider or areas to start benchmarking when +working on your own applications. + +## Introduction + +Common themes include: + +* Give the compiler as much information as possible +* Give the compiler opportunities to batch work together or defer computation +* Keep compute devices saturated with work through pipelining +* Use dense math where possible, particularly for inner loop bodies +* Limit synchronization points between devices like CPUs and GPUs +* Profile early and often, using the right tools for each level of granularity + +## Practices for model authoring + +### Track state within your model when possible + +If your model is stateful prefer to store that state directly within your +program rather than externalizing it through arguments and return values. By +keeping state inside your program the compiler is better able to reason about +it and function calls will have lower overhead. + +If you do externalize state, try to pack that state into a limited number of +arguments. + +See the +[variables and state](https://github.com/google/iree/tree/main/iree/samples/variables_and_state) +sample for further guidance on tracking and using state. + +### Limit uses of dynamic shapes + +While IREE aims to support general dynamic shapes use, it is better able to +optimize parts of programs where shapes are static. Slow varying dimensions +like batch index or timestamp are safer uses of dynamic shapes than faster +varying dimensions like the x/y/channel dimensions of images. + +See the +[dynamic shapes](https://github.com/google/iree/tree/main/iree/samples/dynamic_shapes) +sample for further guidance on using dynamic shapes. + +## Practices for compilation settings + +TODO: mention parameters to tune + +TODO: which compiler targets to use (try both CUDA and Vulkan?) + +TODO: use the most specific LLVM target triple you can? + +## Practices for runtime use + +### Do the minimum amount of work: cache queries and reuse buffers + +Try to front-load queries, particularly queries using strings that look up into +maps like `iree_runtime_session_call_by_name`, so that hot sections of code are +doing the minimum amount of work: routing inputs through buffers, scheduling +runtime calls, and routing outputs through other buffers. + +TODO: sample code, profile numbers
diff --git a/docs/developers/get_started/getting_started_emscripten.md b/docs/developers/get_started/getting_started_emscripten.md new file mode 100644 index 0000000..c5aefb7 --- /dev/null +++ b/docs/developers/get_started/getting_started_emscripten.md
@@ -0,0 +1,67 @@ +# Getting Started With Emscripten + +[Emscripten](https://emscripten.org/index.html) is a complete compiler +toolchain to WebAssembly, using LLVM, with a special focus on speed, size, and +the Web platform. Emscripten can be used to compile parts of IREE to +[WebAssembly](https://webassembly.org/) for execution within web browsers or +other Wasm runtimes. + +## Status + +IREE's _runtime_ can be compiled through Emscripten in some limited +configurations. More of the runtime will be supported over time. + +IREE's _compiler_ can be compiled through Emscripten with local changes. More +work is needed for this to be generally supported. + +## Prerequisites + +Read https://emscripten.org/docs/getting_started/downloads.html and run + +``` +./emsdk install latest +./emsdk activate latest +source ./emsdk_env.sh +``` + +## Building IREE's Runtime with Emscripten + +### Host Configuration + +Build and install at least the compiler tools on your host machine, or install +them from a binary distribution: + +```shell +$ cmake -G Ninja -B ../iree-build-host/ \ + -DCMAKE_C_COMPILER=clang \ + -DCMAKE_CXX_COMPILER=clang++ \ + -DCMAKE_INSTALL_PREFIX=../iree-build-host/install \ + . +$ cmake --build ../iree-build-host/ --target install +``` + +### Target Configuration + +```shell +$ emcmake cmake -G Ninja -B ../iree-build-emscripten/ \ + -DCMake_BUILD_TYPE=Release \ + -DIREE_HOST_BINARY_ROOT=$(realpath ../iree-build-host/install) \ + -DIREE_BUILD_TESTS=OFF \ + -DIREE_BUILD_COMPILER=OFF \ + -DIREE_ENABLE_MLIR=OFF \ + . +``` + +Build: + +``` +cmake --build ../iree-build-emscripten/ \ + --target iree_samples_simple_embedding_simple_embedding_vmvx_sync +``` + +### Load into a WebAssembly Environment + +Copy the outputs from the build process (e.g. `simple_embedding_vmvx_sync.js` +and `simple_embedding_vmvx_sync.wasm`) into your application and follow +instructions at either https://webassembly.org/getting-started/developers-guide/ +or https://developer.mozilla.org/en-US/docs/WebAssembly/Loading_and_running.
diff --git a/docs/website/docs/bindings/python.md b/docs/website/docs/bindings/python.md index b6faade..f279ad6 100644 --- a/docs/website/docs/bindings/python.md +++ b/docs/website/docs/bindings/python.md
@@ -10,11 +10,16 @@ | `iree-tools-tf-snapshot` | Tools for importing from [TensorFlow](https://www.tensorflow.org/) | | `iree-tools-tflite-snapshot` | Tools for importing from [TensorFlow Lite](https://www.tensorflow.org/lite) | | `iree-tools-xla-snapshot` | Tools for importing from [XLA](https://www.tensorflow.org/xla) | +| `iree-jax-snapshot` | Tools for importing from [JAX](https://github.com/google/jax) | Collectively, these packages allow for importing from frontends, compiling towards various targets, and executing compiled code on IREE's backends. -<!-- TODO(??): Which package for JAX? --> +!!! warning + The TensorFlow, TensorFlow Lite, and XLA packages are currently only + available on Linux and macOS. They are not available on Windows yet (see + [this issue](https://github.com/google/iree/issues/6417)). + <!-- TODO(??): API references for packages/modules --> <!-- TODO(??): at least link to source code and sample Colab notebooks for now --> <!-- TODO(??): link to frontend docs -->
diff --git a/docs/website/docs/blog/2021-07-19-tflite-tosa-compilation-diagram.png b/docs/website/docs/blog/2021-07-19-tflite-tosa-compilation-diagram.png new file mode 100644 index 0000000..11b5986 --- /dev/null +++ b/docs/website/docs/blog/2021-07-19-tflite-tosa-compilation-diagram.png Binary files differ
diff --git a/docs/website/docs/blog/2021-07-19-tflite-tosa.md b/docs/website/docs/blog/2021-07-19-tflite-tosa.md new file mode 100644 index 0000000..90929b6 --- /dev/null +++ b/docs/website/docs/blog/2021-07-19-tflite-tosa.md
@@ -0,0 +1,44 @@ + Monday, July 19, 2021<br> + By Rob Suderman and Jenni Kilduff + +## TFLite Support via TOSA + +IREE can now execute [TensorFlow Lite](https://www.tensorflow.org/lite) +(TFLite) models through the use of +[TOSA](https://developer.mlplatform.org/w/tosa/), an open standard of common +tensor operations, and a part of [MLIR](https://mlir.llvm.org/) core. TOSA’s +high-level representation of tensor operations provides a common front-end for +ingesting models from different frameworks. In this case we ingest a TFLite +flatbuffer and compile it to TOSA IR, which IREE takes as an input format to +compile to its various backends. + +{ align=left } + +Using TFLite as a frontend for IREE provides an alternative ingestion method for +already existing models that could benefit from IREE’s design. This enables +models already designed for on-device inference to have an alternative path for +execution without requiring any additional porting, while benefiting from +IREE’s improvements in buffer management, work dispatch system, and compact +binary format. With continued improvements to IREE/MLIR’s compilation +performance, more optimized versions can be compiled and distributed to target +devices without an update to the clientside environment. + +Today, we have validated floating point support for a variety of models, +including +[mobilenet](https://tfhub.dev/s?deployment-format=lite&network-architecture=mobilenet,mobilenet-v2,mobilenet-v3,mobilenet-v1&q=mobilenet) +(v1, v2, and v3) and +[mobilebert](https://tfhub.dev/tensorflow/lite-model/mobilebert/1/default/1). +More work is in progress to support fully quantized models, and TFLite’s hybrid +quantization, along with dynamic shape support. + +## Examples + +TFLite with IREE is available in Python and Java. We have a +[colab notebook](https://colab.sandbox.google.com/github/google/iree/blob/main/colab/tflite_text_classification.ipynb) +that shows how to use IREE’s python bindings and TFLite compiler tools to +compile a pre-trained TFLite model from a flatbuffer and run using IREE. We +also have an +[Android Java app](https://github.com/not-jenni/iree-android-tflite-demo) that +was forked from an existing TFLite demo app, swapping out the TFLite library +for our own AAR. More information on IREE’s TFLite frontend is available +[here](../ml-frameworks/tensorflow-lite.md).
diff --git a/docs/website/docs/deployment-configurations/cpu-dylib.md b/docs/website/docs/deployment-configurations/cpu-dylib.md index 706d039..8bad3be 100644 --- a/docs/website/docs/deployment-configurations/cpu-dylib.md +++ b/docs/website/docs/deployment-configurations/cpu-dylib.md
@@ -131,11 +131,11 @@ <!-- TODO(??): troubleshooting --> -[android-cc]: /building-from-source/android/ -[get-started]: /building-from-source/getting-started/ +[android-cc]: ../building-from-source/android.md +[get-started]: ../building-from-source/getting-started.md [iree-releases]: https://github.com/google/iree/releases/ [llvm]: https://llvm.org/ [mlir]: https://mlir.llvm.org/ [tf-hub-mobilenetv2]: https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification -[tf-import]: /frontends/tensorflow/ -[tflite-import]: /frontends/tensorflow-lite/ +[tf-import]: ../ml-frameworks/tensorflow.md +[tflite-import]: ../ml-frameworks/tensorflow-lite.md
diff --git a/docs/website/docs/deployment-configurations/gpu-vulkan.md b/docs/website/docs/deployment-configurations/gpu-vulkan.md index d86de5c..a973bc3 100644 --- a/docs/website/docs/deployment-configurations/gpu-vulkan.md +++ b/docs/website/docs/deployment-configurations/gpu-vulkan.md
@@ -177,13 +177,13 @@ <!-- TODO(??): troubleshooting --> -[android-cc]: /building-from-source/android/ -[get-started]: /building-from-source/getting-started/ +[android-cc]: ../building-from-source/android.md +[get-started]: ../building-from-source/getting-started.md [iree-releases]: https://github.com/google/iree/releases/ [mlir]: https://mlir.llvm.org/ [spirv]: https://www.khronos.org/registry/spir-v/ [tf-hub-mobilenetv2]: https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification -[tf-import]: /frontends/tensorflow/ -[tflite-import]: /frontends/tensorflow-lite/ +[tf-import]: ../ml-frameworks/tensorflow.md +[tflite-import]: ../ml-frameworks/tensorflow-lite.md [vulkan]: https://www.khronos.org/vulkan/ [vulkan-sdk]: https://vulkan.lunarg.com/sdk/home/
diff --git a/docs/website/docs/ml-frameworks/tensorflow-lite.md b/docs/website/docs/ml-frameworks/tensorflow-lite.md index b91ac3f..e807555 100644 --- a/docs/website/docs/ml-frameworks/tensorflow-lite.md +++ b/docs/website/docs/ml-frameworks/tensorflow-lite.md
@@ -24,6 +24,11 @@ -f https://github.com/google/iree/releases ``` +!!! warning + The TensorFlow Lite package is currently only available on Linux and macOS. + It is not available on Windows yet (see + [this issue](https://github.com/google/iree/issues/6417)). + ## Importing models Fist, import the TFLite model to TOSA MLIR:
diff --git a/docs/website/docs/ml-frameworks/tensorflow.md b/docs/website/docs/ml-frameworks/tensorflow.md index 837f8ec..9e5df96 100644 --- a/docs/website/docs/ml-frameworks/tensorflow.md +++ b/docs/website/docs/ml-frameworks/tensorflow.md
@@ -27,6 +27,11 @@ -f https://github.com/google/iree/releases ``` +!!! warning + The TensorFlow package is currently only available on Linux and macOS. It + is not available on Windows yet (see + [this issue](https://github.com/google/iree/issues/6417)). + ## Importing models IREE compilers transform a model into its final deployable format in several
diff --git a/docs/website/mkdocs.yml b/docs/website/mkdocs.yml index a6dca81..c9253f3 100644 --- a/docs/website/mkdocs.yml +++ b/docs/website/mkdocs.yml
@@ -73,6 +73,7 @@ markdown_extensions: - abbr - admonition + - attr_list - footnotes - meta - pymdownx.details @@ -115,3 +116,5 @@ - TensorFlow Lite: 'bindings/tensorflow-lite.md' - 'Community': - Projects: 'community/projects.md' + - 'Blog': + - TFLite Support via TOSA: 'blog/2021-07-19-tflite-tosa.md'
diff --git a/experimental/rocm/descriptor_set_layout.c b/experimental/rocm/descriptor_set_layout.c index 69fb4af..c9e985e 100644 --- a/experimental/rocm/descriptor_set_layout.c +++ b/experimental/rocm/descriptor_set_layout.c
@@ -13,49 +13,58 @@ typedef struct iree_hal_rocm_descriptor_set_layout_t { iree_hal_resource_t resource; - iree_hal_rocm_context_wrapper_t *context; + iree_hal_rocm_context_wrapper_t* context; + iree_host_size_t binding_count; } iree_hal_rocm_descriptor_set_layout_t; extern const iree_hal_descriptor_set_layout_vtable_t iree_hal_rocm_descriptor_set_layout_vtable; -static iree_hal_rocm_descriptor_set_layout_t * +static iree_hal_rocm_descriptor_set_layout_t* iree_hal_rocm_descriptor_set_layout_cast( - iree_hal_descriptor_set_layout_t *base_value) { + iree_hal_descriptor_set_layout_t* base_value) { IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_rocm_descriptor_set_layout_vtable); - return (iree_hal_rocm_descriptor_set_layout_t *)base_value; + return (iree_hal_rocm_descriptor_set_layout_t*)base_value; } iree_status_t iree_hal_rocm_descriptor_set_layout_create( - iree_hal_rocm_context_wrapper_t *context, + iree_hal_rocm_context_wrapper_t* context, iree_hal_descriptor_set_layout_usage_type_t usage_type, iree_host_size_t binding_count, - const iree_hal_descriptor_set_layout_binding_t *bindings, - iree_hal_descriptor_set_layout_t **out_descriptor_set_layout) { + const iree_hal_descriptor_set_layout_binding_t* bindings, + iree_hal_descriptor_set_layout_t** out_descriptor_set_layout) { IREE_ASSERT_ARGUMENT(context); IREE_ASSERT_ARGUMENT(!binding_count || bindings); IREE_ASSERT_ARGUMENT(out_descriptor_set_layout); *out_descriptor_set_layout = NULL; IREE_TRACE_ZONE_BEGIN(z0); - iree_hal_rocm_descriptor_set_layout_t *descriptor_set_layout = NULL; + iree_hal_rocm_descriptor_set_layout_t* descriptor_set_layout = NULL; iree_status_t status = iree_allocator_malloc(context->host_allocator, sizeof(*descriptor_set_layout), - (void **)&descriptor_set_layout); + (void**)&descriptor_set_layout); if (iree_status_is_ok(status)) { iree_hal_resource_initialize(&iree_hal_rocm_descriptor_set_layout_vtable, &descriptor_set_layout->resource); descriptor_set_layout->context = context; + descriptor_set_layout->binding_count = binding_count; *out_descriptor_set_layout = - (iree_hal_descriptor_set_layout_t *)descriptor_set_layout; + (iree_hal_descriptor_set_layout_t*)descriptor_set_layout; } IREE_TRACE_ZONE_END(z0); return status; } +iree_host_size_t iree_hal_rocm_descriptor_set_layout_binding_count( + iree_hal_descriptor_set_layout_t* base_descriptor_set_layout) { + iree_hal_rocm_descriptor_set_layout_t* descriptor_set_layout = + iree_hal_rocm_descriptor_set_layout_cast(base_descriptor_set_layout); + return descriptor_set_layout->binding_count; +} + static void iree_hal_rocm_descriptor_set_layout_destroy( - iree_hal_descriptor_set_layout_t *base_descriptor_set_layout) { - iree_hal_rocm_descriptor_set_layout_t *descriptor_set_layout = + iree_hal_descriptor_set_layout_t* base_descriptor_set_layout) { + iree_hal_rocm_descriptor_set_layout_t* descriptor_set_layout = iree_hal_rocm_descriptor_set_layout_cast(base_descriptor_set_layout); iree_allocator_t host_allocator = descriptor_set_layout->context->host_allocator;
diff --git a/experimental/rocm/descriptor_set_layout.h b/experimental/rocm/descriptor_set_layout.h index fb07b76..9d58acd 100644 --- a/experimental/rocm/descriptor_set_layout.h +++ b/experimental/rocm/descriptor_set_layout.h
@@ -16,11 +16,15 @@ #endif // __cplusplus iree_status_t iree_hal_rocm_descriptor_set_layout_create( - iree_hal_rocm_context_wrapper_t *context, + iree_hal_rocm_context_wrapper_t* context, iree_hal_descriptor_set_layout_usage_type_t usage_type, iree_host_size_t binding_count, - const iree_hal_descriptor_set_layout_binding_t *bindings, - iree_hal_descriptor_set_layout_t **out_descriptor_set_layout); + const iree_hal_descriptor_set_layout_binding_t* bindings, + iree_hal_descriptor_set_layout_t** out_descriptor_set_layout); + +// Return the binding count for the given descriptor set layout. +iree_host_size_t iree_hal_rocm_descriptor_set_layout_binding_count( + iree_hal_descriptor_set_layout_t* descriptor_set_layout); #ifdef __cplusplus } // extern "C"
diff --git a/experimental/rocm/direct_command_buffer.c b/experimental/rocm/direct_command_buffer.c index 3d8706a..db28d91 100644 --- a/experimental/rocm/direct_command_buffer.c +++ b/experimental/rocm/direct_command_buffer.c
@@ -11,6 +11,7 @@ #include <stdint.h> #include "experimental/rocm/dynamic_symbols.h" +#include "experimental/rocm/executable_layout.h" #include "experimental/rocm/native_executable.h" #include "experimental/rocm/rocm_buffer.h" #include "experimental/rocm/status_util.h" @@ -44,7 +45,7 @@ return (iree_hal_rocm_direct_command_buffer_t*)base_value; } -iree_status_t iree_hal_rocm_direct_command_buffer_allocate( +iree_status_t iree_hal_rocm_direct_command_buffer_create( iree_hal_rocm_context_wrapper_t* context, iree_hal_command_buffer_mode_t mode, iree_hal_command_category_t command_categories, @@ -283,6 +284,8 @@ const iree_hal_descriptor_set_binding_t* bindings) { iree_hal_rocm_direct_command_buffer_t* command_buffer = iree_hal_rocm_direct_command_buffer_cast(base_command_buffer); + iree_host_size_t base_binding = + iree_hal_rocm_base_binding_index(executable_layout, set); // Convention with the compiler side. We map bindings to kernel argument. // We compact the bindings to get a dense set of arguments and keep them order // based on the binding index. @@ -303,7 +306,8 @@ iree_hal_rocm_buffer_device_pointer( iree_hal_buffer_allocated_buffer(binding.buffer)) + iree_hal_buffer_byte_offset(binding.buffer) + binding.offset; - *((hipDeviceptr_t*)command_buffer->current_descriptor[i]) = device_ptr; + *((hipDeviceptr_t*)command_buffer->current_descriptor[i + base_binding]) = + device_ptr; } return iree_ok_status(); }
diff --git a/experimental/rocm/direct_command_buffer.h b/experimental/rocm/direct_command_buffer.h index 0145b15..bd665bf 100644 --- a/experimental/rocm/direct_command_buffer.h +++ b/experimental/rocm/direct_command_buffer.h
@@ -26,16 +26,16 @@ unsigned int blockDimX; unsigned int blockDimY; unsigned int blockDimZ; - void **kernelParams; + void** kernelParams; } hip_launch_params; // Creates a rocm direct command buffer. -iree_status_t iree_hal_rocm_direct_command_buffer_allocate( - iree_hal_rocm_context_wrapper_t *context, +iree_status_t iree_hal_rocm_direct_command_buffer_create( + iree_hal_rocm_context_wrapper_t* context, iree_hal_command_buffer_mode_t mode, iree_hal_command_category_t command_categories, iree_hal_queue_affinity_t queue_affinity, - iree_hal_command_buffer_t **out_command_buffer); + iree_hal_command_buffer_t** out_command_buffer); #ifdef __cplusplus } // extern "C"
diff --git a/experimental/rocm/dynamic_symbol_tables.h b/experimental/rocm/dynamic_symbol_tables.h index fe6cfd5..aa78c5a 100644 --- a/experimental/rocm/dynamic_symbol_tables.h +++ b/experimental/rocm/dynamic_symbol_tables.h
@@ -30,6 +30,7 @@ RC_PFN_DECL(hipMemcpyAsync, void *, const void *, size_t, hipMemcpyKind, hipStream_t) RC_PFN_DECL(hipMalloc, void **, size_t) +RC_PFN_DECL(hipMallocManaged, hipDeviceptr_t *, size_t, unsigned int) RC_PFN_DECL(hipFree, void *) RC_PFN_DECL(hipHostFree, void *) RC_PFN_DECL(hipMemAllocHost, void **, size_t, unsigned int)
diff --git a/experimental/rocm/dynamic_symbols.c b/experimental/rocm/dynamic_symbols.c index d676838..ad7ca6d 100644 --- a/experimental/rocm/dynamic_symbols.c +++ b/experimental/rocm/dynamic_symbols.c
@@ -12,7 +12,7 @@ #include "iree/base/target_platform.h" #include "iree/base/tracing.h" -static const char *kROCMLoaderSearchNames[] = { +static const char* kROCMLoaderSearchNames[] = { #if defined(IREE_PLATFORM_WINDOWS) "amdhip64.dll", #else @@ -21,12 +21,12 @@ }; static iree_status_t iree_hal_rocm_dynamic_symbols_resolve_all( - iree_hal_rocm_dynamic_symbols_t *syms) { -#define RC_PFN_DECL(rocmSymbolName, ...) \ - { \ - static const char *kName = #rocmSymbolName; \ - IREE_RETURN_IF_ERROR(iree_dynamic_library_lookup_symbol( \ - syms->loader_library, kName, (void **)&syms->rocmSymbolName)); \ + iree_hal_rocm_dynamic_symbols_t* syms) { +#define RC_PFN_DECL(rocmSymbolName, ...) \ + { \ + static const char* kName = #rocmSymbolName; \ + IREE_RETURN_IF_ERROR(iree_dynamic_library_lookup_symbol( \ + syms->loader_library, kName, (void**)&syms->rocmSymbolName)); \ } #define RC_PFN_STR_DECL(rocmSymbolName, ...) RC_PFN_DECL(rocmSymbolName, ...) #include "experimental/rocm/dynamic_symbol_tables.h" // IWYU pragma: keep @@ -36,7 +36,7 @@ } iree_status_t iree_hal_rocm_dynamic_symbols_initialize( - iree_allocator_t allocator, iree_hal_rocm_dynamic_symbols_t *out_syms) { + iree_allocator_t allocator, iree_hal_rocm_dynamic_symbols_t* out_syms) { IREE_TRACE_ZONE_BEGIN(z0); memset(out_syms, 0, sizeof(*out_syms)); iree_status_t status = iree_dynamic_library_load_from_files( @@ -59,7 +59,7 @@ } void iree_hal_rocm_dynamic_symbols_deinitialize( - iree_hal_rocm_dynamic_symbols_t *syms) { + iree_hal_rocm_dynamic_symbols_t* syms) { IREE_TRACE_ZONE_BEGIN(z0); iree_dynamic_library_release(syms->loader_library); memset(syms, 0, sizeof(*syms));
diff --git a/experimental/rocm/dynamic_symbols.h b/experimental/rocm/dynamic_symbols.h index 00a58b0..9cbd774 100644 --- a/experimental/rocm/dynamic_symbols.h +++ b/experimental/rocm/dynamic_symbols.h
@@ -20,12 +20,12 @@ // any of the symbol is not available. The functions signatures are matching // the declarations in `hipruntime.h`. typedef struct iree_hal_rocm_dynamic_symbols_t { - iree_dynamic_library_t *loader_library; + iree_dynamic_library_t* loader_library; #define RC_PFN_DECL(rocmSymbolName, ...) \ hipError_t (*rocmSymbolName)(__VA_ARGS__); #define RC_PFN_STR_DECL(rocmSymbolName, ...) \ - const char *(*rocmSymbolName)(__VA_ARGS__); + const char* (*rocmSymbolName)(__VA_ARGS__); #include "experimental/rocm/dynamic_symbol_tables.h" // IWYU pragma: export #undef RC_PFN_DECL #undef RC_PFN_STR_DECL @@ -35,13 +35,13 @@ // iree_hal_rocm_dynamic_symbols_deinitialize must be used to release the // library resources. iree_status_t iree_hal_rocm_dynamic_symbols_initialize( - iree_allocator_t allocator, iree_hal_rocm_dynamic_symbols_t *out_syms); + iree_allocator_t allocator, iree_hal_rocm_dynamic_symbols_t* out_syms); // Deinitializes |syms| by unloading the backing library. All function pointers // will be invalidated. They _may_ still work if there are other reasons the // library remains loaded so be careful. void iree_hal_rocm_dynamic_symbols_deinitialize( - iree_hal_rocm_dynamic_symbols_t *syms); + iree_hal_rocm_dynamic_symbols_t* syms); #ifdef __cplusplus } // extern "C"
diff --git a/experimental/rocm/event_semaphore.c b/experimental/rocm/event_semaphore.c index 7fe6860..d3085fd 100644 --- a/experimental/rocm/event_semaphore.c +++ b/experimental/rocm/event_semaphore.c
@@ -13,34 +13,34 @@ typedef struct iree_hal_rocm_semaphore_t { iree_hal_resource_t resource; - iree_hal_rocm_context_wrapper_t *context; + iree_hal_rocm_context_wrapper_t* context; uint64_t initial_value; } iree_hal_rocm_semaphore_t; extern const iree_hal_semaphore_vtable_t iree_hal_rocm_semaphore_vtable; -static iree_hal_rocm_semaphore_t *iree_hal_rocm_semaphore_cast( - iree_hal_semaphore_t *base_value) { +static iree_hal_rocm_semaphore_t* iree_hal_rocm_semaphore_cast( + iree_hal_semaphore_t* base_value) { IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_rocm_semaphore_vtable); - return (iree_hal_rocm_semaphore_t *)base_value; + return (iree_hal_rocm_semaphore_t*)base_value; } iree_status_t iree_hal_rocm_semaphore_create( - iree_hal_rocm_context_wrapper_t *context, uint64_t initial_value, - iree_hal_semaphore_t **out_semaphore) { + iree_hal_rocm_context_wrapper_t* context, uint64_t initial_value, + iree_hal_semaphore_t** out_semaphore) { IREE_ASSERT_ARGUMENT(context); IREE_ASSERT_ARGUMENT(out_semaphore); IREE_TRACE_ZONE_BEGIN(z0); - iree_hal_rocm_semaphore_t *semaphore = NULL; + iree_hal_rocm_semaphore_t* semaphore = NULL; iree_status_t status = iree_allocator_malloc( - context->host_allocator, sizeof(*semaphore), (void **)&semaphore); + context->host_allocator, sizeof(*semaphore), (void**)&semaphore); if (iree_status_is_ok(status)) { iree_hal_resource_initialize(&iree_hal_rocm_semaphore_vtable, &semaphore->resource); semaphore->context = context; semaphore->initial_value = initial_value; - *out_semaphore = (iree_hal_semaphore_t *)semaphore; + *out_semaphore = (iree_hal_semaphore_t*)semaphore; } IREE_TRACE_ZONE_END(z0); @@ -48,8 +48,8 @@ } static void iree_hal_rocm_semaphore_destroy( - iree_hal_semaphore_t *base_semaphore) { - iree_hal_rocm_semaphore_t *semaphore = + iree_hal_semaphore_t* base_semaphore) { + iree_hal_rocm_semaphore_t* semaphore = iree_hal_rocm_semaphore_cast(base_semaphore); iree_allocator_t host_allocator = semaphore->context->host_allocator; IREE_TRACE_ZONE_BEGIN(z0); @@ -60,24 +60,24 @@ } static iree_status_t iree_hal_rocm_semaphore_query( - iree_hal_semaphore_t *base_semaphore, uint64_t *out_value) { + iree_hal_semaphore_t* base_semaphore, uint64_t* out_value) { // TODO: Support semaphores completely. *out_value = 0; return iree_make_status(IREE_STATUS_UNIMPLEMENTED, "Not impemented on rocm"); } static iree_status_t iree_hal_rocm_semaphore_signal( - iree_hal_semaphore_t *base_semaphore, uint64_t new_value) { + iree_hal_semaphore_t* base_semaphore, uint64_t new_value) { // TODO: Support semaphores completely. Return OK currently as everything is // synchronized for each submit to allow things to run. return iree_ok_status(); } -static void iree_hal_rocm_semaphore_fail(iree_hal_semaphore_t *base_semaphore, +static void iree_hal_rocm_semaphore_fail(iree_hal_semaphore_t* base_semaphore, iree_status_t status) {} static iree_status_t iree_hal_rocm_semaphore_wait( - iree_hal_semaphore_t *base_semaphore, uint64_t value, + iree_hal_semaphore_t* base_semaphore, uint64_t value, iree_timeout_t timeout) { // TODO: Support semaphores completely. Return OK currently as everything is // synchronized for each submit to allow things to run.
diff --git a/experimental/rocm/event_semaphore.h b/experimental/rocm/event_semaphore.h index 5f32492..9e79aa8 100644 --- a/experimental/rocm/event_semaphore.h +++ b/experimental/rocm/event_semaphore.h
@@ -20,8 +20,8 @@ // Create a rocm allocator. iree_status_t iree_hal_rocm_semaphore_create( - iree_hal_rocm_context_wrapper_t *context, uint64_t initial_value, - iree_hal_semaphore_t **out_semaphore); + iree_hal_rocm_context_wrapper_t* context, uint64_t initial_value, + iree_hal_semaphore_t** out_semaphore); #ifdef __cplusplus } // extern "C"
diff --git a/experimental/rocm/executable_layout.c b/experimental/rocm/executable_layout.c index 6960cd2..f293921 100644 --- a/experimental/rocm/executable_layout.c +++ b/experimental/rocm/executable_layout.c
@@ -8,30 +8,31 @@ #include <stddef.h> +#include "experimental/rocm/descriptor_set_layout.h" #include "iree/base/api.h" #include "iree/base/tracing.h" typedef struct iree_hal_rocm_executable_layout_t { iree_hal_resource_t resource; - iree_hal_rocm_context_wrapper_t *context; + iree_hal_rocm_context_wrapper_t* context; iree_host_size_t set_layout_count; - iree_hal_descriptor_set_layout_t *set_layouts[]; + iree_hal_descriptor_set_layout_t* set_layouts[]; } iree_hal_rocm_executable_layout_t; extern const iree_hal_executable_layout_vtable_t iree_hal_rocm_executable_layout_vtable; -static iree_hal_rocm_executable_layout_t *iree_hal_rocm_executable_layout_cast( - iree_hal_executable_layout_t *base_value) { +static iree_hal_rocm_executable_layout_t* iree_hal_rocm_executable_layout_cast( + iree_hal_executable_layout_t* base_value) { IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_rocm_executable_layout_vtable); - return (iree_hal_rocm_executable_layout_t *)base_value; + return (iree_hal_rocm_executable_layout_t*)base_value; } iree_status_t iree_hal_rocm_executable_layout_create( - iree_hal_rocm_context_wrapper_t *context, iree_host_size_t set_layout_count, - iree_hal_descriptor_set_layout_t **set_layouts, + iree_hal_rocm_context_wrapper_t* context, iree_host_size_t set_layout_count, + iree_hal_descriptor_set_layout_t** set_layouts, iree_host_size_t push_constant_count, - iree_hal_executable_layout_t **out_executable_layout) { + iree_hal_executable_layout_t** out_executable_layout) { IREE_ASSERT_ARGUMENT(context); IREE_ASSERT_ARGUMENT(!set_layout_count || set_layouts); IREE_ASSERT_ARGUMENT(out_executable_layout); @@ -40,12 +41,12 @@ // Currently the executable layout doesn't do anything. // TODO: Handle creating the argument layout at that time hadling both push // constant and buffers. - iree_hal_rocm_executable_layout_t *executable_layout = NULL; + iree_hal_rocm_executable_layout_t* executable_layout = NULL; iree_host_size_t total_size = sizeof(*executable_layout) + set_layout_count * sizeof(*executable_layout->set_layouts); iree_status_t status = iree_allocator_malloc( - context->host_allocator, total_size, (void **)&executable_layout); + context->host_allocator, total_size, (void**)&executable_layout); if (iree_status_is_ok(status)) { iree_hal_resource_initialize(&iree_hal_rocm_executable_layout_vtable, &executable_layout->resource); @@ -55,15 +56,15 @@ executable_layout->set_layouts[i] = set_layouts[i]; iree_hal_descriptor_set_layout_retain(set_layouts[i]); } - *out_executable_layout = (iree_hal_executable_layout_t *)executable_layout; + *out_executable_layout = (iree_hal_executable_layout_t*)executable_layout; } IREE_TRACE_ZONE_END(z0); return status; } static void iree_hal_rocm_executable_layout_destroy( - iree_hal_executable_layout_t *base_executable_layout) { - iree_hal_rocm_executable_layout_t *executable_layout = + iree_hal_executable_layout_t* base_executable_layout) { + iree_hal_rocm_executable_layout_t* executable_layout = iree_hal_rocm_executable_layout_cast(base_executable_layout); iree_allocator_t host_allocator = executable_layout->context->host_allocator; IREE_TRACE_ZONE_BEGIN(z0); @@ -76,6 +77,20 @@ IREE_TRACE_ZONE_END(z0); } +iree_host_size_t iree_hal_rocm_base_binding_index( + iree_hal_executable_layout_t* base_executable_layout, uint32_t set) { + iree_hal_rocm_executable_layout_t* executable_layout = + iree_hal_rocm_executable_layout_cast(base_executable_layout); + iree_host_size_t base_binding = 0; + for (iree_host_size_t i = 0; i < set; ++i) { + iree_host_size_t binding_count = + iree_hal_rocm_descriptor_set_layout_binding_count( + executable_layout->set_layouts[i]); + base_binding += binding_count; + } + return base_binding; +} + const iree_hal_executable_layout_vtable_t iree_hal_rocm_executable_layout_vtable = { .destroy = iree_hal_rocm_executable_layout_destroy,
diff --git a/experimental/rocm/executable_layout.h b/experimental/rocm/executable_layout.h index cd0fccf..8287cb8 100644 --- a/experimental/rocm/executable_layout.h +++ b/experimental/rocm/executable_layout.h
@@ -17,10 +17,14 @@ // Creates the kernel arguments. iree_status_t iree_hal_rocm_executable_layout_create( - iree_hal_rocm_context_wrapper_t *context, iree_host_size_t set_layout_count, - iree_hal_descriptor_set_layout_t **set_layouts, + iree_hal_rocm_context_wrapper_t* context, iree_host_size_t set_layout_count, + iree_hal_descriptor_set_layout_t** set_layouts, iree_host_size_t push_constant_count, - iree_hal_executable_layout_t **out_executable_layout); + iree_hal_executable_layout_t** out_executable_layout); + +// Return the base binding index for the given set. +iree_host_size_t iree_hal_rocm_base_binding_index( + iree_hal_executable_layout_t* executable_layout, uint32_t set); #ifdef __cplusplus } // extern "C"
diff --git a/experimental/rocm/native_executable.c b/experimental/rocm/native_executable.c index f72bd4f..680f07f 100644 --- a/experimental/rocm/native_executable.c +++ b/experimental/rocm/native_executable.c
@@ -27,7 +27,7 @@ typedef struct iree_hal_rocm_native_executable_t { iree_hal_resource_t resource; - iree_hal_rocm_context_wrapper_t *context; + iree_hal_rocm_context_wrapper_t* context; iree_host_size_t entry_count; hipModule_t module; iree_hal_rocm_native_executable_function_t entry_functions[]; @@ -36,23 +36,23 @@ extern const iree_hal_executable_vtable_t iree_hal_rocm_native_executable_vtable; -static iree_hal_rocm_native_executable_t *iree_hal_rocm_native_executable_cast( - iree_hal_executable_t *base_value) { +static iree_hal_rocm_native_executable_t* iree_hal_rocm_native_executable_cast( + iree_hal_executable_t* base_value) { IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_rocm_native_executable_vtable); - return (iree_hal_rocm_native_executable_t *)base_value; + return (iree_hal_rocm_native_executable_t*)base_value; } iree_status_t iree_hal_rocm_native_executable_create( - iree_hal_rocm_context_wrapper_t *context, - const iree_hal_executable_spec_t *executable_spec, - iree_hal_executable_t **out_executable) { + iree_hal_rocm_context_wrapper_t* context, + const iree_hal_executable_spec_t* executable_spec, + iree_hal_executable_t** out_executable) { IREE_ASSERT_ARGUMENT(context); IREE_ASSERT_ARGUMENT(executable_spec); IREE_ASSERT_ARGUMENT(out_executable); *out_executable = NULL; IREE_TRACE_ZONE_BEGIN(z0); - iree_hal_rocm_native_executable_t *executable = NULL; + iree_hal_rocm_native_executable_t* executable = NULL; // TODO: Verify the flat buffer. iree_ROCMExecutableDef_table_t executable_def = @@ -69,8 +69,8 @@ iree_host_size_t total_size = sizeof(*executable) + entry_count * sizeof(iree_hal_rocm_native_executable_function_t); - iree_status_t status = iree_allocator_malloc( - context->host_allocator, total_size, (void **)&executable); + iree_status_t status = iree_allocator_malloc(context->host_allocator, + total_size, (void**)&executable); hipModule_t module = NULL; ROCM_RETURN_IF_ERROR(context->syms, hipModuleLoadDataEx(&module, hsaco_image, 0, NULL, NULL), @@ -78,7 +78,7 @@ for (iree_host_size_t i = 0; i < entry_count; i++) { hipFunction_t function = NULL; - const char *entry_name = flatbuffers_string_vec_at(entry_points_vec, i); + const char* entry_name = flatbuffers_string_vec_at(entry_points_vec, i); ROCM_RETURN_IF_ERROR(context->syms, hipModuleGetFunction(&function, module, entry_name), "hipModuleGetFunction"); @@ -92,22 +92,22 @@ &executable->resource); executable->module = module; executable->context = context; - *out_executable = (iree_hal_executable_t *)executable; + *out_executable = (iree_hal_executable_t*)executable; IREE_TRACE_ZONE_END(z0); return iree_ok_status(); } hipFunction_t iree_hal_rocm_native_executable_for_entry_point( - iree_hal_executable_t *base_executable, int32_t entry_point) { - iree_hal_rocm_native_executable_t *executable = + iree_hal_executable_t* base_executable, int32_t entry_point) { + iree_hal_rocm_native_executable_t* executable = iree_hal_rocm_native_executable_cast(base_executable); return executable->entry_functions[entry_point].rocm_function; } iree_status_t iree_hal_rocm_native_executable_block_size( - iree_hal_executable_t *base_executable, int32_t entry_point, uint32_t *x, - uint32_t *y, uint32_t *z) { - iree_hal_rocm_native_executable_t *executable = + iree_hal_executable_t* base_executable, int32_t entry_point, uint32_t* x, + uint32_t* y, uint32_t* z) { + iree_hal_rocm_native_executable_t* executable = iree_hal_rocm_native_executable_cast(base_executable); *x = executable->entry_functions[entry_point].block_size_x; *y = executable->entry_functions[entry_point].block_size_y; @@ -116,8 +116,8 @@ } static void iree_hal_rocm_native_executable_destroy( - iree_hal_executable_t *base_executable) { - iree_hal_rocm_native_executable_t *executable = + iree_hal_executable_t* base_executable) { + iree_hal_rocm_native_executable_t* executable = iree_hal_rocm_native_executable_cast(base_executable); iree_allocator_t host_allocator = executable->context->host_allocator; IREE_TRACE_ZONE_BEGIN(z0);
diff --git a/experimental/rocm/native_executable.h b/experimental/rocm/native_executable.h index 7a9229a..bc671bf 100644 --- a/experimental/rocm/native_executable.h +++ b/experimental/rocm/native_executable.h
@@ -21,17 +21,17 @@ // Creates an executable from a HSACO module. The module may contain several // kernels that can be extracted along with the associated block size. iree_status_t iree_hal_rocm_native_executable_create( - iree_hal_rocm_context_wrapper_t *context, - const iree_hal_executable_spec_t *executable_spec, - iree_hal_executable_t **out_executable); + iree_hal_rocm_context_wrapper_t* context, + const iree_hal_executable_spec_t* executable_spec, + iree_hal_executable_t** out_executable); hipFunction_t iree_hal_rocm_native_executable_for_entry_point( - iree_hal_executable_t *executable, int32_t entry_point); + iree_hal_executable_t* executable, int32_t entry_point); // Return the block size of the given |entry_point| within the executable. iree_status_t iree_hal_rocm_native_executable_block_size( - iree_hal_executable_t *executable, int32_t entry_point, uint32_t *x, - uint32_t *y, uint32_t *z); + iree_hal_executable_t* executable, int32_t entry_point, uint32_t* x, + uint32_t* y, uint32_t* z); #ifdef __cplusplus } // extern "C"
diff --git a/experimental/rocm/nop_executable_cache.c b/experimental/rocm/nop_executable_cache.c index 6592727..54bde49 100644 --- a/experimental/rocm/nop_executable_cache.c +++ b/experimental/rocm/nop_executable_cache.c
@@ -15,44 +15,44 @@ typedef struct iree_hal_rocm_nop_executable_cache_t { iree_hal_resource_t resource; - iree_hal_rocm_context_wrapper_t *context; + iree_hal_rocm_context_wrapper_t* context; } iree_hal_rocm_nop_executable_cache_t; extern const iree_hal_executable_cache_vtable_t iree_hal_rocm_nop_executable_cache_vtable; -static iree_hal_rocm_nop_executable_cache_t * +static iree_hal_rocm_nop_executable_cache_t* iree_hal_rocm_nop_executable_cache_cast( - iree_hal_executable_cache_t *base_value) { + iree_hal_executable_cache_t* base_value) { IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_rocm_nop_executable_cache_vtable); - return (iree_hal_rocm_nop_executable_cache_t *)base_value; + return (iree_hal_rocm_nop_executable_cache_t*)base_value; } iree_status_t iree_hal_rocm_nop_executable_cache_create( - iree_hal_rocm_context_wrapper_t *context, iree_string_view_t identifier, - iree_hal_executable_cache_t **out_executable_cache) { + iree_hal_rocm_context_wrapper_t* context, iree_string_view_t identifier, + iree_hal_executable_cache_t** out_executable_cache) { IREE_ASSERT_ARGUMENT(out_executable_cache); *out_executable_cache = NULL; IREE_TRACE_ZONE_BEGIN(z0); - iree_hal_rocm_nop_executable_cache_t *executable_cache = NULL; + iree_hal_rocm_nop_executable_cache_t* executable_cache = NULL; iree_status_t status = iree_allocator_malloc(context->host_allocator, sizeof(*executable_cache), - (void **)&executable_cache); + (void**)&executable_cache); if (iree_status_is_ok(status)) { iree_hal_resource_initialize(&iree_hal_rocm_nop_executable_cache_vtable, &executable_cache->resource); executable_cache->context = context; - *out_executable_cache = (iree_hal_executable_cache_t *)executable_cache; + *out_executable_cache = (iree_hal_executable_cache_t*)executable_cache; } IREE_TRACE_ZONE_END(z0); return status; } static void iree_hal_rocm_nop_executable_cache_destroy( - iree_hal_executable_cache_t *base_executable_cache) { - iree_hal_rocm_nop_executable_cache_t *executable_cache = + iree_hal_executable_cache_t* base_executable_cache) { + iree_hal_rocm_nop_executable_cache_t* executable_cache = iree_hal_rocm_nop_executable_cache_cast(base_executable_cache); iree_allocator_t host_allocator = executable_cache->context->host_allocator; IREE_TRACE_ZONE_BEGIN(z0); @@ -63,7 +63,7 @@ } static bool iree_hal_rocm_nop_executable_cache_can_prepare_format( - iree_hal_executable_cache_t *base_executable_cache, + iree_hal_executable_cache_t* base_executable_cache, iree_hal_executable_caching_mode_t caching_mode, iree_string_view_t executable_format) { return iree_string_view_equal(executable_format, @@ -71,10 +71,10 @@ } static iree_status_t iree_hal_rocm_nop_executable_cache_prepare_executable( - iree_hal_executable_cache_t *base_executable_cache, - const iree_hal_executable_spec_t *executable_spec, - iree_hal_executable_t **out_executable) { - iree_hal_rocm_nop_executable_cache_t *executable_cache = + iree_hal_executable_cache_t* base_executable_cache, + const iree_hal_executable_spec_t* executable_spec, + iree_hal_executable_t** out_executable) { + iree_hal_rocm_nop_executable_cache_t* executable_cache = iree_hal_rocm_nop_executable_cache_cast(base_executable_cache); return iree_hal_rocm_native_executable_create( executable_cache->context, executable_spec, out_executable);
diff --git a/experimental/rocm/nop_executable_cache.h b/experimental/rocm/nop_executable_cache.h index d1b2fc1..a057466 100644 --- a/experimental/rocm/nop_executable_cache.h +++ b/experimental/rocm/nop_executable_cache.h
@@ -19,8 +19,8 @@ // This is useful to isolate pipeline caching behavior and verify compilation // behavior. iree_status_t iree_hal_rocm_nop_executable_cache_create( - iree_hal_rocm_context_wrapper_t *context, iree_string_view_t identifier, - iree_hal_executable_cache_t **out_executable_cache); + iree_hal_rocm_context_wrapper_t* context, iree_string_view_t identifier, + iree_hal_executable_cache_t** out_executable_cache); #ifdef __cplusplus } // extern "C"
diff --git a/experimental/rocm/rocm_allocator.c b/experimental/rocm/rocm_allocator.c index cccdc9f..4f9742f 100644 --- a/experimental/rocm/rocm_allocator.c +++ b/experimental/rocm/rocm_allocator.c
@@ -16,30 +16,30 @@ typedef struct iree_hal_rocm_allocator_t { iree_hal_resource_t resource; - iree_hal_rocm_context_wrapper_t *context; + iree_hal_rocm_context_wrapper_t* context; } iree_hal_rocm_allocator_t; extern const iree_hal_allocator_vtable_t iree_hal_rocm_allocator_vtable; -static iree_hal_rocm_allocator_t *iree_hal_rocm_allocator_cast( - iree_hal_allocator_t *base_value) { +static iree_hal_rocm_allocator_t* iree_hal_rocm_allocator_cast( + iree_hal_allocator_t* base_value) { IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_rocm_allocator_vtable); - return (iree_hal_rocm_allocator_t *)base_value; + return (iree_hal_rocm_allocator_t*)base_value; } iree_status_t iree_hal_rocm_allocator_create( - iree_hal_rocm_context_wrapper_t *context, - iree_hal_allocator_t **out_allocator) { + iree_hal_rocm_context_wrapper_t* context, + iree_hal_allocator_t** out_allocator) { IREE_ASSERT_ARGUMENT(context); IREE_TRACE_ZONE_BEGIN(z0); - iree_hal_rocm_allocator_t *allocator = NULL; + iree_hal_rocm_allocator_t* allocator = NULL; iree_status_t status = iree_allocator_malloc( - context->host_allocator, sizeof(*allocator), (void **)&allocator); + context->host_allocator, sizeof(*allocator), (void**)&allocator); if (iree_status_is_ok(status)) { iree_hal_resource_initialize(&iree_hal_rocm_allocator_vtable, &allocator->resource); allocator->context = context; - *out_allocator = (iree_hal_allocator_t *)allocator; + *out_allocator = (iree_hal_allocator_t*)allocator; } IREE_TRACE_ZONE_END(z0); @@ -47,8 +47,8 @@ } static void iree_hal_rocm_allocator_destroy( - iree_hal_allocator_t *base_allocator) { - iree_hal_rocm_allocator_t *allocator = + iree_hal_allocator_t* base_allocator) { + iree_hal_rocm_allocator_t* allocator = iree_hal_rocm_allocator_cast(base_allocator); iree_allocator_t host_allocator = allocator->context->host_allocator; IREE_TRACE_ZONE_BEGIN(z0); @@ -59,15 +59,15 @@ } static iree_allocator_t iree_hal_rocm_allocator_host_allocator( - const iree_hal_allocator_t *base_allocator) { - iree_hal_rocm_allocator_t *allocator = - (iree_hal_rocm_allocator_t *)base_allocator; + const iree_hal_allocator_t* base_allocator) { + iree_hal_rocm_allocator_t* allocator = + (iree_hal_rocm_allocator_t*)base_allocator; return allocator->context->host_allocator; } static iree_hal_buffer_compatibility_t iree_hal_rocm_allocator_query_buffer_compatibility( - iree_hal_allocator_t *base_allocator, iree_hal_memory_type_t memory_type, + iree_hal_allocator_t* base_allocator, iree_hal_memory_type_t memory_type, iree_hal_buffer_usage_t allowed_usage, iree_hal_buffer_usage_t intended_usage, iree_device_size_t allocation_size) { @@ -96,19 +96,31 @@ } static iree_status_t iree_hal_rocm_allocator_allocate_buffer( - iree_hal_allocator_t *base_allocator, iree_hal_memory_type_t memory_type, + iree_hal_allocator_t* base_allocator, iree_hal_memory_type_t memory_type, iree_hal_buffer_usage_t allowed_usage, iree_host_size_t allocation_size, - iree_hal_buffer_t **out_buffer) { - iree_hal_rocm_allocator_t *allocator = + iree_hal_buffer_t** out_buffer) { + iree_hal_rocm_allocator_t* allocator = iree_hal_rocm_allocator_cast(base_allocator); // Guard against the corner case where the requested buffer size is 0. The // application is unlikely to do anything when requesting a 0-byte buffer; but // it can happen in real world use cases. So we should at least not crash. if (allocation_size == 0) allocation_size = 4; iree_status_t status; - void *host_ptr = NULL; + void* host_ptr = NULL; hipDeviceptr_t device_ptr = 0; - if (iree_all_bits_set(memory_type, IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)) { + if (iree_all_bits_set(memory_type, IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL)) { + // Device local case. + if (iree_all_bits_set(memory_type, IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)) { + status = ROCM_RESULT_TO_STATUS( + allocator->context->syms, + hipMallocManaged(&device_ptr, allocation_size, hipMemAttachGlobal)); + host_ptr = (void*)device_ptr; + } else { + // Device only. + status = ROCM_RESULT_TO_STATUS(allocator->context->syms, + hipMalloc(&device_ptr, allocation_size)); + } + } else { unsigned int flags = hipHostMallocMapped; if (!iree_all_bits_set(memory_type, IREE_HAL_MEMORY_TYPE_HOST_CACHED)) { flags |= hipHostMallocWriteCombined; @@ -121,14 +133,11 @@ allocator->context->syms, hipHostGetDevicePointer(&device_ptr, host_ptr, /*flags=*/0)); } - } else { - status = ROCM_RESULT_TO_STATUS(allocator->context->syms, - hipMalloc(&device_ptr, allocation_size)); } if (iree_status_is_ok(status)) { status = iree_hal_rocm_buffer_wrap( - (iree_hal_allocator_t *)allocator, memory_type, + (iree_hal_allocator_t*)allocator, memory_type, IREE_HAL_MEMORY_ACCESS_ALL, allowed_usage, allocation_size, /*byte_offset=*/0, /*byte_length=*/allocation_size, device_ptr, host_ptr, out_buffer); @@ -140,23 +149,24 @@ return status; } -void iree_hal_rocm_allocator_free(iree_hal_allocator_t *base_allocator, - hipDeviceptr_t device_ptr, void *host_ptr, +void iree_hal_rocm_allocator_free(iree_hal_allocator_t* base_allocator, + hipDeviceptr_t device_ptr, void* host_ptr, iree_hal_memory_type_t memory_type) { - iree_hal_rocm_allocator_t *allocator = + iree_hal_rocm_allocator_t* allocator = iree_hal_rocm_allocator_cast(base_allocator); - if (iree_all_bits_set(memory_type, IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)) { - ROCM_IGNORE_ERROR(allocator->context->syms, hipHostFree(host_ptr)); - } else { + if (iree_all_bits_set(memory_type, IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL)) { ROCM_IGNORE_ERROR(allocator->context->syms, hipFree(device_ptr)); + } else { + // Host local. + ROCM_IGNORE_ERROR(allocator->context->syms, hipHostFree(host_ptr)); } } static iree_status_t iree_hal_rocm_allocator_wrap_buffer( - iree_hal_allocator_t *base_allocator, iree_hal_memory_type_t memory_type, + iree_hal_allocator_t* base_allocator, iree_hal_memory_type_t memory_type, iree_hal_memory_access_t allowed_access, iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, - iree_allocator_t data_allocator, iree_hal_buffer_t **out_buffer) { + iree_allocator_t data_allocator, iree_hal_buffer_t** out_buffer) { return iree_make_status(IREE_STATUS_UNAVAILABLE, "wrapping of external buffers not supported"); }
diff --git a/experimental/rocm/rocm_allocator.h b/experimental/rocm/rocm_allocator.h index a47480c..37d6710 100644 --- a/experimental/rocm/rocm_allocator.h +++ b/experimental/rocm/rocm_allocator.h
@@ -18,12 +18,12 @@ // Create a ROCM allocator. iree_status_t iree_hal_rocm_allocator_create( - iree_hal_rocm_context_wrapper_t *context, - iree_hal_allocator_t **out_allocator); + iree_hal_rocm_context_wrapper_t* context, + iree_hal_allocator_t** out_allocator); // Free an allocation represent by the given device or host pointer. -void iree_hal_rocm_allocator_free(iree_hal_allocator_t *allocator, - hipDeviceptr_t device_ptr, void *host_ptr, +void iree_hal_rocm_allocator_free(iree_hal_allocator_t* allocator, + hipDeviceptr_t device_ptr, void* host_ptr, iree_hal_memory_type_t memory_type); #ifdef __cplusplus
diff --git a/experimental/rocm/rocm_buffer.c b/experimental/rocm/rocm_buffer.c index 195e897..8822cb5 100644 --- a/experimental/rocm/rocm_buffer.c +++ b/experimental/rocm/rocm_buffer.c
@@ -16,32 +16,32 @@ typedef struct iree_hal_rocm_buffer_t { iree_hal_buffer_t base; - void *host_ptr; + void* host_ptr; hipDeviceptr_t device_ptr; } iree_hal_rocm_buffer_t; extern const iree_hal_buffer_vtable_t iree_hal_rocm_buffer_vtable; -static iree_hal_rocm_buffer_t *iree_hal_rocm_buffer_cast( - iree_hal_buffer_t *base_value) { +static iree_hal_rocm_buffer_t* iree_hal_rocm_buffer_cast( + iree_hal_buffer_t* base_value) { IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_rocm_buffer_vtable); - return (iree_hal_rocm_buffer_t *)base_value; + return (iree_hal_rocm_buffer_t*)base_value; } iree_status_t iree_hal_rocm_buffer_wrap( - iree_hal_allocator_t *allocator, iree_hal_memory_type_t memory_type, + iree_hal_allocator_t* allocator, iree_hal_memory_type_t memory_type, iree_hal_memory_access_t allowed_access, iree_hal_buffer_usage_t allowed_usage, iree_device_size_t allocation_size, iree_device_size_t byte_offset, iree_device_size_t byte_length, - hipDeviceptr_t device_ptr, void *host_ptr, iree_hal_buffer_t **out_buffer) { + hipDeviceptr_t device_ptr, void* host_ptr, iree_hal_buffer_t** out_buffer) { IREE_ASSERT_ARGUMENT(allocator); IREE_ASSERT_ARGUMENT(out_buffer); IREE_TRACE_ZONE_BEGIN(z0); - iree_hal_rocm_buffer_t *buffer = NULL; + iree_hal_rocm_buffer_t* buffer = NULL; iree_status_t status = iree_allocator_malloc(iree_hal_allocator_host_allocator(allocator), - sizeof(*buffer), (void **)&buffer); + sizeof(*buffer), (void**)&buffer); if (iree_status_is_ok(status)) { iree_hal_resource_initialize(&iree_hal_rocm_buffer_vtable, &buffer->base.resource); @@ -59,11 +59,11 @@ } IREE_TRACE_ZONE_END(z0); - return iree_ok_status(); + return status; } -static void iree_hal_rocm_buffer_destroy(iree_hal_buffer_t *base_buffer) { - iree_hal_rocm_buffer_t *buffer = iree_hal_rocm_buffer_cast(base_buffer); +static void iree_hal_rocm_buffer_destroy(iree_hal_buffer_t* base_buffer) { + iree_hal_rocm_buffer_t* buffer = iree_hal_rocm_buffer_cast(base_buffer); iree_allocator_t host_allocator = iree_hal_allocator_host_allocator(iree_hal_buffer_allocator(base_buffer)); IREE_TRACE_ZONE_BEGIN(z0); @@ -76,11 +76,11 @@ } static iree_status_t iree_hal_rocm_buffer_map_range( - iree_hal_buffer_t *base_buffer, iree_hal_mapping_mode_t mapping_mode, + iree_hal_buffer_t* base_buffer, iree_hal_mapping_mode_t mapping_mode, iree_hal_memory_access_t memory_access, iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length, - void **out_data_ptr) { - iree_hal_rocm_buffer_t *buffer = iree_hal_rocm_buffer_cast(base_buffer); + void** out_data_ptr) { + iree_hal_rocm_buffer_t* buffer = iree_hal_rocm_buffer_cast(base_buffer); if (!iree_all_bits_set(buffer->base.memory_type, IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)) { @@ -88,14 +88,14 @@ "trying to map memory not host visible"); } - uint8_t *data_ptr = (uint8_t *)(buffer->host_ptr) + local_byte_offset; + uint8_t* data_ptr = (uint8_t*)(buffer->host_ptr) + local_byte_offset; // If we mapped for discard scribble over the bytes. This is not a mandated // behavior but it will make debugging issues easier. Alternatively for // heap buffers we could reallocate them such that ASAN yells, but that // would only work if the entire buffer was discarded. #ifndef NDEBUG if (iree_any_bit_set(memory_access, IREE_HAL_MEMORY_ACCESS_DISCARD)) { - memset(data_ptr + local_byte_offset, 0xCD, local_byte_length); + memset(data_ptr, 0xCD, local_byte_length); } #endif // !NDEBUG *out_data_ptr = data_ptr; @@ -103,28 +103,28 @@ } static void iree_hal_rocm_buffer_unmap_range( - iree_hal_buffer_t *base_buffer, iree_device_size_t local_byte_offset, - iree_device_size_t local_byte_length, void *data_ptr) { + iree_hal_buffer_t* base_buffer, iree_device_size_t local_byte_offset, + iree_device_size_t local_byte_length, void* data_ptr) { // nothing to do. } static iree_status_t iree_hal_rocm_buffer_invalidate_range( - iree_hal_buffer_t *base_buffer, iree_device_size_t local_byte_offset, + iree_hal_buffer_t* base_buffer, iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length) { // Nothing to do. return iree_ok_status(); } static iree_status_t iree_hal_rocm_buffer_flush_range( - iree_hal_buffer_t *base_buffer, iree_device_size_t local_byte_offset, + iree_hal_buffer_t* base_buffer, iree_device_size_t local_byte_offset, iree_device_size_t local_byte_length) { // Nothing to do. return iree_ok_status(); } hipDeviceptr_t iree_hal_rocm_buffer_device_pointer( - iree_hal_buffer_t *base_buffer) { - iree_hal_rocm_buffer_t *buffer = iree_hal_rocm_buffer_cast(base_buffer); + iree_hal_buffer_t* base_buffer) { + iree_hal_rocm_buffer_t* buffer = iree_hal_rocm_buffer_cast(base_buffer); return buffer->device_ptr; }
diff --git a/experimental/rocm/rocm_buffer.h b/experimental/rocm/rocm_buffer.h index e898630..85cf294 100644 --- a/experimental/rocm/rocm_buffer.h +++ b/experimental/rocm/rocm_buffer.h
@@ -17,16 +17,16 @@ // Wraps a rocm allocation in an iree_hal_buffer_t. iree_status_t iree_hal_rocm_buffer_wrap( - iree_hal_allocator_t *allocator, iree_hal_memory_type_t memory_type, + iree_hal_allocator_t* allocator, iree_hal_memory_type_t memory_type, iree_hal_memory_access_t allowed_access, iree_hal_buffer_usage_t allowed_usage, iree_device_size_t allocation_size, iree_device_size_t byte_offset, iree_device_size_t byte_length, - hipDeviceptr_t device_ptr, void *host_ptr, iree_hal_buffer_t **out_buffer); + hipDeviceptr_t device_ptr, void* host_ptr, iree_hal_buffer_t** out_buffer); // Returns the rocm base pointer for the given |buffer|. // This is the entire allocated_buffer and must be offset by the buffer // byte_offset and byte_length when used. -hipDeviceptr_t iree_hal_rocm_buffer_device_pointer(iree_hal_buffer_t *buffer); +hipDeviceptr_t iree_hal_rocm_buffer_device_pointer(iree_hal_buffer_t* buffer); #ifdef __cplusplus } // extern "C"
diff --git a/experimental/rocm/rocm_device.c b/experimental/rocm/rocm_device.c index cff4d89..fa5a75e 100644 --- a/experimental/rocm/rocm_device.c +++ b/experimental/rocm/rocm_device.c
@@ -32,27 +32,27 @@ // Optional driver that owns the ROCM symbols. We retain it for our lifetime // to ensure the symbols remains valid. - iree_hal_driver_t *driver; + iree_hal_driver_t* driver; hipDevice_t device; // TODO: support multiple streams. hipStream_t stream; iree_hal_rocm_context_wrapper_t context_wrapper; - iree_hal_allocator_t *device_allocator; + iree_hal_allocator_t* device_allocator; } iree_hal_rocm_device_t; extern const iree_hal_device_vtable_t iree_hal_rocm_device_vtable; -static iree_hal_rocm_device_t *iree_hal_rocm_device_cast( - iree_hal_device_t *base_value) { +static iree_hal_rocm_device_t* iree_hal_rocm_device_cast( + iree_hal_device_t* base_value) { IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_rocm_device_vtable); - return (iree_hal_rocm_device_t *)base_value; + return (iree_hal_rocm_device_t*)base_value; } -static void iree_hal_rocm_device_destroy(iree_hal_device_t *base_device) { - iree_hal_rocm_device_t *device = iree_hal_rocm_device_cast(base_device); +static void iree_hal_rocm_device_destroy(iree_hal_device_t* base_device) { + iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); iree_allocator_t host_allocator = iree_hal_device_host_allocator(base_device); IREE_TRACE_ZONE_BEGIN(z0); @@ -70,21 +70,21 @@ } static iree_status_t iree_hal_rocm_device_create_internal( - iree_hal_driver_t *driver, iree_string_view_t identifier, + iree_hal_driver_t* driver, iree_string_view_t identifier, hipDevice_t rocm_device, hipStream_t stream, hipCtx_t context, - iree_hal_rocm_dynamic_symbols_t *syms, iree_allocator_t host_allocator, - iree_hal_device_t **out_device) { - iree_hal_rocm_device_t *device = NULL; + iree_hal_rocm_dynamic_symbols_t* syms, iree_allocator_t host_allocator, + iree_hal_device_t** out_device) { + iree_hal_rocm_device_t* device = NULL; iree_host_size_t total_size = sizeof(*device) + identifier.size; IREE_RETURN_IF_ERROR( - iree_allocator_malloc(host_allocator, total_size, (void **)&device)); + iree_allocator_malloc(host_allocator, total_size, (void**)&device)); memset(device, 0, total_size); iree_hal_resource_initialize(&iree_hal_rocm_device_vtable, &device->resource); device->driver = driver; iree_hal_driver_retain(device->driver); - uint8_t *buffer_ptr = (uint8_t *)device + sizeof(*device); + uint8_t* buffer_ptr = (uint8_t*)device + sizeof(*device); buffer_ptr += iree_string_view_append_to_buffer( - identifier, &device->identifier, (char *)buffer_ptr); + identifier, &device->identifier, (char*)buffer_ptr); device->device = rocm_device; device->stream = stream; device->context_wrapper.rocm_context = context; @@ -93,19 +93,19 @@ iree_status_t status = iree_hal_rocm_allocator_create( &device->context_wrapper, &device->device_allocator); if (iree_status_is_ok(status)) { - *out_device = (iree_hal_device_t *)device; + *out_device = (iree_hal_device_t*)device; } else { - iree_hal_device_release((iree_hal_device_t *)device); + iree_hal_device_release((iree_hal_device_t*)device); } return status; } -iree_status_t iree_hal_rocm_device_create(iree_hal_driver_t *driver, +iree_status_t iree_hal_rocm_device_create(iree_hal_driver_t* driver, iree_string_view_t identifier, - iree_hal_rocm_dynamic_symbols_t *syms, + iree_hal_rocm_dynamic_symbols_t* syms, hipDevice_t device, iree_allocator_t host_allocator, - iree_hal_device_t **out_device) { + iree_hal_device_t** out_device) { IREE_TRACE_ZONE_BEGIN(z0); hipCtx_t context; IREE_RETURN_AND_END_ZONE_IF_ERROR( @@ -130,26 +130,26 @@ } static iree_string_view_t iree_hal_rocm_device_id( - iree_hal_device_t *base_device) { - iree_hal_rocm_device_t *device = iree_hal_rocm_device_cast(base_device); + iree_hal_device_t* base_device) { + iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); return device->identifier; } static iree_allocator_t iree_hal_rocm_device_host_allocator( - iree_hal_device_t *base_device) { - iree_hal_rocm_device_t *device = iree_hal_rocm_device_cast(base_device); + iree_hal_device_t* base_device) { + iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); return device->context_wrapper.host_allocator; } -static iree_hal_allocator_t *iree_hal_rocm_device_allocator( - iree_hal_device_t *base_device) { - iree_hal_rocm_device_t *device = iree_hal_rocm_device_cast(base_device); +static iree_hal_allocator_t* iree_hal_rocm_device_allocator( + iree_hal_device_t* base_device) { + iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); return device->device_allocator; } static iree_status_t iree_hal_rocm_device_query_i32( - iree_hal_device_t *base_device, iree_string_view_t category, - iree_string_view_t key, int32_t *out_value) { + iree_hal_device_t* base_device, iree_string_view_t category, + iree_string_view_t key, int32_t* out_value) { // iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); *out_value = 0; @@ -169,77 +169,77 @@ } static iree_status_t iree_hal_rocm_device_create_command_buffer( - iree_hal_device_t *base_device, iree_hal_command_buffer_mode_t mode, + iree_hal_device_t* base_device, iree_hal_command_buffer_mode_t mode, iree_hal_command_category_t command_categories, iree_hal_queue_affinity_t queue_affinity, - iree_hal_command_buffer_t **out_command_buffer) { - iree_hal_rocm_device_t *device = iree_hal_rocm_device_cast(base_device); - return iree_hal_rocm_direct_command_buffer_allocate( + iree_hal_command_buffer_t** out_command_buffer) { + iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); + return iree_hal_rocm_direct_command_buffer_create( &device->context_wrapper, mode, command_categories, queue_affinity, out_command_buffer); } static iree_status_t iree_hal_rocm_device_create_descriptor_set( - iree_hal_device_t *base_device, - iree_hal_descriptor_set_layout_t *set_layout, + iree_hal_device_t* base_device, + iree_hal_descriptor_set_layout_t* set_layout, iree_host_size_t binding_count, - const iree_hal_descriptor_set_binding_t *bindings, - iree_hal_descriptor_set_t **out_descriptor_set) { + const iree_hal_descriptor_set_binding_t* bindings, + iree_hal_descriptor_set_t** out_descriptor_set) { return iree_make_status(IREE_STATUS_UNIMPLEMENTED, "non-push descriptor sets still need work"); } static iree_status_t iree_hal_rocm_device_create_descriptor_set_layout( - iree_hal_device_t *base_device, + iree_hal_device_t* base_device, iree_hal_descriptor_set_layout_usage_type_t usage_type, iree_host_size_t binding_count, - const iree_hal_descriptor_set_layout_binding_t *bindings, - iree_hal_descriptor_set_layout_t **out_descriptor_set_layout) { - iree_hal_rocm_device_t *device = iree_hal_rocm_device_cast(base_device); + const iree_hal_descriptor_set_layout_binding_t* bindings, + iree_hal_descriptor_set_layout_t** out_descriptor_set_layout) { + iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); return iree_hal_rocm_descriptor_set_layout_create( &device->context_wrapper, usage_type, binding_count, bindings, out_descriptor_set_layout); } static iree_status_t iree_hal_rocm_device_create_event( - iree_hal_device_t *base_device, iree_hal_event_t **out_event) { - iree_hal_rocm_device_t *device = iree_hal_rocm_device_cast(base_device); + iree_hal_device_t* base_device, iree_hal_event_t** out_event) { + iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); return iree_hal_rocm_event_create(&device->context_wrapper, out_event); } static iree_status_t iree_hal_rocm_device_create_executable_cache( - iree_hal_device_t *base_device, iree_string_view_t identifier, - iree_hal_executable_cache_t **out_executable_cache) { - iree_hal_rocm_device_t *device = iree_hal_rocm_device_cast(base_device); + iree_hal_device_t* base_device, iree_string_view_t identifier, + iree_hal_executable_cache_t** out_executable_cache) { + iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); return iree_hal_rocm_nop_executable_cache_create( &device->context_wrapper, identifier, out_executable_cache); } static iree_status_t iree_hal_rocm_device_create_executable_layout( - iree_hal_device_t *base_device, iree_host_size_t push_constants, + iree_hal_device_t* base_device, iree_host_size_t push_constants, iree_host_size_t set_layout_count, - iree_hal_descriptor_set_layout_t **set_layouts, - iree_hal_executable_layout_t **out_executable_layout) { - iree_hal_rocm_device_t *device = iree_hal_rocm_device_cast(base_device); + iree_hal_descriptor_set_layout_t** set_layouts, + iree_hal_executable_layout_t** out_executable_layout) { + iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); return iree_hal_rocm_executable_layout_create( &device->context_wrapper, set_layout_count, set_layouts, push_constants, out_executable_layout); } static iree_status_t iree_hal_rocm_device_create_semaphore( - iree_hal_device_t *base_device, uint64_t initial_value, - iree_hal_semaphore_t **out_semaphore) { - iree_hal_rocm_device_t *device = iree_hal_rocm_device_cast(base_device); + iree_hal_device_t* base_device, uint64_t initial_value, + iree_hal_semaphore_t** out_semaphore) { + iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); return iree_hal_rocm_semaphore_create(&device->context_wrapper, initial_value, out_semaphore); } static iree_status_t iree_hal_rocm_device_queue_submit( - iree_hal_device_t *base_device, + iree_hal_device_t* base_device, iree_hal_command_category_t command_categories, iree_hal_queue_affinity_t queue_affinity, iree_host_size_t batch_count, - const iree_hal_submission_batch_t *batches) { - iree_hal_rocm_device_t *device = iree_hal_rocm_device_cast(base_device); + const iree_hal_submission_batch_t* batches) { + iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); // TODO(raikonenfnu): Once semaphore is implemented wait for semaphores // TODO(thomasraoux): Conservatively syncronize after every submit until we // support semaphores. @@ -251,11 +251,11 @@ } static iree_status_t iree_hal_rocm_device_submit_and_wait( - iree_hal_device_t *base_device, + iree_hal_device_t* base_device, iree_hal_command_category_t command_categories, iree_hal_queue_affinity_t queue_affinity, iree_host_size_t batch_count, - const iree_hal_submission_batch_t *batches, - iree_hal_semaphore_t *wait_semaphore, uint64_t wait_value, + const iree_hal_submission_batch_t* batches, + iree_hal_semaphore_t* wait_semaphore, uint64_t wait_value, iree_timeout_t timeout) { // Submit... IREE_RETURN_IF_ERROR(iree_hal_rocm_device_queue_submit( @@ -266,15 +266,15 @@ } static iree_status_t iree_hal_rocm_device_wait_semaphores( - iree_hal_device_t *base_device, iree_hal_wait_mode_t wait_mode, - const iree_hal_semaphore_list_t *semaphore_list, iree_timeout_t timeout) { + iree_hal_device_t* base_device, iree_hal_wait_mode_t wait_mode, + const iree_hal_semaphore_list_t* semaphore_list, iree_timeout_t timeout) { return iree_make_status(IREE_STATUS_UNIMPLEMENTED, "semaphore not implemented"); } static iree_status_t iree_hal_rocm_device_wait_idle( - iree_hal_device_t *base_device, iree_timeout_t timeout) { - iree_hal_rocm_device_t *device = iree_hal_rocm_device_cast(base_device); + iree_hal_device_t* base_device, iree_timeout_t timeout) { + iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device); // Wait until the stream is done. // TODO(thomasraoux): HIP doesn't support a deadline for wait, figure out how // to handle it better.
diff --git a/experimental/rocm/rocm_device.h b/experimental/rocm/rocm_device.h index 67e2f9b..083f4c7 100644 --- a/experimental/rocm/rocm_device.h +++ b/experimental/rocm/rocm_device.h
@@ -17,12 +17,12 @@ #endif // __cplusplus // Creates a device that owns and manages its own hipContext. -iree_status_t iree_hal_rocm_device_create(iree_hal_driver_t *driver, +iree_status_t iree_hal_rocm_device_create(iree_hal_driver_t* driver, iree_string_view_t identifier, - iree_hal_rocm_dynamic_symbols_t *syms, + iree_hal_rocm_dynamic_symbols_t* syms, hipDevice_t device, iree_allocator_t host_allocator, - iree_hal_device_t **out_device); + iree_hal_device_t** out_device); #ifdef __cplusplus } // extern "C"
diff --git a/experimental/rocm/rocm_driver.c b/experimental/rocm/rocm_driver.c index 8bd8ae6..808642f 100644 --- a/experimental/rocm/rocm_driver.c +++ b/experimental/rocm/rocm_driver.c
@@ -32,44 +32,44 @@ extern const iree_hal_driver_vtable_t iree_hal_rocm_driver_vtable; -static iree_hal_rocm_driver_t *iree_hal_rocm_driver_cast( - iree_hal_driver_t *base_value) { +static iree_hal_rocm_driver_t* iree_hal_rocm_driver_cast( + iree_hal_driver_t* base_value) { IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_rocm_driver_vtable); - return (iree_hal_rocm_driver_t *)base_value; + return (iree_hal_rocm_driver_t*)base_value; } IREE_API_EXPORT void iree_hal_rocm_driver_options_initialize( - iree_hal_rocm_driver_options_t *out_options) { + iree_hal_rocm_driver_options_t* out_options) { memset(out_options, 0, sizeof(*out_options)); out_options->default_device_index = 0; } static iree_status_t iree_hal_rocm_driver_create_internal( iree_string_view_t identifier, - const iree_hal_rocm_driver_options_t *options, - iree_allocator_t host_allocator, iree_hal_driver_t **out_driver) { - iree_hal_rocm_driver_t *driver = NULL; + const iree_hal_rocm_driver_options_t* options, + iree_allocator_t host_allocator, iree_hal_driver_t** out_driver) { + iree_hal_rocm_driver_t* driver = NULL; iree_host_size_t total_size = sizeof(*driver) + identifier.size; IREE_RETURN_IF_ERROR( - iree_allocator_malloc(host_allocator, total_size, (void **)&driver)); + iree_allocator_malloc(host_allocator, total_size, (void**)&driver)); iree_hal_resource_initialize(&iree_hal_rocm_driver_vtable, &driver->resource); driver->host_allocator = host_allocator; iree_string_view_append_to_buffer( identifier, &driver->identifier, - (char *)driver + total_size - identifier.size); + (char*)driver + total_size - identifier.size); driver->default_device_index = options->default_device_index; iree_status_t status = iree_hal_rocm_dynamic_symbols_initialize(host_allocator, &driver->syms); if (iree_status_is_ok(status)) { - *out_driver = (iree_hal_driver_t *)driver; + *out_driver = (iree_hal_driver_t*)driver; } else { - iree_hal_driver_release((iree_hal_driver_t *)driver); + iree_hal_driver_release((iree_hal_driver_t*)driver); } return status; } -static void iree_hal_rocm_driver_destroy(iree_hal_driver_t *base_driver) { - iree_hal_rocm_driver_t *driver = iree_hal_rocm_driver_cast(base_driver); +static void iree_hal_rocm_driver_destroy(iree_hal_driver_t* base_driver) { + iree_hal_rocm_driver_t* driver = iree_hal_rocm_driver_cast(base_driver); iree_allocator_t host_allocator = driver->host_allocator; IREE_TRACE_ZONE_BEGIN(z0); @@ -81,8 +81,8 @@ IREE_API_EXPORT iree_status_t iree_hal_rocm_driver_create( iree_string_view_t identifier, - const iree_hal_rocm_driver_options_t *options, - iree_allocator_t host_allocator, iree_hal_driver_t **out_driver) { + const iree_hal_rocm_driver_options_t* options, + iree_allocator_t host_allocator, iree_hal_driver_t** out_driver) { IREE_ASSERT_ARGUMENT(options); IREE_ASSERT_ARGUMENT(out_driver); IREE_TRACE_ZONE_BEGIN(z0); @@ -97,9 +97,9 @@ // Populates device information from the given ROCM physical device handle. // |out_device_info| must point to valid memory and additional data will be // appended to |buffer_ptr| and the new pointer is returned. -static uint8_t *iree_hal_rocm_populate_device_info( - hipDevice_t device, iree_hal_rocm_dynamic_symbols_t *syms, - uint8_t *buffer_ptr, iree_hal_device_info_t *out_device_info) { +static uint8_t* iree_hal_rocm_populate_device_info( + hipDevice_t device, iree_hal_rocm_dynamic_symbols_t* syms, + uint8_t* buffer_ptr, iree_hal_device_info_t* out_device_info) { char device_name[IREE_MAX_ROCM_DEVICE_NAME_LENGTH]; ROCM_IGNORE_ERROR(syms, hipDeviceGetName(device_name, sizeof(device_name), device)); @@ -109,31 +109,31 @@ iree_string_view_t device_name_string = iree_make_string_view(device_name, strlen(device_name)); buffer_ptr += iree_string_view_append_to_buffer( - device_name_string, &out_device_info->name, (char *)buffer_ptr); + device_name_string, &out_device_info->name, (char*)buffer_ptr); return buffer_ptr; } static iree_status_t iree_hal_rocm_driver_query_available_devices( - iree_hal_driver_t *base_driver, iree_allocator_t host_allocator, - iree_hal_device_info_t **out_device_infos, - iree_host_size_t *out_device_info_count) { - iree_hal_rocm_driver_t *driver = iree_hal_rocm_driver_cast(base_driver); + iree_hal_driver_t* base_driver, iree_allocator_t host_allocator, + iree_hal_device_info_t** out_device_infos, + iree_host_size_t* out_device_info_count) { + iree_hal_rocm_driver_t* driver = iree_hal_rocm_driver_cast(base_driver); // Query the number of available ROCM devices. int device_count = 0; ROCM_RETURN_IF_ERROR(&driver->syms, hipGetDeviceCount(&device_count), "hipGetDeviceCount"); // Allocate the return infos and populate with the devices. - iree_hal_device_info_t *device_infos = NULL; + iree_hal_device_info_t* device_infos = NULL; iree_host_size_t total_size = device_count * sizeof(iree_hal_device_info_t); for (iree_host_size_t i = 0; i < device_count; ++i) { total_size += IREE_MAX_ROCM_DEVICE_NAME_LENGTH * sizeof(char); } iree_status_t status = - iree_allocator_malloc(host_allocator, total_size, (void **)&device_infos); + iree_allocator_malloc(host_allocator, total_size, (void**)&device_infos); if (iree_status_is_ok(status)) { - uint8_t *buffer_ptr = - (uint8_t *)device_infos + device_count * sizeof(iree_hal_device_info_t); + uint8_t* buffer_ptr = + (uint8_t*)device_infos + device_count * sizeof(iree_hal_device_info_t); for (iree_host_size_t i = 0; i < device_count; ++i) { hipDevice_t device; iree_status_t status = ROCM_RESULT_TO_STATUS( @@ -153,8 +153,8 @@ } static iree_status_t iree_hal_rocm_driver_select_default_device( - iree_hal_rocm_dynamic_symbols_t *syms, int default_device_index, - iree_allocator_t host_allocator, hipDevice_t *out_device) { + iree_hal_rocm_dynamic_symbols_t* syms, int default_device_index, + iree_allocator_t host_allocator, hipDevice_t* out_device) { int device_count = 0; ROCM_RETURN_IF_ERROR(syms, hipGetDeviceCount(&device_count), "hipGetDeviceCount"); @@ -173,9 +173,9 @@ } static iree_status_t iree_hal_rocm_driver_create_device( - iree_hal_driver_t *base_driver, iree_hal_device_id_t device_id, - iree_allocator_t host_allocator, iree_hal_device_t **out_device) { - iree_hal_rocm_driver_t *driver = iree_hal_rocm_driver_cast(base_driver); + iree_hal_driver_t* base_driver, iree_hal_device_id_t device_id, + iree_allocator_t host_allocator, iree_hal_device_t** out_device) { + iree_hal_rocm_driver_t* driver = iree_hal_rocm_driver_cast(base_driver); IREE_TRACE_ZONE_BEGIN(z0); IREE_RETURN_AND_END_ZONE_IF_ERROR(
diff --git a/experimental/rocm/rocm_event.c b/experimental/rocm/rocm_event.c index 48b20c3..945a5f7 100644 --- a/experimental/rocm/rocm_event.c +++ b/experimental/rocm/rocm_event.c
@@ -14,40 +14,40 @@ // Dummy events for now, don't do anything. typedef struct iree_hal_rocm_event_t { iree_hal_resource_t resource; - iree_hal_rocm_context_wrapper_t *context_wrapper; + iree_hal_rocm_context_wrapper_t* context_wrapper; } iree_hal_rocm_event_t; extern const iree_hal_event_vtable_t iree_hal_rocm_event_vtable; -static iree_hal_rocm_event_t *iree_hal_rocm_event_cast( - iree_hal_event_t *base_value) { +static iree_hal_rocm_event_t* iree_hal_rocm_event_cast( + iree_hal_event_t* base_value) { IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_rocm_event_vtable); - return (iree_hal_rocm_event_t *)base_value; + return (iree_hal_rocm_event_t*)base_value; } iree_status_t iree_hal_rocm_event_create( - iree_hal_rocm_context_wrapper_t *context_wrapper, - iree_hal_event_t **out_event) { + iree_hal_rocm_context_wrapper_t* context_wrapper, + iree_hal_event_t** out_event) { IREE_ASSERT_ARGUMENT(context_wrapper); IREE_ASSERT_ARGUMENT(out_event); *out_event = NULL; IREE_TRACE_ZONE_BEGIN(z0); - iree_hal_rocm_event_t *event = NULL; + iree_hal_rocm_event_t* event = NULL; iree_status_t status = iree_allocator_malloc(context_wrapper->host_allocator, - sizeof(*event), (void **)&event); + sizeof(*event), (void**)&event); if (iree_status_is_ok(status)) { iree_hal_resource_initialize(&iree_hal_rocm_event_vtable, &event->resource); event->context_wrapper = context_wrapper; - *out_event = (iree_hal_event_t *)event; + *out_event = (iree_hal_event_t*)event; } IREE_TRACE_ZONE_END(z0); return status; } -static void iree_hal_rocm_event_destroy(iree_hal_event_t *base_event) { - iree_hal_rocm_event_t *event = iree_hal_rocm_event_cast(base_event); +static void iree_hal_rocm_event_destroy(iree_hal_event_t* base_event) { + iree_hal_rocm_event_t* event = iree_hal_rocm_event_cast(base_event); iree_allocator_t host_allocator = event->context_wrapper->host_allocator; IREE_TRACE_ZONE_BEGIN(z0);
diff --git a/experimental/rocm/rocm_event.h b/experimental/rocm/rocm_event.h index 73041a4..0bac1a2 100644 --- a/experimental/rocm/rocm_event.h +++ b/experimental/rocm/rocm_event.h
@@ -21,8 +21,8 @@ // command buffer we will add the appropriate edges to enforce the right // synchronization. iree_status_t iree_hal_rocm_event_create( - iree_hal_rocm_context_wrapper_t *context_wrapper, - iree_hal_event_t **out_event); + iree_hal_rocm_context_wrapper_t* context_wrapper, + iree_hal_event_t** out_event); #ifdef __cplusplus } // extern "C"
diff --git a/experimental/rocm/status_util.h b/experimental/rocm/status_util.h index b459591..0f6fcc5 100644 --- a/experimental/rocm/status_util.h +++ b/experimental/rocm/status_util.h
@@ -44,7 +44,7 @@ // Converts a hipError_t to a Status object. iree_status_t iree_hal_rocm_result_to_status( - iree_hal_rocm_dynamic_symbols_t *syms, hipError_t result, const char *file, + iree_hal_rocm_dynamic_symbols_t* syms, hipError_t result, const char* file, uint32_t line); #ifdef __cplusplus
diff --git a/integrations/tensorflow/iree_tf_compiler/iree-import-tf-main.cpp b/integrations/tensorflow/iree_tf_compiler/iree-import-tf-main.cpp index feec3ff..ee95f1e 100644 --- a/integrations/tensorflow/iree_tf_compiler/iree-import-tf-main.cpp +++ b/integrations/tensorflow/iree_tf_compiler/iree-import-tf-main.cpp
@@ -224,8 +224,8 @@ iree_integrations::TF::buildTFImportPassPipeline(pm); if (failed(pm.run(*module))) { - llvm::errs() - << "Running iree-import-tf pass pipeline failed (see diagnostics)\n"; + llvm::errs() << "Running iree-import-tf TF import pass pipeline failed " + "(see diagnostics)\n"; return 2; } if (!saveTempMidLevelImport.empty()) { @@ -237,8 +237,8 @@ applyPassManagerCLOptions(pm); iree_integrations::MHLO::buildMHLOImportPassPipeline(pm); if (failed(pm.run(*module))) { - llvm::errs() - << "Running iree-import-tf pass pipeline failed (see diagnostics)\n"; + llvm::errs() << "Running iree-import-tf MHLO Import pass pipeline failed " + "(see diagnostics)\n"; return 2; } }
diff --git a/integrations/tensorflow/iree_tf_compiler/iree-import-xla-main.cpp b/integrations/tensorflow/iree_tf_compiler/iree-import-xla-main.cpp index 2e11d7f..08f270a 100644 --- a/integrations/tensorflow/iree_tf_compiler/iree-import-xla-main.cpp +++ b/integrations/tensorflow/iree_tf_compiler/iree-import-xla-main.cpp
@@ -100,6 +100,10 @@ static cl::opt<std::string> outputFilename("o", cl::desc("Output filename"), cl::value_desc("filename"), cl::init("-")); + static llvm::cl::opt<std::string> saveTempMhloInput( + "save-temp-mhlo-input", + llvm::cl::desc("Save the MHLO pipeline input IR to this file"), + llvm::cl::init("")); static llvm::cl::opt<std::string> saveTempIreeImport( "save-temp-iree-input", llvm::cl::desc("Save the resultant IR to this file (useful for saving an " @@ -251,6 +255,11 @@ return success(); }; + // Save temp output. + if (!saveTempMhloInput.empty()) { + if (failed(saveToFile(saveTempMhloInput))) return 10; + } + // Run passes. PassManager pm(&context, PassManager::Nesting::Implicit); applyPassManagerCLOptions(pm); @@ -264,8 +273,8 @@ iree_integrations::MHLO::createEmitDefaultIREEABIPass()); if (failed(pm.run(*module))) { - llvm::errs() - << "Running iree-xla-import pass pipeline failed (see diagnostics)\n"; + llvm::errs() << "Running iree-xla-import MHLO import pass pipeline failed " + "(see diagnostics)\n"; return 2; }
diff --git a/iree/base/alignment.h b/iree/base/alignment.h index 0f5953a..5229d97 100644 --- a/iree/base/alignment.h +++ b/iree/base/alignment.h
@@ -11,6 +11,8 @@ #define IREE_BASE_ALIGNMENT_H_ #include <stddef.h> +#include <stdint.h> +#include <string.h> #include "iree/base/config.h" #include "iree/base/target_platform.h" @@ -55,6 +57,178 @@ return (value + (alignment - 1)) & ~(alignment - 1); } +// Returns the size of a struct padded out to iree_max_align_t. +// This must be used when performing manual trailing allocation packing to +// ensure the alignment requirements of the trailing data are satisified. +// +// NOTE: do not use this if using VLAs (`struct { int trailing[]; }`) - those +// must precisely follow the normal sizeof(t) as the compiler does the padding +// for you. +// +// Example: +// some_buffer_ptr_t* p = NULL; +// iree_host_size_t total_size = iree_sizeof_struct(*buffer) + extra_data_size; +// IREE_CHECK_OK(iree_allocator_malloc(allocator, total_size, (void**)&p)); +#define iree_sizeof_struct(t) iree_host_align(sizeof(t), iree_max_align_t) + +//===----------------------------------------------------------------------===// +// Alignment-safe memory accesses +//===----------------------------------------------------------------------===// + +// Map little-endian byte indices in memory to the host memory order indices. +#if defined(IREE_ENDIANNESS_LITTLE) +#define IREE_LE_IDX_1(i) (i) +#define IREE_LE_IDX_2(i) (i) +#define IREE_LE_IDX_4(i) (i) +#define IREE_LE_IDX_8(i) (i) +#else +#define IREE_LE_IDX_1(i) (i) +#define IREE_LE_IDX_2(i) (1 - (i)) +#define IREE_LE_IDX_4(i) (3 - (i)) +#define IREE_LE_IDX_8(i) (7 - (i)) +#endif // IREE_ENDIANNESS_* + +#if IREE_MEMORY_ACCESS_ALIGNMENT_REQUIRED + +static inline uint8_t iree_unaligned_load_le_u8(const uint8_t* ptr) { + return *ptr; +} +static inline uint16_t iree_unaligned_load_le_u16(const uint16_t* ptr) { + const uint8_t* p = (const uint8_t*)ptr; + return ((uint16_t)p[IREE_LE_IDX_2(0)]) | ((uint16_t)p[IREE_LE_IDX_2(1)] << 8); +} +static inline uint32_t iree_unaligned_load_le_u32(const uint32_t* ptr) { + const uint8_t* p = (const uint8_t*)ptr; + return ((uint32_t)p[IREE_LE_IDX_4(0)]) | + ((uint32_t)p[IREE_LE_IDX_4(1)] << 8) | + ((uint32_t)p[IREE_LE_IDX_4(2)] << 16) | + ((uint32_t)p[IREE_LE_IDX_4(3)] << 24); +} +static inline uint64_t iree_unaligned_load_le_u64(const uint64_t* ptr) { + const uint8_t* p = (const uint8_t*)ptr; + return ((uint64_t)p[IREE_LE_IDX_8(0)]) | + ((uint64_t)p[IREE_LE_IDX_8(1)] << 8) | + ((uint64_t)p[IREE_LE_IDX_8(2)] << 16) | + ((uint64_t)p[IREE_LE_IDX_8(3)] << 24) | + ((uint64_t)p[IREE_LE_IDX_8(4)] << 32) | + ((uint64_t)p[IREE_LE_IDX_8(5)] << 40) | + ((uint64_t)p[IREE_LE_IDX_8(6)] << 48) | + ((uint64_t)p[IREE_LE_IDX_8(7)] << 56); +} +static inline float iree_unaligned_load_le_f32(const float* ptr) { + uint32_t uint_value = iree_unaligned_load_le_u32((const uint32_t*)ptr); + float value; + memcpy(&value, &uint_value, sizeof(value)); + return value; +} +static inline double iree_unaligned_load_le_f64(const double* ptr) { + uint64_t uint_value = iree_unaligned_load_le_u64((const uint64_t*)ptr); + double value; + memcpy(&value, &uint_value, sizeof(value)); + return value; +} + +static inline void iree_unaligned_store_le_u8(uint8_t* ptr, uint8_t value) { + *ptr = value; +} +static inline void iree_unaligned_store_le_u16(uint16_t* ptr, uint16_t value) { + uint8_t* p = (uint8_t*)ptr; + p[IREE_LE_IDX_2(0)] = value; + p[IREE_LE_IDX_2(1)] = value >> 8; +} +static inline void iree_unaligned_store_le_u32(uint32_t* ptr, uint32_t value) { + uint8_t* p = (uint8_t*)ptr; + p[IREE_LE_IDX_4(0)] = value; + p[IREE_LE_IDX_4(1)] = value >> 8; + p[IREE_LE_IDX_4(2)] = value >> 16; + p[IREE_LE_IDX_4(3)] = value >> 24; +} +static inline void iree_unaligned_store_le_u64(uint64_t* ptr, uint64_t value) { + uint8_t* p = (uint8_t*)ptr; + p[IREE_LE_IDX_8(0)] = value; + p[IREE_LE_IDX_8(1)] = value >> 8; + p[IREE_LE_IDX_8(2)] = value >> 16; + p[IREE_LE_IDX_8(3)] = value >> 24; + p[IREE_LE_IDX_8(4)] = value >> 32; + p[IREE_LE_IDX_8(5)] = value >> 40; + p[IREE_LE_IDX_8(6)] = value >> 48; + p[IREE_LE_IDX_8(7)] = value >> 56; +} +static inline void iree_unaligned_store_le_f32(float* ptr, float value) { + uint32_t uint_value; + memcpy(&uint_value, &value, sizeof(value)); + iree_unaligned_store_le_u32((uint32_t*)ptr, uint_value); +} +static inline void iree_unaligned_store_le_f64(double* ptr, double value) { + uint64_t uint_value; + memcpy(&uint_value, &value, sizeof(value)); + iree_unaligned_store_le_u64((uint64_t*)ptr, uint_value); +} + +#else + +#if defined(IREE_ENDIANNESS_LITTLE) + +#define iree_unaligned_load_le_u8(ptr) *(ptr) +#define iree_unaligned_load_le_u16(ptr) *(ptr) +#define iree_unaligned_load_le_u32(ptr) *(ptr) +#define iree_unaligned_load_le_u64(ptr) *(ptr) +#define iree_unaligned_load_le_f32(ptr) *(ptr) +#define iree_unaligned_load_le_f64(ptr) *(ptr) + +#define iree_unaligned_store_le_u8(ptr, value) *(ptr) = (value) +#define iree_unaligned_store_le_u16(ptr, value) *(ptr) = (value) +#define iree_unaligned_store_le_u32(ptr, value) *(ptr) = (value) +#define iree_unaligned_store_le_u64(ptr, value) *(ptr) = (value) +#define iree_unaligned_store_le_f32(ptr, value) *(ptr) = (value) +#define iree_unaligned_store_le_f64(ptr, value) *(ptr) = (value) + +#else + +#error "TODO(benvanik): little-endian load/store for big-endian archs" + +#endif // IREE_ENDIANNESS_* + +#endif // IREE_MEMORY_ACCESS_ALIGNMENT_REQUIRED + +// clang-format off + +// Dereferences |ptr| and returns the value. +// Automatically handles unaligned accesses on architectures that may not +// support them natively (or efficiently). Memory is treated as little-endian. +#define iree_unaligned_load_le(ptr) \ + _Generic((ptr), \ + int8_t*: iree_unaligned_load_le_u8((const uint8_t*)(ptr)), \ + uint8_t*: iree_unaligned_load_le_u8((const uint8_t*)(ptr)), \ + int16_t*: iree_unaligned_load_le_u16((const uint16_t*)(ptr)), \ + uint16_t*: iree_unaligned_load_le_u16((const uint16_t*)(ptr)), \ + int32_t*: iree_unaligned_load_le_u32((const uint32_t*)(ptr)), \ + uint32_t*: iree_unaligned_load_le_u32((const uint32_t*)(ptr)), \ + int64_t*: iree_unaligned_load_le_u64((const uint64_t*)(ptr)), \ + uint64_t*: iree_unaligned_load_le_u64((const uint64_t*)(ptr)), \ + float*: iree_unaligned_load_le_f32((const float*)(ptr)), \ + double*: iree_unaligned_load_le_f64((const double*)(ptr)) \ + ) + +// Dereferences |ptr| and writes the given |value|. +// Automatically handles unaligned accesses on architectures that may not +// support them natively (or efficiently). Memory is treated as little-endian. +#define iree_unaligned_store(ptr, value) \ + _Generic((ptr), \ + int8_t*: iree_unaligned_store_le_u8((uint8_t*)(ptr), value), \ + uint8_t*: iree_unaligned_store_le_u8((uint8_t*)(ptr), value), \ + int16_t*: iree_unaligned_store_le_u16((uint16_t*)(ptr), value), \ + uint16_t*: iree_unaligned_store_le_u16((uint16_t*)(ptr), value), \ + int32_t*: iree_unaligned_store_le_u32((uint32_t*)(ptr), value), \ + uint32_t*: iree_unaligned_store_le_u32((uint32_t*)(ptr), value), \ + int64_t*: iree_unaligned_store_le_u64((uint64_t*)(ptr), value), \ + uint64_t*: iree_unaligned_store_le_u64((uint64_t*)(ptr), value), \ + float*: iree_unaligned_store_le_f32((float*)(ptr), value), \ + double*: iree_unaligned_store_le_f64((double*)(ptr), value) \ + ) + +// clang-format on + #ifdef __cplusplus } // extern "C" #endif
diff --git a/iree/base/internal/threading_darwin.c b/iree/base/internal/threading_darwin.c index 0d9f4fd..7f3bd00 100644 --- a/iree/base/internal/threading_darwin.c +++ b/iree/base/internal/threading_darwin.c
@@ -82,7 +82,7 @@ // (including the user-specified entry_arg). iree_thread_t* thread = NULL; iree_status_t status = - iree_allocator_malloc(allocator, sizeof(iree_thread_t), (void**)&thread); + iree_allocator_malloc(allocator, sizeof(*thread), (void**)&thread); if (!iree_status_is_ok(status)) { IREE_TRACE_ZONE_END(z0); return status;
diff --git a/iree/base/internal/threading_pthreads.c b/iree/base/internal/threading_pthreads.c index ed7c59f..f3d66a8 100644 --- a/iree/base/internal/threading_pthreads.c +++ b/iree/base/internal/threading_pthreads.c
@@ -124,7 +124,7 @@ // (including the user-specified entry_arg). iree_thread_t* thread = NULL; iree_status_t status = - iree_allocator_malloc(allocator, sizeof(iree_thread_t), (void**)&thread); + iree_allocator_malloc(allocator, sizeof(*thread), (void**)&thread); if (!iree_status_is_ok(status)) { IREE_TRACE_ZONE_END(z0); return status;
diff --git a/iree/base/internal/threading_win32.c b/iree/base/internal/threading_win32.c index f6e5d24..6e550e3 100644 --- a/iree/base/internal/threading_win32.c +++ b/iree/base/internal/threading_win32.c
@@ -139,7 +139,7 @@ // (including the user-specified entry_arg). iree_thread_t* thread = NULL; iree_status_t status = - iree_allocator_malloc(allocator, sizeof(iree_thread_t), (void**)&thread); + iree_allocator_malloc(allocator, sizeof(*thread), (void**)&thread); if (!iree_status_is_ok(status)) { IREE_TRACE_ZONE_END(z0); return status; @@ -200,7 +200,11 @@ iree_thread_resume(thread); - WaitForSingleObject(thread->handle, INFINITE); + if (thread->id != GetCurrentThreadId()) { + // Join with the thread. Since threads can delete themselves we must ensure + // they don't try to join with themselves and deadlock. + WaitForSingleObject(thread->handle, INFINITE); + } CloseHandle(thread->handle); iree_thread_override_list_deinitialize(&thread->qos_override_list); iree_allocator_free(thread->allocator, thread);
diff --git a/iree/base/internal/wait_handle_poll.c b/iree/base/internal/wait_handle_poll.c index 805a19d..5dd51fc 100644 --- a/iree/base/internal/wait_handle_poll.c +++ b/iree/base/internal/wait_handle_poll.c
@@ -148,10 +148,10 @@ } iree_host_size_t user_handle_list_size = - capacity * sizeof(iree_wait_handle_t); + capacity * iree_sizeof_struct(iree_wait_handle_t); iree_host_size_t poll_fd_list_size = capacity * sizeof(struct pollfd); - iree_host_size_t total_size = - sizeof(iree_wait_set_t) + user_handle_list_size + poll_fd_list_size; + iree_host_size_t total_size = iree_sizeof_struct(iree_wait_set_t) + + user_handle_list_size + poll_fd_list_size; iree_wait_set_t* set = NULL; IREE_RETURN_IF_ERROR( @@ -161,7 +161,8 @@ iree_wait_set_clear(set); set->user_handles = - (iree_wait_handle_t*)((uint8_t*)set + sizeof(iree_wait_set_t)); + (iree_wait_handle_t*)((uint8_t*)set + + iree_sizeof_struct(iree_wait_set_t)); set->poll_fds = (struct pollfd*)((uint8_t*)set->user_handles + user_handle_list_size);
diff --git a/iree/base/internal/wait_handle_win32.c b/iree/base/internal/wait_handle_win32.c index 0a5bb56..af17e4a 100644 --- a/iree/base/internal/wait_handle_win32.c +++ b/iree/base/internal/wait_handle_win32.c
@@ -147,8 +147,8 @@ iree_host_size_t user_handle_list_size = capacity * sizeof(iree_wait_handle_t); iree_host_size_t native_handle_list_size = capacity * sizeof(HANDLE); - iree_host_size_t total_size = - sizeof(iree_wait_set_t) + user_handle_list_size + native_handle_list_size; + iree_host_size_t total_size = iree_sizeof_struct(iree_wait_set_t) + + user_handle_list_size + native_handle_list_size; iree_wait_set_t* set = NULL; IREE_RETURN_IF_ERROR( @@ -158,7 +158,8 @@ iree_wait_set_clear(set); set->user_handles = - (iree_wait_handle_t*)((uint8_t*)set + sizeof(iree_wait_set_t)); + (iree_wait_handle_t*)((uint8_t*)set + + iree_sizeof_struct(iree_wait_set_t)); set->native_handles = (HANDLE*)((uint8_t*)set->user_handles + user_handle_list_size);
diff --git a/iree/base/target_platform.h b/iree/base/target_platform.h index 4277d71..8052967 100644 --- a/iree/base/target_platform.h +++ b/iree/base/target_platform.h
@@ -30,6 +30,8 @@ // IREE_ENDIANNESS_LITTLE // IREE_ENDIANNESS_BIG // +// IREE_MEMORY_ACCESS_ALIGNMENT_REQUIRED (0/1) +// // IREE_COMPILER_CLANG // IREE_COMPILER_GCC // IREE_COMPILER_GCC_COMPAT @@ -138,6 +140,39 @@ #endif // __BYTE_ORDER__ //============================================================================== +// IREE_MEMORY_ACCESS_* +//============================================================================== +// Certain architectures have specific memory access requirements that require +// user-mode code changes to work at all or work at reasonable performance. + +#if !defined(IREE_MEMORY_ACCESS_ALIGNMENT_REQUIRED) + +#if defined(IREE_ARCH_ARM_32) || defined(IREE_ARCH_ARM_64) + +// Armv6‑M and Armv8-M (w/o the main extension) do not support unaligned access. +// The -munaligned-access and -mno-unaligned-access flags control this. +// https://www.keil.com/support/man/docs/armclang_ref/armclang_ref_sam1444138667173.htm +#if !defined(__ARM_FEATURE_UNALIGNED) +#define IREE_MEMORY_ACCESS_ALIGNMENT_REQUIRED 1 +#else +#define IREE_MEMORY_ACCESS_ALIGNMENT_REQUIRED 0 +#endif // !__ARM_FEATURE_UNALIGNED + +#elif defined(IREE_ARCH_RISCV_32) || defined(IREE_ARCH_RISCV_64) + +// Though unaligned access is part of the base spec it is allowed to be +// implemented with trap handlers. Bare-metal systems likely won't have these +// handlers and even on systems that do (linux) we don't want to be trapping for +// every load/store. +#define IREE_MEMORY_ACCESS_ALIGNMENT_REQUIRED 1 + +#endif // IREE_ARCH_* + +#else +#define IREE_MEMORY_ACCESS_ALIGNMENT_REQUIRED 0 +#endif // !IREE_MEMORY_ACCESS_ALIGNMENT_REQUIRED + +//============================================================================== // IREE_COMPILER_* //==============================================================================
diff --git a/iree/compiler/Codegen/SPIRV/ConvertToSPIRVPass.cpp b/iree/compiler/Codegen/SPIRV/ConvertToSPIRVPass.cpp index b79eac9..7c3cbb7 100644 --- a/iree/compiler/Codegen/SPIRV/ConvertToSPIRVPass.cpp +++ b/iree/compiler/Codegen/SPIRV/ConvertToSPIRVPass.cpp
@@ -13,6 +13,8 @@ // //===----------------------------------------------------------------------===// +#include <tuple> + #include "iree/compiler/Codegen/PassDetail.h" #include "iree/compiler/Codegen/Passes.h" #include "iree/compiler/Codegen/Utils/MarkerUtils.h" @@ -51,15 +53,19 @@ // Resource utilities //===----------------------------------------------------------------------===// -/// Inserts a resource evariable of the given `type` at the beginning of +/// Map from hal.interface.binding.subspan ops to their corresponding +/// spv.GlobalVariable ops. +using InterfaceResourceMap = + llvm::DenseMap<Operation *, spirv::GlobalVariableOp>; + +/// Creates a resource evariable of the given `type` at the beginning of /// `moduleOp`'s block via `symbolTable` and bind it to `set` and `binding`. -spirv::GlobalVariableOp insertResourceVariable(Location loc, Type type, +spirv::GlobalVariableOp createResourceVariable(Location loc, Type type, unsigned set, unsigned binding, bool alias, ModuleOp moduleOp, - SymbolTable *symbolTable, - OpBuilder::Listener *listener) { - OpBuilder builder(moduleOp.getContext(), listener); + SymbolTable *symbolTable) { std::string name = llvm::formatv("__resource_var_{0}_{1}_", set, binding); + OpBuilder builder(moduleOp.getContext()); auto variable = builder.create<spirv::GlobalVariableOp>(loc, type, name, set, binding); if (alias) variable->setAttr("aliased", builder.getUnitAttr()); @@ -74,35 +80,66 @@ return {bindingOp.set().getSExtValue(), bindingOp.binding().getSExtValue()}; } -/// Returns the set of resources that should be marked as aliased in SPIR-V. -llvm::DenseSet<Operation *> getAliasedResources(ModuleOp module) { - llvm::DenseSet<Operation *> aliasedResources; +/// Scans all hal.interface.binding.subspan ops in `module`, creates their +/// corresponding spv.GlobalVariables when needed, and returns the map. +/// The created variables need to have their types fixed later. +InterfaceResourceMap createResourceVariables(mlir::ModuleOp module) { + SymbolTable symbolTable(module); + InterfaceResourceMap interfaceToResourceVars; - for (FuncOp func : module.getOps<FuncOp>()) { + auto fns = llvm::to_vector<1>(module.getOps<FuncOp>()); + for (FuncOp func : llvm::reverse(fns)) { // Collect all interface ops and their (set, binding) pairs in this - // function. - SmallVector<Operation *, 4> interfaceOps; - SmallVector<std::pair<uint32_t, uint32_t>, 4> setBindings; - llvm::DenseMap<std::pair<uint32_t, uint32_t>, unsigned> setBindingCount; + // function. Use SmallVector here for a deterministic order. + SmallVector<IREE::HAL::InterfaceBindingSubspanOp, 8> interfaceOps; + SmallVector<std::pair<uint32_t, uint32_t>, 8> setBindings; + + // Use a map to see if we have different types for one (set, binding) pair, + // which will require creating multiple SPIR-V global variables. + llvm::DenseMap<std::pair<uint32_t, uint32_t>, llvm::DenseSet<Type>> + setBindingTypes; + func.walk([&](Operation *op) { - if (isa<IREE::HAL::InterfaceBindingSubspanOp>(op)) { - interfaceOps.emplace_back(op); - setBindings.emplace_back(getInterfaceSetAndBinding(op)); - ++setBindingCount[setBindings.back()]; - } + auto interfaceOp = dyn_cast<IREE::HAL::InterfaceBindingSubspanOp>(op); + if (!interfaceOp || interfaceOp.use_empty()) return; + interfaceOps.emplace_back(interfaceOp); + setBindings.emplace_back(getInterfaceSetAndBinding(interfaceOp)); + setBindingTypes[setBindings.back()].insert(interfaceOp.getType()); }); - // Perform analysis to determine whether we need to mark the resource as - // alias. This should happen when we have multiple resources binding to the - // same (set, binding) pair and they are used in the same function. - for (unsigned i = 0; i < interfaceOps.size(); ++i) { - if (setBindingCount[setBindings[i]] > 1) { - aliasedResources.insert(interfaceOps[i]); + // Keep track of created SPIR-V global variables. This allows us to + // deduplicate when possible to reduce generated SPIR-V blob size. + llvm::DenseMap<std::tuple<uint32_t, uint32_t, Type>, + spirv::GlobalVariableOp> + resourceVars; + + for (int i = interfaceOps.size() - 1; i >= 0; --i) { + auto interfaceOp = interfaceOps[i]; + const auto &setBinding = setBindings[i]; + + auto key = std::make_tuple(setBinding.first, setBinding.second, + interfaceOp.getType()); + auto var = resourceVars.lookup(key); + if (!var) { + // If we have multiple SPIR-V global variables bound to the same (set, + // binding) pair and they are used in the same function, those variables + // need to have alias decoration. + bool alias = setBindingTypes[setBindings[i]].size() > 1; + + // We are using the interface op's type for creating the global + // variable. It's fine. The correctness boundary is the pass. + // We will fix it up during conversion so it won't leak. + var = createResourceVariable( + interfaceOp.getLoc(), interfaceOp.getType(), setBinding.first, + setBinding.second, alias, module, &symbolTable); + resourceVars[key] = var; } + + interfaceToResourceVars[interfaceOp] = var; } } - return aliasedResources; + return interfaceToResourceVars; } } // namespace @@ -169,17 +206,19 @@ : public OpConversionPattern<IREE::HAL::InterfaceBindingSubspanOp> { HALInterfaceBindingSubspanConverter( TypeConverter &typeConverter, MLIRContext *context, - const llvm::DenseSet<Operation *> &aliasedResources, - SymbolTable *symbolTable, PatternBenefit benefit = 1) + const InterfaceResourceMap &interfaceToResourceVars, + PatternBenefit benefit = 1) : OpConversionPattern(typeConverter, context, benefit), - aliasedResources(aliasedResources), - symbolTable(symbolTable) {} + interfaceToResourceVars(interfaceToResourceVars) {} LogicalResult matchAndRewrite( IREE::HAL::InterfaceBindingSubspanOp interfaceOp, ArrayRef<Value> operands, ConversionPatternRewriter &rewriter) const override { - auto moduleOp = interfaceOp->template getParentOfType<ModuleOp>(); + if (interfaceOp.use_empty()) { + rewriter.eraseOp(interfaceOp); + return success(); + } Type resultType = interfaceOp.getOperation()->getResult(0).getType(); Type convertedType = this->getTypeConverter()->convertType(resultType); @@ -187,22 +226,18 @@ return interfaceOp.emitError() << "failed to convert SPIR-V type: " << resultType; } - auto setAndBinding = getInterfaceSetAndBinding(interfaceOp.getOperation()); - // We always create a new resource variable for the interface. - spirv::GlobalVariableOp varOp = insertResourceVariable( - interfaceOp.getLoc(), convertedType, setAndBinding.first, - setAndBinding.second, - aliasedResources.contains(interfaceOp.getOperation()), moduleOp, - symbolTable, rewriter.getListener()); + auto varOp = interfaceToResourceVars.lookup(interfaceOp); + // Fix up the variable's type. + varOp.typeAttr(TypeAttr::get(convertedType)); rewriter.replaceOpWithNewOp<spirv::AddressOfOp>(interfaceOp, varOp); + return success(); } private: - const llvm::DenseSet<Operation *> &aliasedResources; - SymbolTable *symbolTable; + const InterfaceResourceMap &interfaceToResourceVars; }; /// Pattern to lower operations that become a no-ops at this level. @@ -307,18 +342,12 @@ IREE::HAL::InterfaceWorkgroupCountOp, spirv::BuiltIn::NumWorkgroups>>( typeConverter, context); - // Create a symbol table for the current module op. This must be done before - // starting conversion. The conversion framework has a deferred nature. - // Actions like replacing/deleting ops are just recorded instead of directly - // applied. So for function conversion, we will see both the original function - // and the new replacement function during conversion process. That causes an - // issue for creating symbol tables, which scans all symbols and errors out - // when duplicates are found. We should be fine here tough becuase we only use - // this symbol table to insert new SPIR-V global variables. - SymbolTable symbolTable(moduleOp); - auto aliasedResources = getAliasedResources(moduleOp); - patterns.insert<HALInterfaceBindingSubspanConverter>( - typeConverter, context, aliasedResources, &symbolTable); + // Performs a prelimiary step to analyze all hal.interface.binding.subspan ops + // and create spv.GlobalVariables. + auto interfaceToResourceVars = createResourceVariables(moduleOp); + // For using use them in conversion. + patterns.insert<HALInterfaceBindingSubspanConverter>(typeConverter, context, + interfaceToResourceVars); /// Fold certain operations as no-ops: /// - linalg.reshape becomes a no-op since all memrefs are linearized in
diff --git a/iree/compiler/Codegen/SPIRV/test/convert_to_spirv.mlir b/iree/compiler/Codegen/SPIRV/test/convert_to_spirv.mlir index abc817f..a7d78af 100644 --- a/iree/compiler/Codegen/SPIRV/test/convert_to_spirv.mlir +++ b/iree/compiler/Codegen/SPIRV/test/convert_to_spirv.mlir
@@ -24,20 +24,43 @@ module attributes {spv.target_env = #spv.target_env<#spv.vce<v1.3, [Shader], []>, {}>} { // CHECK-LABEL: spv.module - // CHECK: spv.GlobalVariable @__resource_var_3_4_ bind(3, 4) : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer> - // CHECK: spv.GlobalVariable @__resource_var_1_2__0 bind(1, 2) {aliased} : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer> - // CHECK: spv.GlobalVariable @__resource_var_1_2_ bind(1, 2) {aliased} : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer> + // CHECK: spv.GlobalVariable @[[ARG0:.+]] bind(1, 2) : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer> + // CHECK: spv.GlobalVariable @[[ARG1_0:.+]] bind(1, 3) {aliased} : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer> + // CHECK: spv.GlobalVariable @[[ARG1_1:.+]] bind(1, 3) {aliased} : !spv.ptr<!spv.struct<(!spv.array<4 x vector<4xf32>, stride=16> [0])>, StorageBuffer> + // CHECK: spv.GlobalVariable @[[RET0:.+]] bind(3, 4) : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer> // CHECK: spv.func @resource_bindings_in_same_entry_func() func @resource_bindings_in_same_entry_func() { %c0 = constant 0 : index + + // Same type + // CHECK: spv.mlir.addressof @[[ARG0]] + // CHECK: spv.mlir.addressof @[[ARG0]] %0 = hal.interface.binding.subspan @io::@arg0[%c0] : memref<4x4xf32> %1 = hal.interface.binding.subspan @io::@arg0[%c0] : memref<4x4xf32> - %2 = hal.interface.binding.subspan @io::@ret0[%c0] : memref<4x4xf32> + + // Different type + // CHECK: spv.mlir.addressof @[[ARG1_0]] + // CHECK: spv.mlir.addressof @[[ARG1_1]] + %2 = hal.interface.binding.subspan @io::@arg1[%c0] : memref<4x4xf32> + %3 = hal.interface.binding.subspan @io::@arg1[%c0] : memref<4xvector<4xf32>> + + // CHECK: spv.mlir.addressof @[[RET0]] + %4 = hal.interface.binding.subspan @io::@ret0[%c0] : memref<4x4xf32> + + %5 = memref.load %0[%c0, %c0] : memref<4x4xf32> + %6 = memref.load %1[%c0, %c0] : memref<4x4xf32> + + %7 = memref.load %2[%c0, %c0] : memref<4x4xf32> + %8 = memref.load %3[%c0] : memref<4xvector<4xf32>> + + %9 = memref.load %4[%c0, %c0] : memref<4x4xf32> + return } hal.interface @io attributes {push_constants = 5 : index, sym_visibility = "private"} { hal.interface.binding @arg0, set=1, binding=2, type="StorageBuffer", access="Read" + hal.interface.binding @arg1, set=1, binding=3, type="StorageBuffer", access="Read" hal.interface.binding @ret0, set=3, binding=4, type="StorageBuffer", access="Write" } } @@ -46,18 +69,22 @@ module attributes {spv.target_env = #spv.target_env<#spv.vce<v1.3, [Shader], []>, {}>} { // CHECK-LABEL: spv.module - // CHECK: spv.GlobalVariable @[[FUNC2_RET:.+]] bind(3, 4) : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer> - // CHECK: spv.GlobalVariable @[[FUNC2_ARG:.+]] bind(1, 2) : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer> - // CHECK: spv.GlobalVariable @[[FUNC1_RET:.+]] bind(3, 4) : !spv.ptr<!spv.struct<(!spv.array<4 x vector<4xf32>, stride=16> [0])>, StorageBuffer> // CHECK: spv.GlobalVariable @[[FUNC1_ARG:.+]] bind(1, 2) : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer> + // CHECK: spv.GlobalVariable @[[FUNC1_RET:.+]] bind(3, 4) : !spv.ptr<!spv.struct<(!spv.array<4 x vector<4xf32>, stride=16> [0])>, StorageBuffer> + // CHECK: spv.GlobalVariable @[[FUNC2_ARG:.+]] bind(1, 2) : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer> + // CHECK: spv.GlobalVariable @[[FUNC2_RET:.+]] bind(3, 4) : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer> // CHECK: spv.func @resource_bindings_in_entry_func1() func @resource_bindings_in_entry_func1() { - // CHECK: spv.mlir.addressof @[[FUNC1_ARG:.+]] - // CHECK: spv.mlir.addressof @[[FUNC1_RET:.+]] + // CHECK: spv.mlir.addressof @[[FUNC1_ARG]] + // CHECK: spv.mlir.addressof @[[FUNC1_RET]] %c0 = constant 0 : index %0 = hal.interface.binding.subspan @io::@arg0[%c0] : memref<4x4xf32> %1 = hal.interface.binding.subspan @io::@ret0[%c0] : memref<4xvector<4xf32>> + + %2 = memref.load %0[%c0, %c0] : memref<4x4xf32> + %3 = memref.load %1[%c0] : memref<4xvector<4xf32>> + return } @@ -66,8 +93,12 @@ // CHECK: spv.mlir.addressof @[[FUNC2_ARG]] // CHECK: spv.mlir.addressof @[[FUNC2_RET]] %c0 = constant 0 : index - %0 = hal.interface.binding.subspan @io::@arg0[%c0] : memref<4x4xf32> - %1 = hal.interface.binding.subspan @io::@ret0[%c0] : memref<4x4xf32> + %0 = hal.interface.binding.subspan @io::@arg0[%c0] : memref<4x4xf32> // Same type as previous function + %1 = hal.interface.binding.subspan @io::@ret0[%c0] : memref<4x4xf32> // Different type as previous function + + %2 = memref.load %0[%c0, %c0] : memref<4x4xf32> + %3 = memref.load %1[%c0, %c0] : memref<4x4xf32> + return } @@ -85,6 +116,11 @@ %0 = hal.interface.binding.subspan @io::@arg0[%c0] : memref<8x5xf32> %1 = hal.interface.binding.subspan @io::@arg1[%c0] : memref<5xf32> %2 = hal.interface.binding.subspan @io::@ret0[%c0] : memref<8x5xf32> + + %3 = memref.load %0[%c0, %c0] : memref<8x5xf32> + %4 = memref.load %1[%c0] : memref<5xf32> + %5 = memref.load %2[%c0, %c0] : memref<8x5xf32> + return } hal.interface @io attributes {sym_visibility = "private"} { @@ -93,14 +129,17 @@ hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write|Discard" } } + +// Explicitly check the variable symbols + // CHECK-LABEL: spv.module -// CHECK-DAG: spv.GlobalVariable @[[RET0:.+]] bind(0, 2) -// CHECK-DAG: spv.GlobalVariable @[[ARG1:.+]] bind(0, 1) -// CHECK-DAG: spv.GlobalVariable @[[ARG0:.+]] bind(0, 0) +// CHECK: spv.GlobalVariable @__resource_var_0_0_ bind(0, 0) +// CHECK: spv.GlobalVariable @__resource_var_0_1_ bind(0, 1) +// CHECK: spv.GlobalVariable @__resource_var_0_2_ bind(0, 2) // CHECK: spv.func -// CHECK-DAG: %{{.+}} = spv.mlir.addressof @[[RET0]] -// CHECK-DAG: %{{.+}} = spv.mlir.addressof @[[ARG0]] -// CHECK-DAG: %{{.+}} = spv.mlir.addressof @[[ARG1]] +// CHECK: %{{.+}} = spv.mlir.addressof @__resource_var_0_0_ +// CHECK: %{{.+}} = spv.mlir.addressof @__resource_var_0_1_ +// CHECK: %{{.+}} = spv.mlir.addressof @__resource_var_0_2_ // -----
diff --git a/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp b/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp index 8bf33ff..ce2f8cb 100644 --- a/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp +++ b/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp
@@ -695,6 +695,43 @@ } }; +// Replace `flow.tensor.splat`-`flow.tensor.load` op-pairs by the input +// primitive value for the splat op. +struct FoldSplatLoadIntoPrimitive : public OpRewritePattern<TensorLoadOp> { + using OpRewritePattern<TensorLoadOp>::OpRewritePattern; + + LogicalResult matchAndRewrite(TensorLoadOp loadOp, + PatternRewriter &rewriter) const override { + auto sourceOp = + dyn_cast_or_null<TensorSplatOp>(loadOp.source().getDefiningOp()); + + if (!sourceOp) return failure(); + + rewriter.replaceOp(loadOp, sourceOp.value()); + return success(); + } +}; + +struct FoldSplatReshapeIntoSplat : public OpRewritePattern<TensorSplatOp> { + using OpRewritePattern<TensorSplatOp>::OpRewritePattern; + + LogicalResult matchAndRewrite(TensorSplatOp splatOp, + PatternRewriter &rewriter) const override { + if (!splatOp.result().hasOneUse()) return failure(); + + auto reshapeOp = dyn_cast_or_null<TensorReshapeOp>( + splatOp.result().use_begin()->getOwner()); + if (!reshapeOp) return failure(); + + rewriter.replaceOpWithNewOp<TensorSplatOp>( + reshapeOp, reshapeOp.result().getType(), splatOp.value(), + reshapeOp.result_dims()); + rewriter.eraseOp(splatOp); + + return success(); + } +}; + } // namespace void TensorReshapeOp::getCanonicalizationPatterns( @@ -702,6 +739,11 @@ results.insert<FlattenTensorReshapeChain>(context); } +void TensorLoadOp::getCanonicalizationPatterns( + OwningRewritePatternList &results, MLIRContext *context) { + results.insert<FoldSplatLoadIntoPrimitive>(context); +} + OpFoldResult TensorLoadOp::fold(ArrayRef<Attribute> operands) { if (auto source = operands[0].dyn_cast_or_null<ElementsAttr>()) { // Load directly from the constant source tensor. @@ -739,6 +781,12 @@ return {}; } +void TensorSplatOp::getCanonicalizationPatterns( + OwningRewritePatternList &results, MLIRContext *context) { + // TODO(benvanik): canonicalize splat+slice to smaller splat. + results.insert<FoldSplatReshapeIntoSplat>(context); +} + OpFoldResult TensorSplatOp::fold(ArrayRef<Attribute> operands) { if (operands.size() == 1 && operands.front()) { // Splat value is constant and we can fold the operation.
diff --git a/iree/compiler/Dialect/Flow/IR/FlowOps.td b/iree/compiler/Dialect/Flow/IR/FlowOps.td index 3b14102..cd1b2f4 100644 --- a/iree/compiler/Dialect/Flow/IR/FlowOps.td +++ b/iree/compiler/Dialect/Flow/IR/FlowOps.td
@@ -85,7 +85,7 @@ }]; let arguments = (ins - FLOW_VariableRefAttr:$variable + Arg<FLOW_VariableRefAttr, "", [MemRead]>:$variable ); let results = (outs AnyType:$result @@ -135,7 +135,7 @@ let arguments = (ins AnyType:$value, - FLOW_VariableRefAttr:$variable + Arg<FLOW_VariableRefAttr, "", [MemWrite]>:$variable ); let assemblyFormat = "$value `,` $variable attr-dict `:` type($value)"; @@ -822,6 +822,7 @@ // TODO(benvanik): canonicalize to slice+load if dims are known. let hasFolder = 1; + let hasCanonicalizer = 1; } def FLOW_TensorStoreOp : FLOW_PureOp<"tensor.store", [ @@ -899,7 +900,7 @@ bool isTransfer() { return true; } }]; - // TODO(benvanik): canonicalize splat+slice to smaller splat. + let hasCanonicalizer = 1; let hasFolder = 1; }
diff --git a/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir b/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir index 62cab58..724b8e5 100644 --- a/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir +++ b/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir
@@ -411,3 +411,34 @@ // CHECK: return %[[RESULT]] return %1 : tensor<?x?xf32> } + +// ----- + +// CHECK-LABEL: @foldSplatLoadIntoPrimitive +// CHECK-SAME: (%[[arg0:.+]]: f32, %[[arg1:.+]]: index, %[[arg2:.+]]: index) +func @foldSplatLoadIntoPrimitive(%arg0 : f32, %arg1 : index, %arg2 : index) -> f32 { + // CHECK-NEXT: return %[[arg0]] : f32 + %0 = flow.tensor.splat %arg0 : tensor<4x4xf32> + %1 = flow.tensor.load %0[%arg1, %arg2] : tensor<4x4xf32> + return %1 : f32 +} + +// ----- + +// CHECK-LABEL: @foldSplatReshapeIntoSplat +func @foldSplatReshapeIntoSplat(%arg0 : f32) -> tensor<16xf32> { + // CHECK-NEXT: %0 = flow.tensor.splat %arg0 : tensor<16xf32> + // CHECK-NEXT: return %0 : tensor<16xf32> + %0 = flow.tensor.splat %arg0 : tensor<4x4xf32> + %1 = flow.tensor.reshape %0 : tensor<4x4xf32> -> tensor<16xf32> + return %1 : tensor<16xf32> +} + +// CHECK-LABEL: @foldSplatReshapeIntoSplatDynamic +func @foldSplatReshapeIntoSplatDynamic(%arg0 : f32, %arg1 : index, %arg2 : index, %arg3 : index) -> tensor<?x?xf32> { + // CHECK-NEXT: %0 = flow.tensor.splat %arg0 : tensor<?x?xf32>{%arg2, %arg3} + // CHECK-NEXT: return %0 : tensor<?x?xf32> + %0 = flow.tensor.splat %arg0 : tensor<?x4xf32>{%arg1} + %1 = flow.tensor.reshape %0 : tensor<?x4xf32>{%arg1} -> tensor<?x?xf32>{%arg2, %arg3} + return %1 : tensor<?x?xf32> +}
diff --git a/iree/compiler/Dialect/Flow/Transforms/BUILD b/iree/compiler/Dialect/Flow/Transforms/BUILD index 167a506..8b26a4b 100644 --- a/iree/compiler/Dialect/Flow/Transforms/BUILD +++ b/iree/compiler/Dialect/Flow/Transforms/BUILD
@@ -53,6 +53,7 @@ "Passes.cpp", "PromoteI1ToI8Pass.cpp", "PromoteTensorLoads.cpp", + "SimplifyVariableAccesses.cpp", "StripAndSplatConstantVariables.cpp", "TypeConverter.cpp", ],
diff --git a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt index 57484a1..e8c00f3 100644 --- a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt +++ b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -50,6 +50,7 @@ "Passes.cpp" "PromoteI1ToI8Pass.cpp" "PromoteTensorLoads.cpp" + "SimplifyVariableAccesses.cpp" "StripAndSplatConstantVariables.cpp" "TypeConverter.cpp" DEPS
diff --git a/iree/compiler/Dialect/Flow/Transforms/DestructiveUpdateUtils.cpp b/iree/compiler/Dialect/Flow/Transforms/DestructiveUpdateUtils.cpp index 69249fc..81483a4 100644 --- a/iree/compiler/Dialect/Flow/Transforms/DestructiveUpdateUtils.cpp +++ b/iree/compiler/Dialect/Flow/Transforms/DestructiveUpdateUtils.cpp
@@ -14,6 +14,7 @@ #include "iree/compiler/Dialect/Flow/IR/FlowOps.h" #include "iree/compiler/Dialect/HAL/IR/HALOps.h" #include "iree/compiler/Dialect/IREE/IR/IREEDialect.h" +#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtInterfaces.h" #include "llvm/ADT/TypeSwitch.h" #include "llvm/Support/Debug.h" #include "mlir/Dialect/Linalg/IR/LinalgOps.h" @@ -55,17 +56,28 @@ }; // TODO(nicolasvasilache): Use some interface instead of op names directly. -static bool hasDestructiveUpdateSubTensorUses( - BlockArgument arg, SpecialTerminatorOpCapture &capture) { +static bool hasDestructiveUpdateUses(BlockArgument arg, + SpecialTerminatorOpCapture &capture) { SmallVector<Operation *> reads; - SmallVector<tensor::InsertSliceOp> writes; + SmallVector<Operation *> writes; for (OpOperand &u : arg.getUses()) { - if (auto subTensorInsertOp = - dyn_cast<tensor::InsertSliceOp>(u.getOwner())) { - writes.push_back(subTensorInsertOp); - } else { - reads.push_back(u.getOwner()); - } + TypeSwitch<Operation *, void>(u.getOwner()) + .Case<linalg::LinalgOp, linalg_ext::LinalgExtOp>( + [&](auto linalgLikeOp) { + if (linalgLikeOp.isOutputTensor(&u)) { + writes.push_back(linalgLikeOp); + } else { + reads.push_back(linalgLikeOp); + } + }) + .Case<tensor::InsertSliceOp>([&](tensor::InsertSliceOp sliceOp) { + if (sliceOp.dest() == u.get()) { + writes.push_back(sliceOp); + } else { + reads.push_back(sliceOp); + } + }) + .Default([&](Operation *op) { reads.push_back(op); }); } // For now, only allow exactly a single SubTensorInsertOp that must be // dominated by all SubTensorOp. @@ -77,7 +89,7 @@ if (!domInfo.properlyDominates(read, writes.front())) { LLVM_DEBUG(llvm::dbgs() << "non-destructive use-def: " << *read << " does not properly dominate " - << *(writes.front().getOperation()) << "\n"); + << *(writes.front()) << "\n"); return false; } } @@ -133,8 +145,7 @@ // Case 2: multiple uses from an scf::ForOp then this must be used only by // SubTensorOp / SubTensorInsertOp with proper dominance. if (!regionArg.hasOneUse()) { - if (!hasDestructiveUpdateSubTensorUses(regionArg, capture)) - return nullptr; + if (!hasDestructiveUpdateUses(regionArg, capture)) return nullptr; return returnValue; } @@ -145,8 +156,7 @@ // Case 3a: Single use which is not an scf::ForOp, it may still be a // single SubTensor / SubTensorInsertOp. if (!innerForOp) { - if (!hasDestructiveUpdateSubTensorUses(regionArg, capture)) - return nullptr; + if (!hasDestructiveUpdateUses(regionArg, capture)) return nullptr; return returnValue; } @@ -216,31 +226,67 @@ return success(); } -static LogicalResult rewriteSubTensorInsertInPlace(OpBuilder &b, - tensor::InsertSliceOp op, - Value target) { - LLVM_DEBUG(llvm::dbgs() << "RewriteSubTensorInsertInPlace: " - << *(op.getOperation()) << "\n"); - if (!op.getResult().hasOneUse()) { - return op.emitError("not a single use operation"); +template <typename OpTy> +static LogicalResult rewriteDestructiveUpdateInPlace(OpBuilder &b, + OpTy linalgLikeOp, + Value target) { + LLVM_DEBUG(llvm::dbgs() << "RewriteDestructiveUpdateInPlace: " + << *linalgLikeOp.getOperation() << "\n"); + if (!linalgLikeOp->hasOneUse()) { + return linalgLikeOp.emitError("not a single use operation"); } - Operation *user = *(op.getResult().getUsers().begin()); - if (isa<scf::YieldOp>(user)) { - auto dest = op.dest(); - if (!dest.isa<BlockArgument>()) { - return op.emitError("dest is not a argument to the loop"); + OpOperand &use = *(linalgLikeOp->use_begin()); + if (isa<scf::YieldOp>(use.getOwner())) { + OpResult usedResult = use.get().cast<OpResult>(); + Value dest = + linalgLikeOp.getOutputOperand(usedResult.getResultNumber())->get(); + if (!dest || !dest.isa<BlockArgument>()) { + return linalgLikeOp.emitError("dest is not a argument to the loop"); } OpBuilder::InsertionGuard g(b); - b.setInsertionPoint(op); + b.setInsertionPointAfter(linalgLikeOp); // Kills the SSA use-def chain. - op.replaceAllUsesWith(dest); + usedResult.replaceAllUsesWith(dest); + + b.create<IREE::Flow::DispatchTensorStoreOp>(linalgLikeOp.getLoc(), + usedResult, target); + + return success(); + } + return failure(); +} + +/// Rewrites destructive in-place updates with the update operation being +/// tensor.insert_slice. +template <> +LogicalResult rewriteDestructiveUpdateInPlace<tensor::InsertSliceOp>( + OpBuilder &b, tensor::InsertSliceOp insertSliceOp, Value target) { + LLVM_DEBUG(llvm::dbgs() << "RewriteDestructiveUpdateInPlace: " + << *insertSliceOp.getOperation() << "\n"); + if (!insertSliceOp->hasOneUse()) { + return insertSliceOp.emitError("not a single use operation"); + } + + OpOperand &use = *(insertSliceOp->use_begin()); + if (isa<scf::YieldOp>(use.getOwner())) { + OpResult usedResult = use.get().cast<OpResult>(); + Value dest = insertSliceOp.dest(); + if (!dest || !dest.isa<BlockArgument>()) { + return insertSliceOp.emitError("dest is not a argument to the loop"); + } + OpBuilder::InsertionGuard g(b); + b.setInsertionPoint(insertSliceOp); + + // Kills the SSA use-def chain. + usedResult.replaceAllUsesWith(dest); b.create<IREE::Flow::DispatchTensorStoreOp>( - op.getLoc(), op.source(), target, op.offsets(), op.sizes(), - op.strides(), op.static_offsets(), op.static_sizes(), - op.static_strides()); + insertSliceOp->getLoc(), insertSliceOp.source(), target, + insertSliceOp.offsets(), insertSliceOp.sizes(), insertSliceOp.strides(), + insertSliceOp.static_offsets(), insertSliceOp.static_sizes(), + insertSliceOp.static_strides()); return success(); } @@ -366,11 +412,17 @@ << "\n"); // Try to rewrite inplace. - if (failed(rewriteSubTensorInsertInPlace( - b, cast<tensor::InsertSliceOp>(capture.rootDestructiveUpdate), - target))) { - return failure(); - } + auto status = + TypeSwitch<Operation *, LogicalResult>(capture.rootDestructiveUpdate) + .Case<linalg::LinalgOp, linalg_ext::LinalgExtOp, + tensor::InsertSliceOp>([&](auto op) { + if (failed(rewriteDestructiveUpdateInPlace(b, op, target))) { + return failure(); + } + return success(); + }) + .Default([&](Operation *) { return failure(); }); + if (failed(status)) return failure(); if (scf::ForOp loopOp = dyn_cast<scf::ForOp>(outermostProducingOp)) loopOp.walk( @@ -415,9 +467,9 @@ capture.initValue = op.value(); Value sourceValue = isADestructiveUpdatePattern(capture.initValue, capture); - if (!sourceValue || !isa_and_nonnull<tensor::InsertSliceOp>( - capture.rootDestructiveUpdate)) + if (!sourceValue) { return WalkResult::advance(); + } if (failed(rewriteDestructiveUpdateInPlace(b, capture, op.target()))) { return WalkResult::interrupt();
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp index 6185eac..5fca973 100644 --- a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp +++ b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
@@ -65,9 +65,14 @@ namespace Flow { void buildFlowTransformPassPipeline(OpPassManager &passManager) { - // Perform initial cleanup. - // NOTE: There is no principled reason to be doing this here. But also ensures - // some consistency at the tool boundary. + // Simplify flow.variable accesses early on; this can help with dispatch + // region formation as redundant store-loads are removed. + passManager.addNestedPass<FuncOp>( + IREE::Flow::createSimplifyVariableAccessesPass()); + + // Perform cleanup after variable simplification as more canonicalizers may be + // able to kick in. + passManager.addNestedPass<FuncOp>(mlir::createCanonicalizerPass()); passManager.addNestedPass<FuncOp>(mlir::createCSEPass()); // Replaces variables with !shapex.ranked_shape types with individual @@ -169,9 +174,11 @@ // Reorder blocks to increase the grouping of streamable ops. passManager.addNestedPass<FuncOp>( IREE::Flow::createHoistUnstreamableOpsPass()); + // The hoisting pass does some reordering. Canonicalize to avoid unnecessary // arbitrary ordering. passManager.addNestedPass<FuncOp>(mlir::createCanonicalizerPass()); + passManager.addNestedPass<FuncOp>(mlir::createCSEPass()); // Clone constants that escape basic blocks until we have better analysis. passManager.addNestedPass<FuncOp>(
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.h b/iree/compiler/Dialect/Flow/Transforms/Passes.h index 5f4403a..445190e 100644 --- a/iree/compiler/Dialect/Flow/Transforms/Passes.h +++ b/iree/compiler/Dialect/Flow/Transforms/Passes.h
@@ -131,6 +131,10 @@ // Reorders blocks to hoist ops that cannot be put into streams. std::unique_ptr<OperationPass<FuncOp>> createHoistUnstreamableOpsPass(); +// Hoists loads and sinks stores to variables to decrease data dependency +// regions. +std::unique_ptr<OperationPass<FuncOp>> createSimplifyVariableAccessesPass(); + // TODO(benvanik): cross-function stream flows. // Inserts clones of constant values where they may be required.
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.td b/iree/compiler/Dialect/Flow/Transforms/Passes.td index c8fa51c..4dfc814 100644 --- a/iree/compiler/Dialect/Flow/Transforms/Passes.td +++ b/iree/compiler/Dialect/Flow/Transforms/Passes.td
@@ -57,7 +57,7 @@ def FormStreams : Pass<"iree-flow-form-streams", "FuncOp"> { - let summary = "Identifies dispatches that can be grouped into streams within functions"; + let summary = "Identifies dispatches that can be grouped into streams within functions."; let constructor = "mlir::iree_compiler::IREE::Flow::createFormStreamsPass()"; } @@ -75,7 +75,7 @@ def InjectDispatchTracing : Pass<"iree-flow-inject-dispatch-tracing", "FuncOp"> { - let summary = "Injects dispatch region tracing"; + let summary = "Injects dispatch region tracing."; let constructor = "mlir::iree_compiler::IREE::Flow::createInjectDispatchTracingPass()"; } @@ -128,6 +128,12 @@ let constructor = "mlir::iree_compiler::IREE::Flow::createPromoteTensorLoadsPass()"; } +def SimplifyVariableAccesses : + Pass<"iree-flow-simplify-variable-accesses", "FuncOp"> { + let summary = "Hoist loads and sinks stores to variables to decrease data dependency regions."; + let constructor = "mlir::iree_compiler::IREE::Flow::createSimplifyVariableAccessesPass()"; +} + def StripAndSplatConstantVariables : Pass<"iree-flow-strip-and-splat-constant-variables", "ModuleOp"> { let summary = "Strips constant flow.variables and replaces them with splats.";
diff --git a/iree/compiler/Dialect/Flow/Transforms/SimplifyVariableAccesses.cpp b/iree/compiler/Dialect/Flow/Transforms/SimplifyVariableAccesses.cpp new file mode 100644 index 0000000..f190956 --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/SimplifyVariableAccesses.cpp
@@ -0,0 +1,270 @@ +// Copyright 2021 The IREE Authors +// +// 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 + +#include <algorithm> +#include <iterator> + +#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h" +#include "iree/compiler/Dialect/Flow/IR/FlowOps.h" +#include "iree/compiler/Dialect/Flow/Transforms/PassDetail.h" +#include "iree/compiler/Dialect/Flow/Transforms/Passes.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassRegistry.h" + +#define DEBUG_TYPE "iree-flow-simplify-variable-accesses" + +namespace mlir { +namespace iree_compiler { +namespace IREE { +namespace Flow { + +// Builds symbol ref set for all immutable variables in |moduleOp|. +static DenseSet<StringRef> gatherImmutableVariables(ModuleOp moduleOp) { + DenseSet<StringRef> set; + moduleOp.walk([&](IREE::Flow::VariableOp variableOp) { + if (!variableOp.is_mutable()) { + set.insert(variableOp.sym_name()); + } + }); + return set; +} + +// Hoists all loads of immutable variables in |funcOp| to the entry block. +// |immutableVariables| is used for lookups of which variables are immutable. +static void hoistImmutableLoads(FuncOp funcOp, + DenseSet<StringRef> &immutableVariables) { + // Since CSE of loads isn't a thing yet we perform a basic deduping here by + // folding all subsequent loads into the first one found. This works only for + // immutable variables as otherwise we'd have to ensure stores and + // side-effects were properly observed. + DenseMap<Attribute, Operation *> loadOps; + auto *entryBlock = &funcOp.getBlocks().front(); + Operation *lastEntryOp = nullptr; + for (auto &block : funcOp) { + for (auto op : llvm::make_early_inc_range( + block.getOps<IREE::Flow::VariableLoadOp>())) { + if (immutableVariables.contains(op.variable())) { + auto variableRef = op.variableAttr().cast<Attribute>(); + auto it = loadOps.find(variableRef); + if (it == loadOps.end()) { + // Move to entry block; even if it's already there (so loads are + // hoisted at the same time). + LLVM_DEBUG(llvm::dbgs() + << "moving immutable variable " << op.variable() + << " load to the entry block\n"); + if (lastEntryOp) { + op->moveAfter(lastEntryOp); + } else { + op->moveBefore(entryBlock, entryBlock->begin()); + } + loadOps[variableRef] = op; + lastEntryOp = op; + } else { + LLVM_DEBUG(llvm::dbgs() + << "CSE'ing immutable variable " << op.variable() << "\n"); + op->replaceAllUsesWith(it->getSecond()); + op->erase(); + } + } + } + } +} + +static bool doesOpBlockMotion(Operation *op) { + return isa<mlir::CallOpInterface>(op) || + op->hasTrait<OpTrait::IREE::YieldPoint>(); +} + +static void moveOpUpInBlock(Block &block, Operation *op) { + while (op->getPrevNode()) { + if (doesOpBlockMotion(op->getPrevNode())) break; + op->moveBefore(op->getPrevNode()); + } +} + +static void moveOpDownInBlock(Block &block, Operation *op) { + while (op->getNextNode() != block.getTerminator()) { + if (doesOpBlockMotion(op->getNextNode())) break; + op->moveAfter(op->getNextNode()); + } +} + +// Optimizes the load/store ops for each given bucket. +// Returns true if any op was removed. +static bool optimizeBuckets( + Block &block, std::map<StringRef, SmallVector<Operation *>> &buckets) { + bool didRemoveAny = false; + for (auto &bucket : buckets) { + // First perform basic load-store forwarding and such. + auto &ops = bucket.second; + for (int i = ops.size() - 1; i >= 1; --i) { + auto previous = ops[i - 1]; + auto current = ops[i]; + if (isa<IREE::Flow::VariableStoreOp>(previous) && + isa<IREE::Flow::VariableLoadOp>(current)) { + // RAW - forward the stored variable to the following use. + auto storedValue = previous->getOperand(0); + LLVM_DEBUG({ + llvm::dbgs() << "RAW: replacing load with previous store value:\n"; + current->dump(); + llvm::dbgs() << "->\n"; + storedValue.dump(); + }); + current->replaceAllUsesWith(ValueRange{storedValue}); + ops.erase(ops.begin() + i); + current->erase(); + didRemoveAny = true; + } else if (isa<IREE::Flow::VariableLoadOp>(previous) && + isa<IREE::Flow::VariableLoadOp>(current)) { + // RAR - forward the loaded variable to the following use. + LLVM_DEBUG({ + llvm::dbgs() << "RAR: replacing subsequent load with op:\n"; + current->dump(); + llvm::dbgs() << "->\n"; + previous->dump(); + }); + current->replaceAllUsesWith(previous); + ops.erase(ops.begin() + i); + current->erase(); + didRemoveAny = true; + } else if (isa<IREE::Flow::VariableStoreOp>(previous) && + isa<IREE::Flow::VariableStoreOp>(current)) { + // WAW - remove the first store. + LLVM_DEBUG({ + llvm::dbgs() << "WAW: erasing source op:\n"; + previous->dump(); + llvm::dbgs() << "\nand keeping subsequent op:\n"; + current->dump(); + }); + ops.erase(ops.begin() + i - 1); + previous->erase(); + didRemoveAny = true; + } + } + assert(!ops.empty()); + + if (auto loadOp = dyn_cast<IREE::Flow::VariableLoadOp>(ops.front())) { + // If the head op is a load we can move that to the top of the block. + LLVM_DEBUG(llvm::dbgs() << "moving mutable variable " << loadOp.variable() + << " load upward\n"); + moveOpUpInBlock(block, ops.front()); + } + if (auto storeOp = dyn_cast<IREE::Flow::VariableStoreOp>(ops.back())) { + // If the tail op is a store we can move that to the bottom of the block. + LLVM_DEBUG(llvm::dbgs() << "moving mutable variable " + << storeOp.variable() << " store downward\n"); + moveOpDownInBlock(block, ops.back()); + } + } + return didRemoveAny; +} + +// Hoists loads and sinks stores to the boundary of |block| when safe. +// |immutableVariables| is used for lookups of which variables are immutable. +// +// Basic algorithm (repeat until no op removals): +// for each op: +// if immutable: skip +// add to load/store buckets (sorted vector) +// for each bucket (symbol): +// walk ops in reverse: +// if (prev == store && this == load) // RAW +// replace load with store source +// if (prev == load && this == load) // RAR +// replace with first load +// if (prev == store && this == store) // WAW +// remove first store +// if (head == load) move load to front +// if (tail == store) move store to back +// +// Returns true if there were any removals and the block should be reprocessed. +static bool rearrangeBlockVariableAccesses( + Block &block, DenseSet<StringRef> &immutableVariables) { + // Produce [symbol_name, [op, op, op, ...]] buckets. + // NOTE: we use a map here so that we are deterministically ordered. This may + // not be needed but the variable count is low and it's nice to not care about + // op order issues. + // + // Because there may be ops that we can't optimize across (calls/etc) we + // handle flushing buckets on demand. + std::map<StringRef, SmallVector<Operation *>> buckets; + bool didRemoveAny = false; + for (auto &op : block) { + if (auto loadOp = dyn_cast<IREE::Flow::VariableLoadOp>(op)) { + if (immutableVariables.contains(loadOp.variable())) continue; + buckets[loadOp.variable()].push_back(&op); + } else if (auto storeOp = dyn_cast<IREE::Flow::VariableStoreOp>(op)) { + buckets[storeOp.variable()].push_back(&op); + } else if (doesOpBlockMotion(&op)) { + // Split point - all accesses after this point must not assume anything + // about accesses before it. + didRemoveAny = optimizeBuckets(block, buckets) || didRemoveAny; + buckets.clear(); + } + } + didRemoveAny = optimizeBuckets(block, buckets) || didRemoveAny; + return didRemoveAny; +} + +namespace { + +class SimplifyVariableAccessesPass + : public SimplifyVariableAccessesBase<SimplifyVariableAccessesPass> { + public: + void runOnOperation() override { + auto funcOp = getOperation(); + if (funcOp.empty()) return; + + auto moduleOp = funcOp->getParentOfType<mlir::ModuleOp>(); + assert(moduleOp && "func not in a module"); + + // Build a set of all immutable variables for fast lookup. + auto immutableVariables = gatherImmutableVariables(moduleOp); + + // Hoist immutable variables first. These have no hazards and don't care + // about control flow - like `constant` - so getting them handled first + // avoids the need for us to do the full analysis. + hoistImmutableLoads(funcOp, immutableVariables); + + // We can't optimize the function if there are indirect loads/stores. + // Note that constant loads are still ok above. + for (auto &block : funcOp) { + for (auto &op : block) { + if (isa<IREE::Flow::VariableLoadIndirectOp>(op) || + isa<IREE::Flow::VariableStoreIndirectOp>(op)) { + LLVM_DEBUG(llvm::dbgs() + << "bailing on variable access simplification: indirect " + "accesses present in function\n"); + return; + } + } + } + + // For each block in the function hoist loads and sink stores. + // This does no cross-block movement, though it really should. Maybe when a + // real compiler engineer sees this they'll be inspired to do this properly. + for (auto &block : funcOp) { + LLVM_DEBUG(llvm::dbgs() << "==== REARRANGING BLOCK ACCESSES ====\n"); + while (rearrangeBlockVariableAccesses(block, immutableVariables)) { + // NOTE: block is processed until no more ops are removed. Will always + // end in a fixed amount of time as ops are only removed from the block. + } + } + } +}; + +} // namespace + +std::unique_ptr<OperationPass<FuncOp>> createSimplifyVariableAccessesPass() { + return std::make_unique<SimplifyVariableAccessesPass>(); +} + +} // namespace Flow +} // namespace IREE +} // namespace iree_compiler +} // namespace mlir
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/BUILD b/iree/compiler/Dialect/Flow/Transforms/test/BUILD index b8e2e4f..b16ee25 100644 --- a/iree/compiler/Dialect/Flow/Transforms/test/BUILD +++ b/iree/compiler/Dialect/Flow/Transforms/test/BUILD
@@ -38,6 +38,7 @@ "pad_tensor_to_tensor.mlir", "promote_i1_to_i8.mlir", "promote_tensor_loads.mlir", + "simplify_variable_accesses.mlir", "strip_and_splat_constant_variables.mlir", "transformation.mlir", ],
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt b/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt index 1d3a4d1..01a4294 100644 --- a/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt +++ b/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
@@ -35,6 +35,7 @@ "pad_tensor_to_tensor.mlir" "promote_i1_to_i8.mlir" "promote_tensor_loads.mlir" + "simplify_variable_accesses.mlir" "strip_and_splat_constant_variables.mlir" "transformation.mlir" DATA
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/dispatch_linalg_on_tensors.mlir b/iree/compiler/Dialect/Flow/Transforms/test/dispatch_linalg_on_tensors.mlir index 5e30bed..3a7999d 100644 --- a/iree/compiler/Dialect/Flow/Transforms/test/dispatch_linalg_on_tensors.mlir +++ b/iree/compiler/Dialect/Flow/Transforms/test/dispatch_linalg_on_tensors.mlir
@@ -910,7 +910,7 @@ // CHECK-SAME: {__internal_linalg_transform__ = "workgroup"} // CHECK-SAME: ins(%[[UPDATE_TILE]], %[[INDICES_TILE]] : tensor<?x?xf32>, tensor<?x1xi32>) // CHECK-SAME: outs(%[[ORIGINAL]] : tensor<?x?xf32>) -// CHECK: flow.dispatch.tensor.store %[[RESULT_TILE]], %[[ARG5]], offsets = [0, 0] +// CHECK: flow.dispatch.tensor.store %[[RESULT_TILE]], %[[ARG5]], offsets = [], sizes = [], strides = [] // CHECK: return %[[RESULT]] : tensor<?x?xf32> // ----- @@ -1007,3 +1007,33 @@ // CHECK: flow.return // CHECK: } // CHECK: return %[[RESULT]]#0, %[[RESULT]]#1 + +// ----- + +func @scatter_static(%arg0 : tensor<4xi32>, %arg1 : tensor<4x1xi32>, %arg2 : tensor<8xi32>) + -> tensor<8xi32>{ + %cst = constant dense<[0, 9, 0, 10, 11, 0, 0, 12]> : tensor<8xi32> + %cst_0 = constant dense<[9, 10, 11, 12]> : tensor<4xi32> + %cst_1 = constant dense<[[1], [3], [4], [7]]> : tensor<4x1xi32> + %cst_2 = constant dense<0> : tensor<8xi32> + %0 = linalg_ext.scatter + ins(%arg0, %arg1 : tensor<4xi32>, tensor<4x1xi32>) + outs(%arg2 : tensor<8xi32>) { + ^bb0(%arg3: i32, %arg4: i32): // no predecessors + linalg_ext.yield %arg3 : i32 + } -> tensor<8xi32> + return %0 : tensor<8xi32> +} +// CHECK: func @scatter_static +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<4xi32> +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<4x1xi32> +// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<8xi32> +// CHECK: %[[RESULT:.+]] = flow.dispatch.workgroups +// CHECK-NEXT: %[[ARG3:[a-zA-Z0-9_]+]]: !flow.dispatch.tensor<readonly:4xi32> +// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: !flow.dispatch.tensor<readonly:4x1xi32> +// CHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: !flow.dispatch.tensor<readwrite:8xi32> +// CHECK: scf.for %[[IV:.+]] = %{{.+}} to %{{.+}} step %{{.+}} { +// CHECK: %[[SCATTER_TILE:.+]] = linalg_ext.scatter +// CHECK: flow.dispatch.tensor.store %[[SCATTER_TILE]], %[[ARG5]], offsets = [], sizes = [], strides = [] +// CHECK-NEXT: } +// CHECK: return %[[RESULT]]
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/simplify_variable_accesses.mlir b/iree/compiler/Dialect/Flow/Transforms/test/simplify_variable_accesses.mlir new file mode 100644 index 0000000..927d193 --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/test/simplify_variable_accesses.mlir
@@ -0,0 +1,136 @@ +// RUN: iree-opt -split-input-file -iree-flow-simplify-variable-accesses %s | IreeFileCheck %s + +flow.variable @varA dense<1> : tensor<2xi32> attributes {sym_visibility = "private"} +flow.variable @varB dense<3> : tensor<2x4xi32> attributes {sym_visibility = "private"} + +// CHECK-LABEL: @constants() +func @constants() { + // CHECK-NEXT: %[[VAR_A:.+]] = flow.variable.load @varA : tensor<2xi32> + // CHECK-NEXT: %[[VAR_B:.+]] = flow.variable.load @varB : tensor<2x4xi32> + // CHECK-NEXT: constant 10 + %w = constant 10 : index + %varA = flow.variable.load @varA : tensor<2xi32> + // CHECK-NEXT: %[[T:.+]] = flow.dispatch @ex::@dispatch0{{.+}}(%[[VAR_A]]) + %d0 = flow.dispatch @ex::@dispatch0[%w](%varA) : (tensor<2xi32>) -> tensor<2xi32> + %varB = flow.variable.load @varB : tensor<2x4xi32> + // CHECK-NEXT: flow.dispatch @ex::@dispatch1{{.+}}(%[[T]], %[[VAR_B]]) + %d1 = flow.dispatch @ex::@dispatch1[%w](%d0, %varB) : (tensor<2xi32>, tensor<2x4xi32>) -> tensor<2xi32> + return +} + +// ----- + +flow.variable @varA 1 : i32 attributes {sym_visibility = "private"} +flow.variable @varB 2 : i32 attributes {sym_visibility = "private"} + +// CHECK-LABEL: @constants_in_cfg +func @constants_in_cfg(%start: i32, %bound: i32) -> i32 { + // CHECK-NEXT: %[[VAR_A:.+]] = flow.variable.load @varA : i32 + // CHECK-NEXT: %[[VAR_B:.+]] = flow.variable.load @varB : i32 + // CHECK-NEXT: br ^bb1 + br ^bb1(%start : i32) +// CHECK: ^bb1(%[[BB1_ARG:.+]]: i32): +^bb1(%2: i32): + %cmp = cmpi slt, %2, %bound : i32 + cond_br %cmp, ^bb2(%2 : i32), ^bb3(%2 : i32) +// CHECK: ^bb2(%[[BB2_ARG:.+]]: i32): +^bb2(%5: i32): + %6 = flow.variable.load @varA : i32 + // CHECK-NEXT: = addi %[[BB2_ARG]], %[[VAR_A]] : i32 + %7 = addi %5, %6 : i32 + br ^bb1(%7 : i32) +// CHECK: ^bb3(%[[BB3_ARG:.+]]: i32): +^bb3(%8: i32): + %9 = flow.variable.load @varA : i32 + // CHECK-NEXT: %[[T0:.+]] = muli %[[BB3_ARG]], %[[VAR_A]] : i32 + %10 = muli %8, %9 : i32 + %11 = flow.variable.load @varB : i32 + // CHECK-NEXT: %[[T1:.+]] = subi %[[T0]], %[[VAR_B]] + %12 = subi %10, %11 : i32 + // CHECK-NEXT: return %[[T1]] + return %12 : i32 +} + +// ----- + +flow.variable @varA mutable dense<1> : tensor<2xi32> attributes {sym_visibility = "private"} +flow.variable @varB dense<3> : tensor<2x4xi32> attributes {sym_visibility = "private"} + +// CHECK-LABEL: @mixed_mutability +func @mixed_mutability() { + // CHECK-DAG: %[[VAR_A:.+]] = flow.variable.load @varA : tensor<2xi32> + // CHECK-DAG: %[[VAR_B:.+]] = flow.variable.load @varB : tensor<2x4xi32> + // CHECK-NEXT: constant 10 + %w = constant 10 : index + %varA = flow.variable.load @varA : tensor<2xi32> + // CHECK-NEXT: %[[T0:.+]] = flow.dispatch @ex::@dispatch0{{.+}}(%[[VAR_A]]) + %d0 = flow.dispatch @ex::@dispatch0[%w](%varA) : (tensor<2xi32>) -> tensor<2xi32> + %varB = flow.variable.load @varB : tensor<2x4xi32> + // CHECK-NEXT: %[[T1:.+]] = flow.dispatch @ex::@dispatch1{{.+}}(%[[T0]], %[[VAR_B]]) + %d1 = flow.dispatch @ex::@dispatch1[%w](%d0, %varB) : (tensor<2xi32>, tensor<2x4xi32>) -> tensor<2xi32> + // CHECK-NEXT: flow.variable.store %[[T1]], @varA : tensor<2xi32> + flow.variable.store %d1, @varA : tensor<2xi32> + return +} + +// ----- + +flow.variable @varA mutable dense<1> : tensor<2xi32> attributes {sym_visibility = "private"} + +// CHECK-LABEL: @raw +func @raw() { + // CHECK: %[[T:.+]] = flow.variable.load @varA {id = 0 + %varA_0 = flow.variable.load @varA {id = 0} : tensor<2xi32> + flow.variable.store %varA_0, @varA {id = 0} : tensor<2xi32> + %varA_1 = flow.variable.load @varA {id = 1} : tensor<2xi32> + // CHECK-NEXT: flow.variable.store %[[T]], @varA {id = 1 + flow.variable.store %varA_1, @varA {id = 1} : tensor<2xi32> + return +} + +// ----- + +flow.variable @varA mutable dense<1> : tensor<2xi32> attributes {sym_visibility = "private"} + +// CHECK-LABEL: @rar +func @rar() -> (tensor<2xi32>, tensor<2xi32>) { + // CHECK: %[[T:.+]] = flow.variable.load @varA {id = 0 + %varA_0 = flow.variable.load @varA {id = 0} : tensor<2xi32> + %varA_1 = flow.variable.load @varA {id = 1} : tensor<2xi32> + // CHECK-NEXT: return %[[T]], %[[T]] + return %varA_0, %varA_1 : tensor<2xi32>, tensor<2xi32> +} + +// ----- + +flow.variable @varA mutable dense<1> : tensor<2xi32> attributes {sym_visibility = "private"} + +// CHECK-LABEL: @waw +// CHECK-SAME: (%[[ARG0:.+]]: tensor<2xi32>, %[[ARG1:.+]]: tensor<2xi32>) +func @waw(%varA_0: tensor<2xi32>, %varA_1: tensor<2xi32>) { + flow.variable.store %varA_0, @varA : tensor<2xi32> + // CHECK-NEXT: flow.variable.store %[[ARG1]], @varA + flow.variable.store %varA_1, @varA : tensor<2xi32> + return +} + +// ----- + +flow.variable @varA mutable dense<1> : tensor<2xi32> attributes {sym_visibility = "private"} + +// CHECK-LABEL: @side_effects( +func @side_effects() { + // CHECK-NEXT: %[[T0:.+]] = flow.variable.load @varA + %varA_0 = flow.variable.load @varA : tensor<2xi32> + // CHECK-NEXT: flow.variable.store %[[T0]], @varA + flow.variable.store %varA_0, @varA : tensor<2xi32> + // CHECK-NEXT: call @other_fn() + call @other_fn() : () -> () + // CHECK-NEXT: %[[T1:.+]] = flow.variable.load @varA + %varA_1 = flow.variable.load @varA : tensor<2xi32> + // CHECK-NEXT: flow.variable.store %[[T1]], @varA + flow.variable.store %varA_1, @varA : tensor<2xi32> + return +} + +func private @other_fn()
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/BUILD b/iree/compiler/Dialect/HAL/Target/LLVM/BUILD index 01c1c13..3359b98 100644 --- a/iree/compiler/Dialect/HAL/Target/LLVM/BUILD +++ b/iree/compiler/Dialect/HAL/Target/LLVM/BUILD
@@ -41,12 +41,14 @@ "//iree/compiler/Codegen/LLVMCPU", "//iree/compiler/Codegen/Utils", "//iree/compiler/Dialect/HAL/Target", + "//iree/compiler/Dialect/HAL/Target/LLVM/librt", "//iree/compiler/Utils", "//iree/schemas:dylib_executable_def_c_fbs", "@llvm-project//llvm:AArch64AsmParser", "@llvm-project//llvm:AArch64CodeGen", "@llvm-project//llvm:ARMAsmParser", "@llvm-project//llvm:ARMCodeGen", + "@llvm-project//llvm:BitReader", "@llvm-project//llvm:BitWriter", "@llvm-project//llvm:Core", "@llvm-project//llvm:RISCVAsmParser",
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/CMakeLists.txt b/iree/compiler/Dialect/HAL/Target/LLVM/CMakeLists.txt index f12301b..29c9514 100644 --- a/iree/compiler/Dialect/HAL/Target/LLVM/CMakeLists.txt +++ b/iree/compiler/Dialect/HAL/Target/LLVM/CMakeLists.txt
@@ -32,6 +32,7 @@ LLVMAArch64CodeGen LLVMARMAsmParser LLVMARMCodeGen + LLVMBitReader LLVMBitWriter LLVMCore LLVMRISCVAsmParser @@ -50,6 +51,7 @@ iree::compiler::Codegen::PassHeaders iree::compiler::Codegen::Utils iree::compiler::Dialect::HAL::Target + iree::compiler::Dialect::HAL::Target::LLVM::librt iree::compiler::Utils iree::schemas::dylib_executable_def_c_fbs PUBLIC
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp index f80e47c..fdefab2 100644 --- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp +++ b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp
@@ -13,9 +13,11 @@ #include "iree/compiler/Dialect/HAL/Target/LLVM/LibraryBuilder.h" #include "iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.h" #include "iree/compiler/Dialect/HAL/Target/LLVM/StaticLibraryGenerator.h" +#include "iree/compiler/Dialect/HAL/Target/LLVM/librt/librt.h" #include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h" #include "iree/compiler/Utils/FlatbufferUtils.h" #include "iree/schemas/dylib_executable_def_builder.h" +#include "llvm/Bitcode/BitcodeReader.h" #include "llvm/Bitcode/BitcodeWriter.h" #include "llvm/IR/LLVMContext.h" #include "llvm/IR/Module.h" @@ -30,11 +32,10 @@ namespace IREE { namespace HAL { -namespace { +static constexpr char kQueryFunctionName[] = + "iree_hal_executable_library_query"; -constexpr char kQueryFunctionName[] = "iree_hal_executable_library_query"; - -llvm::Optional<FileLineColLoc> findFirstFileLoc(Location baseLoc) { +static llvm::Optional<FileLineColLoc> findFirstFileLoc(Location baseLoc) { if (auto loc = baseLoc.dyn_cast<FusedLoc>()) { for (auto &childLoc : loc.getLocations()) { auto childResult = findFirstFileLoc(childLoc); @@ -46,7 +47,7 @@ return llvm::None; } -std::string guessModuleName(mlir::ModuleOp moduleOp) { +static std::string guessModuleName(mlir::ModuleOp moduleOp) { std::string moduleName = moduleOp.getName().hasValue() ? moduleOp.getName().getValue().str() : ""; if (!moduleName.empty()) return moduleName; @@ -58,8 +59,6 @@ } } -} // namespace - class LLVMAOTTargetBackend final : public TargetBackend { public: explicit LLVMAOTTargetBackend(LLVMTargetOptions options) @@ -284,8 +283,10 @@ << options_.targetTriple << "'"; } - // Emit object files. - SmallVector<Artifact, 4> objectFiles; + SmallVector<Artifact> objectFiles; + + // Emit the base object file containing the bulk of our code. + // This must come first such that we have the proper library linking order. { // NOTE: today we just use a single object file, however if we wanted to // scale code generation and linking we'd want to generate one per @@ -306,6 +307,16 @@ objectFiles.push_back(std::move(objectFile)); } + // Optionally append additional object files that provide functionality that + // may otherwise have been runtime-dynamic (like libc/libm calls). + // For now we only do this for embedded uses. + if (options_.linkEmbedded) { + if (failed(buildLibraryObjects(variantOp.getLoc(), targetMachine.get(), + objectFiles, context))) { + return variantOp.emitError() << "failed generating library objects"; + } + } + // If we are keeping artifacts then let's also add the bitcode for easier // debugging (vs just the binary object file). if (options_.keepLinkerArtifacts) { @@ -329,7 +340,6 @@ // Copy the static object file to the specified output along with // generated header file. const std::string &libraryPath = options_.staticLibraryOutput; - const auto library_name = objectFiles[0].path; if (!outputStaticLibrary(libraryName, queryFunctionName, libraryPath, objectFiles[0].path)) { return variantOp.emitError() << "static library generation failed"; @@ -506,6 +516,76 @@ return IREE::HAL::ExecutableTargetAttr::get(context, "llvm", format); } + static void overridePlatformGlobal(llvm::Module &module, StringRef globalName, + uint32_t newValue) { + // NOTE: the global will not be defined if it is not used in the module. + auto *globalValue = module.getNamedGlobal(globalName); + if (!globalValue) return; + globalValue->setLinkage(llvm::GlobalValue::PrivateLinkage); + globalValue->setDSOLocal(true); + globalValue->setConstant(true); + globalValue->setInitializer(llvm::ConstantInt::get( + globalValue->getValueType(), APInt(32, newValue))); + } + + // Builds an object file for the librt embedded runtime library. + // This is done per link operation so that we can match the precise target + // configuration. Since we (mostly) link once per user-level compilation + // this is fine today. If in the future we invoke the compiler for thousands + // of modules we'd want to (carefully) cache this. + LogicalResult buildLibraryObjects(Location loc, + llvm::TargetMachine *targetMachine, + SmallVector<Artifact> &objectFiles, + llvm::LLVMContext &context) { + assert(!objectFiles.empty() && "libraries must come after the base object"); + + // Load the generic bitcode file contents. + llvm::MemoryBufferRef bitcodeBufferRef( + llvm::StringRef(iree_compiler_librt_create()->data, + iree_compiler_librt_create()->size), + "librt.bc"); + auto bitcodeModuleValue = llvm::parseBitcodeFile(bitcodeBufferRef, context); + if (!bitcodeModuleValue) { + return mlir::emitError(loc) + << "failed to parse librt bitcode: " + << llvm::toString(bitcodeModuleValue.takeError()); + } + auto bitcodeModule = std::move(bitcodeModuleValue.get()); + bitcodeModule->setDataLayout(targetMachine->createDataLayout()); + bitcodeModule->setTargetTriple(targetMachine->getTargetTriple().str()); + + // Inject target-specific flags. + // TODO(benvanik): move this entire function to another file that can do + // more complex logic cleanly. This is just an example. + overridePlatformGlobal(*bitcodeModule, "librt_platform_example_flag", 0u); + + // Run the LLVM passes to optimize it for the current target. + if (failed(runLLVMIRPasses(options_, targetMachine, bitcodeModule.get()))) { + return mlir::emitError(loc) + << "failed to run librt LLVM-IR opt passes targeting '" + << options_.targetTriple << "'"; + } + + // Emit an object file we can pass to the linker. + std::string objectData; + if (failed(runEmitObjFilePasses(targetMachine, bitcodeModule.get(), + &objectData))) { + return mlir::emitError(loc) + << "failed to compile librt LLVM-IR module to an object file"; + } + + // Write the object file to disk with a similar name to the base file. + auto objectFile = + Artifact::createVariant(objectFiles.front().path, ".librt.o"); + auto &os = objectFile.outputFile->os(); + os << objectData; + os.flush(); + os.close(); + objectFiles.push_back(std::move(objectFile)); + + return success(); + } + LLVMTargetOptions options_; };
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/internal/EmbeddedLinkerTool.cpp b/iree/compiler/Dialect/HAL/Target/LLVM/internal/EmbeddedLinkerTool.cpp index f2847c2..cda219e 100644 --- a/iree/compiler/Dialect/HAL/Target/LLVM/internal/EmbeddedLinkerTool.cpp +++ b/iree/compiler/Dialect/HAL/Target/LLVM/internal/EmbeddedLinkerTool.cpp
@@ -80,7 +80,7 @@ // Create the shared object name; if we only have a single input object we // can just reuse that. - if (objectFiles.size() == 1) { + if (!objectFiles.empty()) { artifacts.libraryFile = Artifact::createVariant(objectFiles.front().path, "so"); } else {
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/librt/BUILD b/iree/compiler/Dialect/HAL/Target/LLVM/librt/BUILD new file mode 100644 index 0000000..f9f6feb --- /dev/null +++ b/iree/compiler/Dialect/HAL/Target/LLVM/librt/BUILD
@@ -0,0 +1,31 @@ +# Copyright 2021 The IREE Authors +# +# 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 + +load("//build_tools/embed_data:build_defs.bzl", "c_embed_data") +load("//iree:build_defs.oss.bzl", "iree_cmake_extra_content") + +package( + default_visibility = ["//visibility:public"], + features = ["layering_check"], + licenses = ["notice"], # Apache 2.0 +) + +iree_cmake_extra_content( + content = """ +if(NOT "${IREE_TARGET_BACKEND_DYLIB-LLVM-AOT}" AND NOT "${IREE_TARGET_BACKEND_WASM-LLVM-AOT}") + return() +endif() +""", +) + +c_embed_data( + name = "librt", + srcs = ["bin/librt.bc"], + c_file_output = "librt.c", + flatten = True, + h_file_output = "librt.h", + identifier = "iree_compiler_librt", +)
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/librt/CMakeLists.txt b/iree/compiler/Dialect/HAL/Target/LLVM/librt/CMakeLists.txt new file mode 100644 index 0000000..c6dabe0 --- /dev/null +++ b/iree/compiler/Dialect/HAL/Target/LLVM/librt/CMakeLists.txt
@@ -0,0 +1,32 @@ +################################################################################ +# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from # +# iree/compiler/Dialect/HAL/Target/LLVM/librt/BUILD # +# # +# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary # +# CMake-only content. # +# # +# To disable autogeneration for this file entirely, delete this header. # +################################################################################ + +if(NOT "${IREE_TARGET_BACKEND_DYLIB-LLVM-AOT}" AND NOT "${IREE_TARGET_BACKEND_WASM-LLVM-AOT}") + return() +endif() + +iree_add_all_subdirs() + +iree_c_embed_data( + NAME + librt + SRCS + "bin/librt.bc" + C_FILE_OUTPUT + "librt.c" + H_FILE_OUTPUT + "librt.h" + IDENTIFIER + "iree_compiler_librt" + FLATTEN + PUBLIC +) + +### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/librt/bin/librt.bc b/iree/compiler/Dialect/HAL/Target/LLVM/librt/bin/librt.bc new file mode 100644 index 0000000..a484f31 --- /dev/null +++ b/iree/compiler/Dialect/HAL/Target/LLVM/librt/bin/librt.bc Binary files differ
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/librt/build.sh b/iree/compiler/Dialect/HAL/Target/LLVM/librt/build.sh new file mode 100644 index 0000000..c671e4d --- /dev/null +++ b/iree/compiler/Dialect/HAL/Target/LLVM/librt/build.sh
@@ -0,0 +1,45 @@ +# Copyright 2021 The IREE Authors +# +# 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 -e + +SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )" +OUT="${SCRIPT_DIR}/bin" +SRC="${SCRIPT_DIR}/src" +LL_FILE="${OUT}/librt.ll" +BC_FILE="${OUT}/librt.bc" + +# Generate an LLVM IR assembly listing so we can easily read the file. +# This is not checked in or used by the compiler. +clang \ + -target wasm32 \ + -std=c17 \ + -O2 \ + -Xclang -disable-llvm-passes \ + -fno-ident \ + -fvisibility=hidden \ + -nostdinc \ + -g0 \ + -S \ + -emit-llvm \ + -fno-verbose-asm \ + -fdiscard-value-names \ + -o "${LL_FILE}" \ + -c \ + "${SRC}/libm.c" + +# Clang adds a bunch of bad attributes and host-specific information that we +# don't want (so we get at least somewhat deterministic builds). +sed -i 's/^;.*$//' ${LL_FILE} +sed -i 's/^source_filename.*$//' ${LL_FILE} +sed -i 's/^target datalayout.*$//' ${LL_FILE} +sed -i 's/^target triple.*$//' ${LL_FILE} +sed -i 's/^\(attributes #[0-9]* = {\).*$/\1 inlinehint }/' ${LL_FILE} + +# Generate a binary bitcode file embedded into the compiler binary. +# NOTE: we do this from stdin so that the filename on the user's system is not +# embedded in the bitcode file (making it non-deterministic). +cat ${LL_FILE} | llvm-as -o=${BC_FILE}
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/librt/src/libm.c b/iree/compiler/Dialect/HAL/Target/LLVM/librt/src/libm.c new file mode 100644 index 0000000..bdbe0ae --- /dev/null +++ b/iree/compiler/Dialect/HAL/Target/LLVM/librt/src/libm.c
@@ -0,0 +1,13 @@ +// Copyright 2021 The IREE Authors +// +// 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 + +#include "libm.h" + +// https://en.cppreference.com/w/c/numeric/math/fma +LIBRT_EXPORT float fmaf(float x, float y, float z) { + // TODO(*): a real implementation :) + return (x * y) + z; +}
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/librt/src/libm.h b/iree/compiler/Dialect/HAL/Target/LLVM/librt/src/libm.h new file mode 100644 index 0000000..10dfe76 --- /dev/null +++ b/iree/compiler/Dialect/HAL/Target/LLVM/librt/src/libm.h
@@ -0,0 +1,15 @@ +// Copyright 2021 The IREE Authors +// +// 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 + +#ifndef IREE_COMPILER_DIALECT_HAL_TARGET_LLVM_LIBRT_SRC_LIBM_H_ +#define IREE_COMPILER_DIALECT_HAL_TARGET_LLVM_LIBRT_SRC_LIBM_H_ + +#include "librt.h" + +// https://en.cppreference.com/w/c/numeric/math/fma +LIBRT_EXPORT float fmaf(float x, float y, float z); + +#endif // IREE_COMPILER_DIALECT_HAL_TARGET_LLVM_LIBRT_SRC_LIBM_H_
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/librt/src/librt.h b/iree/compiler/Dialect/HAL/Target/LLVM/librt/src/librt.h new file mode 100644 index 0000000..0af8b21 --- /dev/null +++ b/iree/compiler/Dialect/HAL/Target/LLVM/librt/src/librt.h
@@ -0,0 +1,75 @@ +// Copyright 2021 The IREE Authors +// +// 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 + +//===----------------------------------------------------------------------===// +// A simplified libc/libm-alike that is designed to compile to portable LLVM IR. +//===----------------------------------------------------------------------===// +// This library is focused on supporting the subset of LLVM's RuntimeLibcalls +// that we need in our embedded executable binaries. This means that things like +// printf, malloc, etc are excluded. +// +// See the full list of possible functions here: +// third_party/llvm-project/llvm/include/llvm/IR/RuntimeLibcalls.def +// +// Code here must not use any system headers - as almost all pull in bits/ and +// various other target-dependent definitions that make the resulting IR +// non-portable. This means there is no size_t, etc. Any definitions that may +// come from an std* file must be redefined here with care. +// +// Code must also not use any mutable global or thread-local state ala +// errno/rounding modes/etc. Each of the functions in the library will be called +// concurrently from multiple threads and from multiple source modules. There +// must be no mutable static values anywhere. +// +// Avoid #ifdef entirely: they indicate a leakage of host build configuration +// into what is supposed to be a portable module. Anything that requires +// target-specific conditional logic must be implemented via an extern that +// can be substituted by the IREE compiler when producing the final +// target-specific module. + +//===----------------------------------------------------------------------===// +// Attributes and metadata +//===----------------------------------------------------------------------===// + +// Tagged on functions that are part of the public API. +#define LIBRT_EXPORT __attribute__((visibility("hidden"))) + +//===----------------------------------------------------------------------===// +// stdint.h +//===----------------------------------------------------------------------===// +// https://pubs.opengroup.org/onlinepubs/009604599/basedefs/stdint.h.html +// NOTE: no size_t/ptrdiff_t/etc (as they are target dependent). + +typedef signed char int8_t; +typedef short int16_t; +typedef int int32_t; +typedef long long int64_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; + +#define INT8_MIN (-127i8 - 1) +#define INT16_MIN (-32767i16 - 1) +#define INT32_MIN (-2147483647i32 - 1) +#define INT64_MIN (-9223372036854775807i64 - 1) +#define INT8_MAX 127i8 +#define INT16_MAX 32767i16 +#define INT32_MAX 2147483647i32 +#define INT64_MAX 9223372036854775807i64 +#define UINT8_MAX 0xffui8 +#define UINT16_MAX 0xffffui16 +#define UINT32_MAX 0xffffffffui32 +#define UINT64_MAX 0xffffffffffffffffui64 + +//===----------------------------------------------------------------------===// +// Target-specific queries +//===----------------------------------------------------------------------===// +// These are substituted with values from the compiler and must not be specified +// here in C before we generate the IR. + +// Do not use: here as an example. Remove once we have any other flag. +extern int librt_platform_example_flag;
diff --git a/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp b/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp index c7d61c8..3e47a35 100644 --- a/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp +++ b/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp
@@ -203,11 +203,11 @@ return {Range{zero, ub, one}}; } -Operation *ScatterOp::getTiledImplementation( - OpBuilder &builder, ValueRange outputs, ArrayRef<OpFoldResult> offsets, - ArrayRef<OpFoldResult> sizes, - SmallVectorImpl<SmallVector<OpFoldResult, 4>> &resultOffsets, - SmallVectorImpl<SmallVector<OpFoldResult, 4>> &resultSizes) { +Operation *ScatterOp::getTiledImplementation(OpBuilder &builder, + ValueRange outputs, + ArrayRef<OpFoldResult> offsets, + ArrayRef<OpFoldResult> sizes, + SmallVectorImpl<Value> &results) { assert(outputs.size() == 1 && offsets.size() == 1 && sizes.size() == 1); Location loc = getLoc(); auto zeroAttr = builder.getI64IntegerAttr(0); @@ -241,20 +241,18 @@ indicesSizes, indicesStrides); assert(tiledIndices && "failed to get slice of indices"); - resultOffsets.resize(1); - resultOffsets[0].resize(updateRank, zeroAttr); - resultSizes.resize(1); - resultSizes[0].resize(updateRank); - for (auto dim : llvm::seq<int64_t>(0, updateRank)) { - resultSizes[0][dim] = getDim(builder, loc, original(), dim); - } SmallVector<Type> resultTypes; if (getNumResults()) { resultTypes.push_back(getResultTypes()[0]); } - return cast<LinalgExtOp>(getOperation()) - .clone(builder, loc, resultTypes, - ValueRange{tiledUpdate, tiledIndices, outputs[0]}); + Operation *tiledScatterOp = + cast<LinalgExtOp>(getOperation()) + .clone(builder, loc, resultTypes, + ValueRange{tiledUpdate, tiledIndices, outputs[0]}); + if (getNumResults()) { + results.push_back(tiledScatterOp->getResult(0)); + } + return tiledScatterOp; } void ScatterOp::generateScalarUpdateLoopBody(OpBuilder &b, Location loc, @@ -466,11 +464,11 @@ return partitionableLoops; } -Operation *SortOp::getTiledImplementation( - OpBuilder &builder, ValueRange outputs, ArrayRef<OpFoldResult> offsets, - ArrayRef<OpFoldResult> sizes, - SmallVectorImpl<SmallVector<OpFoldResult, 4>> &resultOffsets, - SmallVectorImpl<SmallVector<OpFoldResult, 4>> &resultSizes) { +Operation *SortOp::getTiledImplementation(OpBuilder &builder, + ValueRange outputs, + ArrayRef<OpFoldResult> offsets, + ArrayRef<OpFoldResult> sizes, + SmallVectorImpl<Value> &results) { assert(outputs.size() == this->outputs().size()); int64_t rank = getOperandRank(); assert(offsets.size() == static_cast<size_t>(rank) && @@ -479,22 +477,24 @@ SmallVector<OpFoldResult> strides(rank, oneAttr); Location loc = getLoc(); SmallVector<Value> tiledOperands(outputs.size()); - resultOffsets.resize(outputs.size()); - resultSizes.resize(outputs.size()); for (auto en : llvm::enumerate(outputs)) { tiledOperands[en.index()] = getSlice(builder, getLoc(), en.value(), offsets, sizes, strides); assert(tiledOperands[en.index()] && "failed to get slice of operand"); - resultOffsets[en.index()].assign(offsets.begin(), offsets.end()); - resultSizes[en.index()].assign(sizes.begin(), sizes.end()); } SmallVector<Type, 4> resultTypes; if (getNumResults()) { resultTypes = llvm::to_vector<4>( llvm::map_range(tiledOperands, [&](Value v) { return v.getType(); })); } - return cast<LinalgExtOp>(getOperation()) - .clone(builder, loc, resultTypes, tiledOperands); + Operation *tiledSortOp = cast<LinalgExtOp>(getOperation()) + .clone(builder, loc, resultTypes, tiledOperands); + for (auto result : llvm::enumerate(tiledSortOp->getResults())) { + auto insertSliceOp = builder.create<tensor::InsertSliceOp>( + loc, result.value(), outputs[result.index()], offsets, sizes, strides); + results.push_back(insertSliceOp.getResult()); + } + return tiledSortOp; } LogicalResult SortOp::generateScalarImplementation(OpBuilder &b, Location loc,
diff --git a/iree/compiler/Dialect/LinalgExt/IR/TiledOpInterface.td b/iree/compiler/Dialect/LinalgExt/IR/TiledOpInterface.td index 232c6aa..337b273 100644 --- a/iree/compiler/Dialect/LinalgExt/IR/TiledOpInterface.td +++ b/iree/compiler/Dialect/LinalgExt/IR/TiledOpInterface.td
@@ -85,6 +85,10 @@ /*desc=*/[{ Generates a tiled version of the operation given the tile size for the loops. + + Returns the tiled operation generated. If the operation has + tensor semantics then the result of the tiled values are to + be inserted into the `outputs` and return in `results`. }], /*retType=*/"Operation *", /*methodName=*/"getTiledImplementation", @@ -93,8 +97,7 @@ "ValueRange ":$outputs, "ArrayRef<OpFoldResult> ":$offsets, "ArrayRef<OpFoldResult> ":$sizes, - "SmallVectorImpl<SmallVector<OpFoldResult, 4>> &":$resultOffsets, - "SmallVectorImpl<SmallVector<OpFoldResult, 4>> &":$resultSizes), + "SmallVectorImpl<Value> &":$results), /*methodBody=*/"", /*defaultImplementation=*/[{ return nullptr;
diff --git a/iree/compiler/Dialect/LinalgExt/Transforms/Tiling.cpp b/iree/compiler/Dialect/LinalgExt/Transforms/Tiling.cpp index 829c306..bf6f1be 100644 --- a/iree/compiler/Dialect/LinalgExt/Transforms/Tiling.cpp +++ b/iree/compiler/Dialect/LinalgExt/Transforms/Tiling.cpp
@@ -132,36 +132,13 @@ // If this is the innermost loop, then generated the tiled implementation of // the op by invoking the TiledOpInterface methods. if (loopDepth == tileSizes.size()) { - SmallVector<SmallVector<OpFoldResult, 4>> resultOffsets; - SmallVector<SmallVector<OpFoldResult, 4>> resultSizes; - Operation *tiledOp = tilableOp.getTiledImplementation( - builder, outputs, offsets, tileSizes, resultOffsets, resultSizes); - if (!tiledOp) { + TiledOp ret; + ret.op = tilableOp.getTiledImplementation(builder, outputs, offsets, + tileSizes, ret.results); + if (!ret.op) { return static_cast<LogicalResult>( tilableOp.emitOpError("failed to get tiled implementation")); } - assert(tiledOp->getNumResults() == 0 || - (resultOffsets.size() == tiledOp->getNumResults())); - TiledOp ret; - ret.op = tiledOp; - - // If the operation has results, then the result of the tiled operation is - // to be inserted into the `initValues` and returned. - if (tiledOp->getNumResults()) { - SmallVector<Value> results; - auto oneAttr = builder.getI64IntegerAttr(1); - results.reserve(tiledOp->getNumResults()); - for (auto en : llvm::enumerate(tiledOp->getResults())) { - Value result = en.value(); - ArrayRef<OpFoldResult> offsets(resultOffsets[en.index()]); - ArrayRef<OpFoldResult> sizes(resultSizes[en.index()]); - SmallVector<OpFoldResult> strides(offsets.size(), oneAttr); - Value insert = builder.create<tensor::InsertSliceOp>( - loc, result, outputs[en.index()], offsets, sizes, strides); - results.push_back(insert); - } - std::swap(ret.results, results); - } return ret; } @@ -318,15 +295,13 @@ return loopBounds; } - Operation *getTiledImplementation( - Operation *op, OpBuilder &b, ValueRange outputs, - ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes, - SmallVectorImpl<SmallVector<OpFoldResult, 4>> &resultOffsets, - SmallVectorImpl<SmallVector<OpFoldResult, 4>> &resultSizes) const { - // Compute a subtensor of the source based on the offsets. + Operation *getTiledImplementation(Operation *op, OpBuilder &b, + ValueRange outputs, + ArrayRef<OpFoldResult> offsets, + ArrayRef<OpFoldResult> sizes, + SmallVectorImpl<Value> &results) const { auto insertOp = cast<tensor::InsertSliceOp>(op); - auto opOffsets = insertOp.getMixedOffsets(); - auto opSizes = insertOp.getMixedSizes(); + // Compute a subtensor of the source based on the offsets. auto opStrides = insertOp.getMixedStrides(); if (!llvm::all_of(opStrides, [&](OpFoldResult valueOrAttr) { return isValue(valueOrAttr, 1); @@ -334,23 +309,22 @@ op->emitOpError("unable to tile operation with non-unit stride"); return nullptr; } - // The operation returned is just a tensor.extract_slice of the source with - // the given offsets, sizes and strides. Setting the correct result offset - // will make the sure the tiling algorithm will insert this slice into the - // correct place in the destination. - // The result offset is just the offset passed in plus the offset specified - // in the op (since all strides are checked to be 1). + Location loc = insertOp.getLoc(); + auto oneAttr = b.getI64IntegerAttr(1); + SmallVector<OpFoldResult> strides(offsets.size(), oneAttr); + auto extractSliceOp = b.create<tensor::ExtractSliceOp>( + loc, insertOp.source(), offsets, sizes, strides); + + // The offsets for the insert is based on the op offsets plus the offsets of + // the loops passed in. + auto opOffsets = insertOp.getMixedOffsets(); + auto opSizes = insertOp.getMixedSizes(); unsigned offsetIndex = 0; ArrayRef<int64_t> sourceShape = insertOp.getSourceType().getShape(); int64_t destRank = insertOp.getType().getRank(); - resultOffsets.resize(1); - resultOffsets[0].resize(destRank); - resultSizes.resize(1); - resultSizes[0].resize(destRank); - Location loc = insertOp.getLoc(); + SmallVector<OpFoldResult> resultOffsets(destRank); + SmallVector<OpFoldResult> resultSizes(destRank); auto zeroAttr = b.getI64IntegerAttr(0); - auto oneAttr = b.getI64IntegerAttr(1); - SmallVector<OpFoldResult> strides(offsets.size(), oneAttr); for (auto opOffset : llvm::enumerate(opOffsets)) { // Check for rank-reducing by checking that // 1) The corresponding opSize value is 1 @@ -358,29 +332,33 @@ // Then the opOffset is for the rank-reduced dimension. Skip. unsigned opOffsetIndex = opOffset.index(); if (isValue(opSizes[opOffsetIndex], 1) && sourceShape[offsetIndex] != 1) { - resultOffsets[0][opOffsetIndex] = zeroAttr; - resultSizes[0][opOffsetIndex] = oneAttr; + resultOffsets[opOffsetIndex] = zeroAttr; + resultSizes[opOffsetIndex] = oneAttr; continue; } OpFoldResult opOffsetVal = opOffset.value(); OpFoldResult offset = offsets[offsetIndex]; if (opOffsetVal.is<Attribute>() && offset.is<Attribute>()) { - resultOffsets[0][opOffsetIndex] = b.getI64IntegerAttr( + resultOffsets[opOffsetIndex] = b.getI64IntegerAttr( *getConstantValue(opOffsetVal) + *getConstantValue(offset)); } else { AffineMap map = AffineMap::get( 1, 1, {b.getAffineDimExpr(0) + b.getAffineSymbolExpr(0)}); - resultOffsets[0][opOffsetIndex] = + resultOffsets[opOffsetIndex] = b.create<AffineApplyOp>(loc, map, ValueRange{getValue(b, loc, offset), getValue(b, loc, opOffsetVal)}) .getResult(); } - resultSizes[0][opOffsetIndex] = sizes[offsetIndex]; + resultSizes[opOffsetIndex] = sizes[offsetIndex]; offsetIndex++; } - return b.create<tensor::ExtractSliceOp>(loc, insertOp.source(), offsets, - sizes, strides); + SmallVector<OpFoldResult> resultStrides(destRank, oneAttr); + auto tiledInsertOp = b.create<tensor::InsertSliceOp>( + loc, extractSliceOp.result(), outputs[0], resultOffsets, resultSizes, + resultStrides); + results.push_back(tiledInsertOp.result()); + return extractSliceOp; } }; } // namespace
diff --git a/iree/compiler/Dialect/LinalgExt/Transforms/test/tiling.mlir b/iree/compiler/Dialect/LinalgExt/Transforms/test/tiling.mlir index 5818c01..41a1cd0 100644 --- a/iree/compiler/Dialect/LinalgExt/Transforms/test/tiling.mlir +++ b/iree/compiler/Dialect/LinalgExt/Transforms/test/tiling.mlir
@@ -30,15 +30,11 @@ // CHECK-SAME: [%[[USED_TILESIZE]], %[[D1]]] // CHECK: %[[INDEX_SLICE:.+]] = tensor.extract_slice %[[INDICES]][%[[IV]], 0] // CHECK-SAME: [%[[USED_TILESIZE]], 1] -// CHECK-DAG: %[[SLICE_D0:.+]] = tensor.dim %[[ORIGINAL]], %[[C0]] -// CHECK-DAG: %[[SLICE_D1:.+]] = tensor.dim %[[ORIGINAL]], %[[C1]] // CHECK: %[[SCATTER_TILE:.+]] = linalg_ext.scatter // CHECK-SAME: __internal_linalg_transform__ = "tiling_output" // CHECK-SAME: ins(%[[UPDATE_SLICE]], %[[INDEX_SLICE]] // CHECK-SAME: outs(%[[INIT]] -// CHECK: %[[YIELD:.+]] = tensor.insert_slice %[[SCATTER_TILE]] into %[[INIT]][0, 0] -// CHECK-SAME: [%[[SLICE_D0]], %[[SLICE_D1]]] -// CHECK: scf.yield %[[YIELD]] +// CHECK: scf.yield %[[SCATTER_TILE]] // CHECK: return %[[RESULT]] // ----- @@ -114,15 +110,11 @@ // CHECK-SAME: [%[[USED_TILESIZE]], %[[D1]]] // CHECK: %[[INDEX_SLICE:.+]] = tensor.extract_slice %[[INDICES]][%[[IV]], 0] // CHECK-SAME: [%[[USED_TILESIZE]], 1] -// CHECK-DAG: %[[SLICE_D0:.+]] = tensor.dim %[[ORIGINAL]], %[[C0]] -// CHECK-DAG: %[[SLICE_D1:.+]] = tensor.dim %[[ORIGINAL]], %[[C1]] // CHECK: %[[SCATTER_TILE:.+]] = linalg_ext.scatter // CHECK-SAME: __internal_linalg_transform__ = "distribute_output" // CHECK-SAME: ins(%[[UPDATE_SLICE]], %[[INDEX_SLICE]] // CHECK-SAME: outs(%[[INIT]] -// CHECK: %[[YIELD:.+]] = tensor.insert_slice %[[SCATTER_TILE]] into %[[INIT]][0, 0] -// CHECK-SAME: [%[[SLICE_D0]], %[[SLICE_D1]]] -// CHECK: scf.yield %[[YIELD]] +// CHECK: scf.yield %[[SCATTER_TILE]] // CHECK: return %[[RESULT]] // ----- @@ -435,9 +427,9 @@ // CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: index // CHECK: %[[RESULT:.+]] = scf.for %[[IV0:[a-zA-Z0-9]+]] = // CHECK: %[[YIELD1:.+]] = scf.for %[[IV1:[a-zA-Z0-9]+]] = +// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], %[[IV1]]] // CHECK: %[[OFFSET0:.+]] = affine.apply #[[MAP2]](%[[IV0]])[%[[ARG2]]] // CHECK: %[[OFFSET1:.+]] = affine.apply #[[MAP2]](%[[IV1]])[%[[ARG3]]] -// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], %[[IV1]]] // CHECK: %[[UPDATE:.+]] = tensor.insert_slice %[[SLICE]] // CHECK-SAME: into %{{.+}}[%[[OFFSET0]], %[[OFFSET1]]] // CHECK: scf.yield %[[UPDATE]] @@ -466,9 +458,9 @@ // CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: index // CHECK: %[[RESULT:.+]] = scf.for %[[IV0:[a-zA-Z0-9]+]] = // CHECK: %[[YIELD1:.+]] = scf.for %[[IV1:[a-zA-Z0-9]+]] = +// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], %[[IV1]]] // CHECK: %[[OFFSET0:.+]] = affine.apply #[[MAP2]](%[[IV0]])[%[[ARG2]]] // CHECK: %[[OFFSET1:.+]] = affine.apply #[[MAP2]](%[[IV1]])[%[[ARG3]]] -// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], %[[IV1]]] // CHECK: %[[UPDATE:.+]] = tensor.insert_slice %[[SLICE]] // CHECK-SAME: into %{{.+}}[%[[OFFSET0]], 0, %[[OFFSET1]]] // CHECK: scf.yield %[[UPDATE]]
diff --git a/iree/compiler/InputConversion/MHLO/BUILD b/iree/compiler/InputConversion/MHLO/BUILD index c1dc464..aaf1475 100644 --- a/iree/compiler/InputConversion/MHLO/BUILD +++ b/iree/compiler/InputConversion/MHLO/BUILD
@@ -46,10 +46,10 @@ name = "MHLO", srcs = [ "BroadcastingToLinalgPatterns.cpp", - "ConvertAndDistributeMHLOToLinalgExt.cpp", "ConvertComplexToReal.cpp", "ConvertMHLOToFlow.cpp", "ConvertMHLOToFlow.h", + "ConvertMHLOToLinalgExt.cpp", "LegalizeInputTypes.cpp", "MHLOToLinalgOnTensors.cpp", "MHLOToMHLOPreprocessing.cpp",
diff --git a/iree/compiler/InputConversion/MHLO/CMakeLists.txt b/iree/compiler/InputConversion/MHLO/CMakeLists.txt index 36d2a4c..66276fc 100644 --- a/iree/compiler/InputConversion/MHLO/CMakeLists.txt +++ b/iree/compiler/InputConversion/MHLO/CMakeLists.txt
@@ -41,10 +41,10 @@ "Passes.h" SRCS "BroadcastingToLinalgPatterns.cpp" - "ConvertAndDistributeMHLOToLinalgExt.cpp" "ConvertComplexToReal.cpp" "ConvertMHLOToFlow.cpp" "ConvertMHLOToFlow.h" + "ConvertMHLOToLinalgExt.cpp" "LegalizeInputTypes.cpp" "MHLOToLinalgOnTensors.cpp" "MHLOToMHLOPreprocessing.cpp"
diff --git a/iree/compiler/InputConversion/MHLO/ConvertAndDistributeMHLOToLinalgExt.cpp b/iree/compiler/InputConversion/MHLO/ConvertMHLOToLinalgExt.cpp similarity index 68% rename from iree/compiler/InputConversion/MHLO/ConvertAndDistributeMHLOToLinalgExt.cpp rename to iree/compiler/InputConversion/MHLO/ConvertMHLOToLinalgExt.cpp index b24c76d..571b854 100644 --- a/iree/compiler/InputConversion/MHLO/ConvertAndDistributeMHLOToLinalgExt.cpp +++ b/iree/compiler/InputConversion/MHLO/ConvertMHLOToLinalgExt.cpp
@@ -47,51 +47,6 @@ namespace { //===----------------------------------------------------------------------===// -// Base classes. -//===----------------------------------------------------------------------===// - -template <typename Derived, typename OpTy> -struct ConvertToLinalgExtPattern : public OpConversionPattern<OpTy> { - using OpConversionPattern<OpTy>::OpConversionPattern; - - LogicalResult matchAndRewrite( - OpTy op, ArrayRef<Value> args, - ConversionPatternRewriter &rewriter) const final { - Value one = rewriter.create<ConstantIndexOp>(op.getLoc(), 1); - SmallVector<Value> workload(3, one); - - // Gather the dynamic dimensions for all operands. - SmallVector<Value> operandDynamicDims; - for (Value arg : args) { - if (auto rt = arg.getType().dyn_cast<RankedTensorType>()) { - for (unsigned i = 0; i < rt.getRank(); ++i) { - if (!rt.isDynamicDim(i)) continue; - auto dim = rewriter.createOrFold<tensor::DimOp>(op.getLoc(), arg, i); - operandDynamicDims.push_back(dim); - } - } - } - - auto dispatchOp = rewriter.create<IREE::Flow::DispatchWorkgroupsOp>( - op.getLoc(), workload, op->getResultTypes(), - /*result_dims=*/ValueRange{}, - /*operands=*/args, - /*operand_dims=*/operandDynamicDims, - /*tied_operands=*/Derived::getTiedResultOperandIndices(args)); - { - OpBuilder::InsertionGuard guard(rewriter); - rewriter.setInsertionPointToStart(&dispatchOp.getRegion().front()); - if (failed(Derived::lowerMHLOOp(dispatchOp, op, args, rewriter))) { - return failure(); - } - rewriter.create<IREE::Flow::ReturnOp>(op.getLoc()); - } - rewriter.replaceOp(op, dispatchOp.getResults()); - return success(); - } -}; - -//===----------------------------------------------------------------------===// // Region operations lowering. //===----------------------------------------------------------------------===// @@ -131,32 +86,15 @@ // SortOp //===----------------------------------------------------------------------===// -struct SortOpConversion - : public ConvertToLinalgExtPattern<SortOpConversion, mhlo::SortOp> { - using ConvertToLinalgExtPattern<SortOpConversion, - mhlo::SortOp>::ConvertToLinalgExtPattern; +struct SortOpConversion : public OpConversionPattern<mhlo::SortOp> { + using OpConversionPattern<mhlo::SortOp>::OpConversionPattern; - static SmallVector<int64_t> getTiedResultOperandIndices( - ArrayRef<Value> args) { - return llvm::to_vector<4>(llvm::seq<int64_t>(0, args.size())); - } - - static LogicalResult lowerMHLOOp(IREE::Flow::DispatchWorkgroupsOp dispatchOp, - mhlo::SortOp op, ArrayRef<Value> args, - ConversionPatternRewriter &rewriter) { - auto blockArgs = dispatchOp.getClosureBodyRegion().getArguments(); - SmallVector<Value> initValues; - ImplicitLocOpBuilder b(op.getLoc(), rewriter); - for (auto it : llvm::zip(args, blockArgs)) { - auto argTy = std::get<0>(it).getType().cast<RankedTensorType>(); - auto blockArg = std::get<1>(it); - initValues.push_back( - b.create<IREE::Flow::DispatchTensorLoadOp>(argTy, blockArg)); - } - - auto sortOp = b.create<linalg_ext::SortOp>(op.getResultTypes(), - /*inputs=*/ValueRange{}, - initValues, op.dimensionAttr()); + LogicalResult matchAndRewrite( + mhlo::SortOp op, ArrayRef<Value> args, + ConversionPatternRewriter &rewriter) const final { + auto sortOp = rewriter.create<linalg_ext::SortOp>( + op.getLoc(), op.getResultTypes(), + /*inputs=*/ValueRange{}, args, op.dimensionAttr()); rewriter.inlineRegionBefore(op.comparator(), sortOp.region(), sortOp.region().begin()); Region ®ion = sortOp.region(); @@ -169,12 +107,7 @@ } rewriter.applySignatureConversion(®ion, signature_converter); - for (auto it : llvm::zip(sortOp.getResults(), blockArgs)) { - auto value = std::get<0>(it); - auto target = std::get<1>(it); - b.create<IREE::Flow::DispatchTensorStoreOp>(value, target); - } - + rewriter.replaceOp(op, sortOp->getResults()); return success(); } }; @@ -183,10 +116,8 @@ // ScatterOp //===----------------------------------------------------------------------===// -struct ScatterOpConversion - : public ConvertToLinalgExtPattern<ScatterOpConversion, mhlo::ScatterOp> { - using ConvertToLinalgExtPattern<ScatterOpConversion, - mhlo::ScatterOp>::ConvertToLinalgExtPattern; +struct ScatterOpConversion : public OpConversionPattern<mhlo::ScatterOp> { + using OpConversionPattern<mhlo::ScatterOp>::OpConversionPattern; /// Returns true if the `dimensionNumbers` from the mhlo.scatter op follows a /// canonical form: @@ -275,28 +206,23 @@ return success(); } - static LogicalResult lowerMHLOOp(IREE::Flow::DispatchWorkgroupsOp dispatchOp, - mhlo::ScatterOp op, ArrayRef<Value> args, - ConversionPatternRewriter &rewriter) { + LogicalResult matchAndRewrite( + mhlo::ScatterOp op, ArrayRef<Value> args, + ConversionPatternRewriter &rewriter) const final { if (!hasCanonicalDimensionNumbers(op)) return failure(); ImplicitLocOpBuilder b(op.getLoc(), rewriter); mhlo::ScatterOpAdaptor adaptor(args); - auto blockArgs = dispatchOp.getClosureBodyRegion().getArguments(); - Value original = b.create<IREE::Flow::DispatchTensorLoadOp>( - adaptor.operand().getType().cast<RankedTensorType>(), blockArgs[0]); - Value indices = b.create<IREE::Flow::DispatchTensorLoadOp>( - adaptor.scatter_indices().getType().cast<RankedTensorType>(), - blockArgs[1]); - Value updates = b.create<IREE::Flow::DispatchTensorLoadOp>( - adaptor.updates().getType().cast<RankedTensorType>(), blockArgs[2]); + Value original = adaptor.operand(); + Value indices = adaptor.scatter_indices(); + Value updates = adaptor.updates(); if (failed(collapseBatchDimsIfNeeded(indices, updates, b))) { return failure(); } - auto scatterOp = b.create<linalg_ext::ScatterOp>( - op->getResultTypes(), ValueRange{updates, indices}, + auto scatterOp = rewriter.create<linalg_ext::ScatterOp>( + op.getLoc(), op->getResultTypes(), ValueRange{updates, indices}, ValueRange{original}); rewriter.inlineRegionBefore(op.update_computation(), scatterOp.region(), @@ -311,8 +237,7 @@ signatureConverter.addInputs(0, argType); rewriter.applySignatureConversion(®ion, signatureConverter); - b.create<IREE::Flow::DispatchTensorStoreOp>(scatterOp.getResult(0), - blockArgs[0]); + rewriter.replaceOp(op, scatterOp->getResults()); return success(); } }; @@ -321,9 +246,8 @@ // Pass //===----------------------------------------------------------------------===// -struct ConvertAndDistributeMHLOToLinalgExtPass - : public ConvertAndDistributeMHLOToLinalgExtBase< - ConvertAndDistributeMHLOToLinalgExtPass> { +struct ConvertMHLOToLinalgExtPass + : public ConvertMHLOToLinalgExtBase<ConvertMHLOToLinalgExtPass> { void getDependentDialects(DialectRegistry ®istry) const override { registry.insert<linalg_ext::LinalgExtDialect, linalg::LinalgDialect, IREE::Flow::FlowDialect, StandardOpsDialect, @@ -354,9 +278,8 @@ }; } // namespace -std::unique_ptr<OperationPass<FuncOp>> -createConvertAndDistributeMHLOToLinalgExtPass() { - return std::make_unique<ConvertAndDistributeMHLOToLinalgExtPass>(); +std::unique_ptr<OperationPass<FuncOp>> createConvertMHLOToLinalgExtPass() { + return std::make_unique<ConvertMHLOToLinalgExtPass>(); } } // namespace iree_compiler
diff --git a/iree/compiler/InputConversion/MHLO/Passes.cpp b/iree/compiler/InputConversion/MHLO/Passes.cpp index ebf3147..2b589cb 100644 --- a/iree/compiler/InputConversion/MHLO/Passes.cpp +++ b/iree/compiler/InputConversion/MHLO/Passes.cpp
@@ -66,7 +66,7 @@ // Convert to Linalg. After this point, MHLO will be eliminated. passManager.addNestedPass<FuncOp>( - mlir::iree_compiler::createConvertAndDistributeMHLOToLinalgExtPass()); + mlir::iree_compiler::createConvertMHLOToLinalgExtPass()); passManager.addNestedPass<FuncOp>( mlir::iree_compiler::createMHLOToLinalgOnTensorsPass());
diff --git a/iree/compiler/InputConversion/MHLO/Passes.h b/iree/compiler/InputConversion/MHLO/Passes.h index 5761952..8bc4222 100644 --- a/iree/compiler/InputConversion/MHLO/Passes.h +++ b/iree/compiler/InputConversion/MHLO/Passes.h
@@ -34,9 +34,8 @@ /// Creates XLA-HLO to Linalg on tensors transformation pass. std::unique_ptr<OperationPass<FuncOp>> createMHLOToLinalgOnTensorsPass(); -/// Creates XLA-HLO to LinalgExt and Flow transformation pass. -std::unique_ptr<OperationPass<FuncOp>> -createConvertAndDistributeMHLOToLinalgExtPass(); +/// Creates XLA-HLO to LinalgExt pass. +std::unique_ptr<OperationPass<FuncOp>> createConvertMHLOToLinalgExtPass(); /// Creates XLA-HLO preprocessing transformation pass. In this pass we should /// have all mhlo -> mhlo transformations that are shared between all
diff --git a/iree/compiler/InputConversion/MHLO/Passes.td b/iree/compiler/InputConversion/MHLO/Passes.td index a08fa08..b0082b2 100644 --- a/iree/compiler/InputConversion/MHLO/Passes.td +++ b/iree/compiler/InputConversion/MHLO/Passes.td
@@ -15,12 +15,12 @@ let constructor = "mlir::iree_compiler::createMHLOToLinalgOnTensorsPass()"; } -def ConvertAndDistributeMHLOToLinalgExt +def ConvertMHLOToLinalgExt : Pass<"iree-mhlo-to-linalg-ext", "FuncOp"> { let summary = "Convert from XLA-HLO ops to LinalgExt ops and distribute to Flow ops"; let constructor = - "mlir::iree_compiler::createConvertAndDistributeMHLOToLinalgExtPass()"; + "mlir::iree_compiler::createConvertMHLOToLinalgExtPass()"; } def LegalizeInputTypes :
diff --git a/iree/compiler/InputConversion/MHLO/test/BUILD b/iree/compiler/InputConversion/MHLO/test/BUILD index c3f4485..cef5ec5 100644 --- a/iree/compiler/InputConversion/MHLO/test/BUILD +++ b/iree/compiler/InputConversion/MHLO/test/BUILD
@@ -20,7 +20,7 @@ srcs = enforce_glob( [ "broadcasting.mlir", - "convert_and_distribute_mhlo_to_linalg_ext.mlir", + "convert_mhlo_to_linalg_ext.mlir", "convert_complex_to_real.mlir", "dynamic_shape.mlir", "fft.mlir",
diff --git a/iree/compiler/InputConversion/MHLO/test/CMakeLists.txt b/iree/compiler/InputConversion/MHLO/test/CMakeLists.txt index 9f3a453..d7405b8 100644 --- a/iree/compiler/InputConversion/MHLO/test/CMakeLists.txt +++ b/iree/compiler/InputConversion/MHLO/test/CMakeLists.txt
@@ -15,8 +15,8 @@ lit SRCS "broadcasting.mlir" - "convert_and_distribute_mhlo_to_linalg_ext.mlir" "convert_complex_to_real.mlir" + "convert_mhlo_to_linalg_ext.mlir" "dynamic_shape.mlir" "fft.mlir" "legalize_input_types.mlir"
diff --git a/iree/compiler/InputConversion/MHLO/test/convert_and_distribute_mhlo_to_linalg_ext.mlir b/iree/compiler/InputConversion/MHLO/test/convert_and_distribute_mhlo_to_linalg_ext.mlir deleted file mode 100644 index 8b80c34..0000000 --- a/iree/compiler/InputConversion/MHLO/test/convert_and_distribute_mhlo_to_linalg_ext.mlir +++ /dev/null
@@ -1,342 +0,0 @@ -// RUN: iree-opt -split-input-file -iree-mhlo-to-linalg-ext %s | IreeFileCheck %s - -func @sort_1d(%arg0: tensor<128xi32>) -> (tensor<128xi32>) { - %0 = "mhlo.sort"(%arg0) ( { - ^bb0(%arg2: tensor<i32>, %arg3: tensor<i32>): // no predecessors - %1 = "mhlo.compare"(%arg2, %arg3) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1> - "mhlo.return"(%1) : (tensor<i1>) -> () - }) {dimension = 0 : i64, is_stable = false} : (tensor<128xi32>) -> (tensor<128xi32>) - return %0 : tensor<128xi32> -} -// CHECK-LABEL: func @sort_1d -// CHECK: %[[ARG0:[a-zA-Z0-9]+]] -// CHECK-DAG: %[[C1:.+]] = constant 1 : index -// CHECK: %[[RES:.+]] = flow.dispatch.workgroups -// CHECK-SAME: [%[[C1]], %[[C1]], %[[C1]]](%[[ARG0]]) : (tensor<128xi32>) -> %[[ARG0]] -// CHECK: %[[ARG1:.+]]: !flow.dispatch.tensor<readwrite:128xi32> -// CHECK: %[[IN:.+]] = flow.dispatch.tensor.load %[[ARG1]] -// CHECK: %[[SORT:.+]] = linalg_ext.sort -// CHECK-SAME: dimension(0) -// CHECK-SAME: outs(%[[IN]] : tensor<128xi32>) -// CHECK: ^bb0(%[[ARG2:.+]]: i32, %[[ARG3:.+]]: i32) -// CHECK: %[[CMP:.+]] = cmpi sgt, %[[ARG2]], %[[ARG3]] -// CHECK: linalg_ext.yield %[[CMP]] -// CHECK: flow.dispatch.tensor.store %[[SORT]], %[[ARG1]] -// CHECK: return %[[RES]] - -// ----- - -func @sort_2d(%arg0: tensor<16x32xi32>) -> (tensor<16x32xi32>) { - %0 = "mhlo.sort"(%arg0) ( { - ^bb0(%arg2: tensor<i32>, %arg3: tensor<i32>): // no predecessors - %1 = "mhlo.compare"(%arg2, %arg3) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1> - "mhlo.return"(%1) : (tensor<i1>) -> () - }) {dimension = 0 : i64, is_stable = false} : (tensor<16x32xi32>) -> (tensor<16x32xi32>) - return %0 : tensor<16x32xi32> -} -// CHECK-LABEL: func @sort_2d -// CHECK: %[[ARG0:[a-zA-Z0-9]+]] -// CHECK-DAG: %[[C1:.+]] = constant 1 : index -// CHECK: %[[RES:.+]] = flow.dispatch.workgroups -// CHECK-SAME: [%[[C1]], %[[C1]], %[[C1]]](%[[ARG0]]) : (tensor<16x32xi32>) -> %[[ARG0]] -// CHECK: %[[ARG1:.+]]: !flow.dispatch.tensor<readwrite:16x32xi32> -// CHECK: %[[IN:.+]] = flow.dispatch.tensor.load %[[ARG1]] -// CHECK: %[[SORT:.+]] = linalg_ext.sort -// CHECK-SAME: dimension(0) -// CHECK-SAME: outs(%[[IN]] : tensor<16x32xi32>) -// CHECK: ^bb0(%[[ARG2:.+]]: i32, %[[ARG3:.+]]: i32) -// CHECK: %[[CMP:.+]] = cmpi sgt, %[[ARG2]], %[[ARG3]] -// CHECK: linalg_ext.yield %[[CMP]] -// CHECK: flow.dispatch.tensor.store %[[SORT]], %[[ARG1]] -// CHECK: return %[[RES]] - -// ----- - -func @topk(%arg0: tensor<128xi32>, %arg1: tensor<128xi32>) -> (tensor<128xi32>) { - %0:2 = "mhlo.sort"(%arg0, %arg1) ( { - ^bb0(%arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<i32>, %arg5: tensor<i32>): // no predecessors - %1 = "mhlo.compare"(%arg2, %arg3) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1> - "mhlo.return"(%1) : (tensor<i1>) -> () - }) {dimension = 0 : i64, is_stable = false} : (tensor<128xi32>, tensor<128xi32>) -> (tensor<128xi32>, tensor<128xi32>) - return %0#0 : tensor<128xi32> -} -// CHECK-LABEL: func @topk -// CHECK: %[[ARG0:[a-zA-Z0-9]+]] -// CHECK: %[[ARG1:[a-zA-Z0-9]+]] -// CHECK-DAG: %[[C1:.+]] = constant 1 : index -// CHECK: %[[RES:.+]]:2 = flow.dispatch.workgroups -// CHECK-SAME: [%[[C1]], %[[C1]], %[[C1]]](%[[ARG0]], %[[ARG1]]) -// CHECK-SAME: : (tensor<128xi32>, tensor<128xi32>) -> (%[[ARG0]], %[[ARG1]]) -// CHECK: %[[ARG2:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readwrite:128xi32> -// CHECK: %[[ARG3:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readwrite:128xi32> -// CHECK: %[[IN1:.+]] = flow.dispatch.tensor.load %[[ARG2]] -// CHECK: %[[IN2:.+]] = flow.dispatch.tensor.load %[[ARG3]] -// CHECK: %[[SORT:.+]]:2 = linalg_ext.sort -// CHECK-SAME: dimension(0) -// CHECK-SAME: outs(%[[IN1]], %[[IN2]] : tensor<128xi32>, tensor<128xi32>) -// CHECK: ^bb0(%[[ARG4:.+]]: i32, %[[ARG5:.+]]: i32, %{{.*}}: i32, %{{.*}}: i32) -// CHECK: %[[CMP:.+]] = cmpi sgt, %[[ARG4]], %[[ARG5]] -// CHECK: linalg_ext.yield %[[CMP]] -// CHECK: flow.dispatch.tensor.store %[[SORT]]#0, %[[ARG2]] -// CHECK: flow.dispatch.tensor.store %[[SORT]]#1, %[[ARG3]] -// CHECK: return %[[RES]]#0 - -// ----- - -func @scatter_update_scalar_1D(%arg0: tensor<8xi32>, %arg1: tensor<4x1xi32>, - %arg2: tensor<4xi32>) -> tensor<8xi32> { - %0 = "mhlo.scatter"(%arg0, %arg1, %arg2) ( { - ^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors - "mhlo.return"(%arg4) : (tensor<i32>) -> () - }) { - indices_are_sorted = false, - scatter_dimension_numbers = { - index_vector_dim = 1 : i64, - inserted_window_dims = dense<0> : tensor<1xi64>, - scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, - update_window_dims = dense<> : tensor<0xi64> - }, - unique_indices = false - } : (tensor<8xi32>, tensor<4x1xi32>, tensor<4xi32>) -> tensor<8xi32> - return %0 : tensor<8xi32> -} -// CHECK-LABEL: func @scatter_update_scalar_1D -// CHECK: %[[ARG0:[a-zA-Z0-9]+]] -// CHECK: %[[ARG1:[a-zA-Z0-9]+]] -// CHECK: %[[ARG2:[a-zA-Z0-9]+]] -// CHECK-DAG: %[[C1:.+]] = constant 1 : index -// CHECK: %[[RES:.+]] = flow.dispatch.workgroups -// CHECK-SAME: [%[[C1]], %[[C1]], %[[C1]]](%[[ARG0]], %[[ARG1]], %[[ARG2]]) -// CHECK-SAME: : (tensor<8xi32>, tensor<4x1xi32>, tensor<4xi32>) -> %[[ARG0]] -// CHECK: %[[ARG3:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readwrite:8xi32> -// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:4x1xi32> -// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:4xi32> -// CHECK: %[[ORIGINAL:.+]] = flow.dispatch.tensor.load %[[ARG3]] -// CHECK: %[[INDICES:.+]] = flow.dispatch.tensor.load %[[ARG4]] -// CHECK: %[[UPDATES:.+]] = flow.dispatch.tensor.load %[[ARG5]] -// CHECK: %[[SCATTER:.+]] = linalg_ext.scatter -// CHECK-SAME: ins(%[[UPDATES]], %[[INDICES]] : tensor<4xi32>, tensor<4x1xi32>) -// CHECK-SAME: outs(%[[ORIGINAL]] : tensor<8xi32>) -// CHECK: ^bb0(%[[V1:.+]]: i32, %[[V2:.+]]: i32): // no predecessors -// CEECK: linalg.yield %[[V1]] -// CHECK: flow.dispatch.tensor.store %[[SCATTER]], %[[ARG3]] -// CHECK: return %[[RES]] - -// ----- - -func @scatter_update_scalar_2D(%arg0: tensor<4x3xi32>, %arg1: tensor<3x2xi32>, - %arg2: tensor<3xi32>) -> tensor<4x3xi32> { - %0 = "mhlo.scatter"(%arg0, %arg1, %arg2) ( { - ^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors - "mhlo.return"(%arg4) : (tensor<i32>) -> () - }) {indices_are_sorted = false, - scatter_dimension_numbers = { - index_vector_dim = 1 : i64, - inserted_window_dims = dense<[0, 1]> : tensor<2xi64>, - scatter_dims_to_operand_dims = dense<[0, 1]> : tensor<2xi64>, - update_window_dims = dense<> : tensor<0xi64> - }, - unique_indices = false - } : (tensor<4x3xi32>, tensor<3x2xi32>, tensor<3xi32>) -> tensor<4x3xi32> - return %0 : tensor<4x3xi32> -} -// CHECK-LABEL: func @scatter_update_scalar_2D -// CHECK: %[[ARG0:[a-zA-Z0-9]+]] -// CHECK: %[[ARG1:[a-zA-Z0-9]+]] -// CHECK: %[[ARG2:[a-zA-Z0-9]+]] -// CHECK-DAG: %[[C1:.+]] = constant 1 : index -// CHECK: %[[RES:.+]] = flow.dispatch.workgroups -// CHECK-SAME: [%[[C1]], %[[C1]], %[[C1]]](%[[ARG0]], %[[ARG1]], %[[ARG2]]) -// CHECK-SAME: : (tensor<4x3xi32>, tensor<3x2xi32>, tensor<3xi32>) -> %[[ARG0]] -// CHECK: %[[ARG3:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readwrite:4x3xi32> -// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:3x2xi32> -// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:3xi32> -// CHECK: %[[ORIGINAL:.+]] = flow.dispatch.tensor.load %[[ARG3]] -// CHECK: %[[INDICES:.+]] = flow.dispatch.tensor.load %[[ARG4]] -// CHECK: %[[UPDATES:.+]] = flow.dispatch.tensor.load %[[ARG5]] -// CHECK: %[[SCATTER:.+]] = linalg_ext.scatter -// CHECK-SAME: ins(%[[UPDATES]], %[[INDICES]] : tensor<3xi32>, tensor<3x2xi32>) -// CHECK-SAME: outs(%[[ORIGINAL]] : tensor<4x3xi32>) -// CHECK: ^bb0(%[[V1:.+]]: i32, %[[V2:.+]]: i32): // no predecessors -// CEECK: linalg.yield %[[V1]] -// CHECK: flow.dispatch.tensor.store %[[SCATTER]], %[[ARG3]] -// CHECK: return %[[RES]] - -// ----- - -func @scatter_update_slice_2D(%arg0: tensor<6x3xi32>, %arg1: tensor<2x1xi32>, - %arg2: tensor<2x3xi32>) -> tensor<6x3xi32> { - %0 = "mhlo.scatter"(%arg0, %arg1, %arg2) ( { - ^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors - "mhlo.return"(%arg4) : (tensor<i32>) -> () - }) { - indices_are_sorted = false, - scatter_dimension_numbers = { - index_vector_dim = 1 : i64, - inserted_window_dims = dense<0> : tensor<1xi64>, - scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, - update_window_dims = dense<1> : tensor<1xi64> - }, - unique_indices = false - } : (tensor<6x3xi32>, tensor<2x1xi32>, tensor<2x3xi32>) -> tensor<6x3xi32> - return %0 : tensor<6x3xi32> -} -// CHECK-LABEL: func @scatter_update_slice_2D -// CHECK: %[[ARG0:[a-zA-Z0-9]+]] -// CHECK: %[[ARG1:[a-zA-Z0-9]+]] -// CHECK: %[[ARG2:[a-zA-Z0-9]+]] -// CHECK-DAG: %[[C1:.+]] = constant 1 : index -// CHECK: %[[RES:.+]] = flow.dispatch.workgroups -// CHECK-SAME: [%[[C1]], %[[C1]], %[[C1]]](%[[ARG0]], %[[ARG1]], %[[ARG2]]) -// CHECK-SAME: : (tensor<6x3xi32>, tensor<2x1xi32>, tensor<2x3xi32>) -> %[[ARG0]] -// CHECK: %[[ARG3:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readwrite:6x3xi32> -// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:2x1xi32> -// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:2x3xi32> -// CHECK: %[[ORIGINAL:.+]] = flow.dispatch.tensor.load %[[ARG3]] -// CHECK: %[[INDICES:.+]] = flow.dispatch.tensor.load %[[ARG4]] -// CHECK: %[[UPDATES:.+]] = flow.dispatch.tensor.load %[[ARG5]] -// CHECK: %[[SCATTER:.+]] = linalg_ext.scatter -// CHECK-SAME: ins(%[[UPDATES]], %[[INDICES]] : tensor<2x3xi32>, tensor<2x1xi32>) -// CHECK-SAME: outs(%[[ORIGINAL]] : tensor<6x3xi32>) -// CHECK: ^bb0(%[[V1:.+]]: i32, %[[V2:.+]]: i32): // no predecessors -// CEECK: linalg.yield %[[V1]] -// CHECK: flow.dispatch.tensor.store %[[SCATTER]], %[[ARG3]] -// CHECK: return %[[RES]] - -// ----- - -func @scatter_add_slice_2D(%arg0: tensor<6x3xi32>, %arg1: tensor<2x1xi32>, - %arg2: tensor<2x3xi32>) -> tensor<6x3xi32> { - %0 = "mhlo.scatter"(%arg0, %arg1, %arg2) ( { - ^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors - %1 = mhlo.add %arg3, %arg4 : tensor<i32> - "mhlo.return"(%1) : (tensor<i32>) -> () - }) { - indices_are_sorted = false, - scatter_dimension_numbers = { - index_vector_dim = 1 : i64, - inserted_window_dims = dense<0> : tensor<1xi64>, - scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, - update_window_dims = dense<1> : tensor<1xi64> - }, - unique_indices = false - } : (tensor<6x3xi32>, tensor<2x1xi32>, tensor<2x3xi32>) -> tensor<6x3xi32> - return %0 : tensor<6x3xi32> -} -// CHECK-LABEL: func @scatter_add_slice_2D -// CHECK: %[[ARG0:[a-zA-Z0-9]+]] -// CHECK: %[[ARG1:[a-zA-Z0-9]+]] -// CHECK: %[[ARG2:[a-zA-Z0-9]+]] -// CHECK-DAG: %[[C1:.+]] = constant 1 : index -// CHECK: %[[RES:.+]] = flow.dispatch.workgroups -// CHECK-SAME: [%[[C1]], %[[C1]], %[[C1]]](%[[ARG0]], %[[ARG1]], %[[ARG2]]) -// CHECK-SAME: : (tensor<6x3xi32>, tensor<2x1xi32>, tensor<2x3xi32>) -> %[[ARG0]] -// CHECK: %[[ARG3:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readwrite:6x3xi32> -// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:2x1xi32> -// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:2x3xi32> -// CHECK: %[[ORIGINAL:.+]] = flow.dispatch.tensor.load %[[ARG3]] -// CHECK: %[[INDICES:.+]] = flow.dispatch.tensor.load %[[ARG4]] -// CHECK: %[[UPDATES:.+]] = flow.dispatch.tensor.load %[[ARG5]] -// CHECK: %[[SCATTER:.+]] = linalg_ext.scatter -// CHECK-SAME: ins(%[[UPDATES]], %[[INDICES]] : tensor<2x3xi32>, tensor<2x1xi32>) -// CHECK-SAME: outs(%[[ORIGINAL]] : tensor<6x3xi32>) -// CHECK: ^bb0(%[[V1:.+]]: i32, %[[V2:.+]]: i32): // no predecessors -// -// The order is reverse. -// CHECK: %[[V3:.+]] = addi %[[V2]], %[[V1]] -// CEECK: linalg.yield %[[V3]] -// CHECK: flow.dispatch.tensor.store %[[SCATTER]], %[[ARG3]] -// CHECK: return %[[RES]] - -// ----- - -func @scatter_update_batch_scalar_1D(%arg0: tensor<8xi32>, - %arg1: tensor<3x4x1xi32>, %arg2: tensor<3x4xi32>) -> tensor<8xi32> { - %0 = "mhlo.scatter"(%arg0, %arg1, %arg2) ( { - ^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors - "mhlo.return"(%arg4) : (tensor<i32>) -> () - }) { - indices_are_sorted = false, - scatter_dimension_numbers = { - index_vector_dim = 2 : i64, - inserted_window_dims = dense<0> : tensor<i64>, - scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, - update_window_dims = dense<> : tensor<0xi64> - }, - unique_indices = false - } : (tensor<8xi32>, tensor<3x4x1xi32>, tensor<3x4xi32>) -> tensor<8xi32> - return %0 : tensor<8xi32> -} -// CHECK-LABEL: func @scatter_update_batch_scalar_1D -// CHECK: %[[ARG0:[a-zA-Z0-9]+]] -// CHECK: %[[ARG1:[a-zA-Z0-9]+]] -// CHECK: %[[ARG2:[a-zA-Z0-9]+]] -// CHECK-DAG: %[[C1:.+]] = constant 1 : index -// CHECK: %[[RES:.+]] = flow.dispatch.workgroups -// CHECK-SAME: [%[[C1]], %[[C1]], %[[C1]]](%[[ARG0]], %[[ARG1]], %[[ARG2]]) -// CHECK-SAME: : (tensor<8xi32>, tensor<3x4x1xi32>, tensor<3x4xi32>) -> %[[ARG0]] -// CHECK: %[[ARG3:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readwrite:8xi32> -// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:3x4x1xi32> -// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:3x4xi32> -// CHECK: %[[ORIGINAL:.+]] = flow.dispatch.tensor.load %[[ARG3]] -// CHECK: %[[INDICES:.+]] = flow.dispatch.tensor.load %[[ARG4]] -// CHECK: %[[UPDATES:.+]] = flow.dispatch.tensor.load %[[ARG5]] -// CHECK: %[[COLLAPSED_INDICES:.+]] = linalg.tensor_collapse_shape -// CHECK-SAME: %[[INDICES]] {{\[}}[0, 1], [2]] : tensor<3x4x1xi32> into tensor<12x1xi32> -// CHECK: %[[COLLAPSED_UPDATES:.+]] = linalg.tensor_collapse_shape -// CHECK-SAME: %[[UPDATES]] {{\[}}[0, 1]] : tensor<3x4xi32> into tensor<12xi32> -// CHECK: %[[SCATTER:.+]] = linalg_ext.scatter -// CHECK-SAME: ins(%[[COLLAPSED_UPDATES]], %[[COLLAPSED_INDICES]] : tensor<12xi32>, tensor<12x1xi32>) -// CHECK-SAME: outs(%[[ORIGINAL]] : tensor<8xi32>) -// CHECK: ^bb0(%[[V1:.+]]: i32, %[[V2:.+]]: i32): // no predecessors -// CEECK: linalg.yield %[[V1]] -// CHECK: flow.dispatch.tensor.store %[[SCATTER]], %[[ARG3]] -// CHECK: return %[[RES]] - -// ----- - -func @scatter_update_batch_slice_3D_dynamic(%arg0: tensor<1x24x512xi32>, - %arg1: tensor<?x3x2xi32>, %arg2: tensor<?x3x512xi32>) -> tensor<1x24x512xi32> { - %0 = "mhlo.scatter"(%arg0, %arg1, %arg2) ( { - ^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors - "mhlo.return"(%arg4) : (tensor<i32>) -> () - }) {indices_are_sorted = false, - scatter_dimension_numbers = { - index_vector_dim = 2 : i64, - inserted_window_dims = dense<[0, 1]> : tensor<2xi64>, - scatter_dims_to_operand_dims = dense<[0, 1]> : tensor<2xi64>, - update_window_dims = dense<2> : tensor<1xi64> - }, - unique_indices = false - } : (tensor<1x24x512xi32>, tensor<?x3x2xi32>, tensor<?x3x512xi32>) -> tensor<1x24x512xi32> - return %0 : tensor<1x24x512xi32> -} -// CHECK-LABEL: func @scatter_update_batch_slice_3D_dynamic -// CHECK: %[[ARG0:[a-zA-Z0-9]+]] -// CHECK: %[[ARG1:[a-zA-Z0-9]+]] -// CHECK: %[[ARG2:[a-zA-Z0-9]+]] -// CHECK-DAG: %[[C1:.+]] = constant 1 : index -// CHECK-DAG: %[[C0:.+]] = constant 0 : index -// CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG1]], %[[C0]] : tensor<?x3x2xi32> -// CHECK-DAG: %[[C0:.+]] = constant 0 : index -// CHECK-DAG: %[[DIM2:.+]] = tensor.dim %[[ARG2]], %[[C0]] : tensor<?x3x512xi32> -// CHECK: %[[RES:.+]] = flow.dispatch.workgroups -// CHECK-SAME: [%[[C1]], %[[C1]], %[[C1]]](%[[ARG0]], %[[ARG1]], %[[ARG2]]) -// CHECK-SAME: : (tensor<1x24x512xi32>, tensor<?x3x2xi32>{%[[DIM1]]}, tensor<?x3x512xi32>{%[[DIM2]]}) -> %[[ARG0]] -// CHECK: %[[ARG3:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readwrite:1x24x512xi32> -// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:?x3x2xi32> -// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:?x3x512xi32> -// CHECK: %[[ORIGINAL:.+]] = flow.dispatch.tensor.load %[[ARG3]] -// CHECK: %[[INDICES:.+]] = flow.dispatch.tensor.load %[[ARG4]] -// CHECK: %[[UPDATES:.+]] = flow.dispatch.tensor.load %[[ARG5]] -// CHECK: %[[COLLAPSED_INDICES:.+]] = linalg.tensor_collapse_shape -// CHECK-SAME: %[[INDICES]] {{\[}}[0, 1], [2]] : tensor<?x3x2xi32> into tensor<?x2xi32> -// CHECK: %[[COLLAPSED_UPDATES:.+]] = linalg.tensor_collapse_shape -// CHECK-SAME: %[[UPDATES]] {{\[}}[0, 1], [2]] : tensor<?x3x512xi32> into tensor<?x512xi32> -// CHECK: %[[SCATTER:.+]] = linalg_ext.scatter -// CHECK-SAME: ins(%[[COLLAPSED_UPDATES]], %[[COLLAPSED_INDICES]] : tensor<?x512xi32>, tensor<?x2xi32>) -// CHECK-SAME: outs(%[[ORIGINAL]] : tensor<1x24x512xi32>) -// CHECK: ^bb0(%[[V1:.+]]: i32, %[[V2:.+]]: i32): // no predecessors -// CEECK: linalg.yield %[[V1]] -// CHECK: flow.dispatch.tensor.store %[[SCATTER]], %[[ARG3]] -// CHECK: return %[[RES]]
diff --git a/iree/compiler/InputConversion/MHLO/test/convert_mhlo_to_linalg_ext.mlir b/iree/compiler/InputConversion/MHLO/test/convert_mhlo_to_linalg_ext.mlir new file mode 100644 index 0000000..e25865b --- /dev/null +++ b/iree/compiler/InputConversion/MHLO/test/convert_mhlo_to_linalg_ext.mlir
@@ -0,0 +1,250 @@ +// RUN: iree-opt -split-input-file -iree-mhlo-to-linalg-ext %s | IreeFileCheck %s + +func @sort_1d(%arg0: tensor<128xi32>) -> (tensor<128xi32>) { + %0 = "mhlo.sort"(%arg0) ( { + ^bb0(%arg2: tensor<i32>, %arg3: tensor<i32>): // no predecessors + %1 = "mhlo.compare"(%arg2, %arg3) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1> + "mhlo.return"(%1) : (tensor<i1>) -> () + }) {dimension = 0 : i64, is_stable = false} : (tensor<128xi32>) -> (tensor<128xi32>) + return %0 : tensor<128xi32> +} +// CHECK-LABEL: func @sort_1d +// CHECK: %[[ARG0:[a-zA-Z0-9]+]] +// CHECK: %[[SORT:.+]] = linalg_ext.sort +// CHECK-SAME: dimension(0) +// CHECK-SAME: outs(%[[ARG0]] : tensor<128xi32>) +// CHECK: ^bb0(%[[ARG1:.+]]: i32, %[[ARG2:.+]]: i32) +// CHECK: %[[CMP:.+]] = cmpi sgt, %[[ARG1]], %[[ARG2]] +// CHECK: linalg_ext.yield %[[CMP]] +// CHECK: return %[[SORT]] + +// ----- + +func @sort_2d(%arg0: tensor<16x32xi32>) -> (tensor<16x32xi32>) { + %0 = "mhlo.sort"(%arg0) ( { + ^bb0(%arg2: tensor<i32>, %arg3: tensor<i32>): // no predecessors + %1 = "mhlo.compare"(%arg2, %arg3) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1> + "mhlo.return"(%1) : (tensor<i1>) -> () + }) {dimension = 0 : i64, is_stable = false} : (tensor<16x32xi32>) -> (tensor<16x32xi32>) + return %0 : tensor<16x32xi32> +} +// CHECK-LABEL: func @sort_2d +// CHECK: %[[ARG0:[a-zA-Z0-9]+]] +// CHECK: %[[SORT:.+]] = linalg_ext.sort +// CHECK-SAME: dimension(0) +// CHECK-SAME: outs(%[[ARG0]] : tensor<16x32xi32>) +// CHECK: ^bb0(%[[ARG1:.+]]: i32, %[[ARG2:.+]]: i32) +// CHECK: %[[CMP:.+]] = cmpi sgt, %[[ARG1]], %[[ARG2]] +// CHECK: linalg_ext.yield %[[CMP]] +// CHECK: return %[[SORT]] + +// ----- + +func @topk(%arg0: tensor<128xi32>, %arg1: tensor<128xi32>) -> (tensor<128xi32>) { + %0:2 = "mhlo.sort"(%arg0, %arg1) ( { + ^bb0(%arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<i32>, %arg5: tensor<i32>): // no predecessors + %1 = "mhlo.compare"(%arg2, %arg3) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1> + "mhlo.return"(%1) : (tensor<i1>) -> () + }) {dimension = 0 : i64, is_stable = false} : (tensor<128xi32>, tensor<128xi32>) -> (tensor<128xi32>, tensor<128xi32>) + return %0#0 : tensor<128xi32> +} +// CHECK-LABEL: func @topk +// CHECK: %[[ARG0:[a-zA-Z0-9]+]] +// CHECK: %[[ARG1:[a-zA-Z0-9]+]] +// CHECK: %[[SORT:.+]]:2 = linalg_ext.sort +// CHECK-SAME: dimension(0) +// CHECK-SAME: outs(%[[ARG0]], %[[ARG1]] : tensor<128xi32>, tensor<128xi32>) +// CHECK: ^bb0(%[[ARG2:.+]]: i32, %[[ARG3:.+]]: i32, %{{.*}}: i32, %{{.*}}: i32) +// CHECK: %[[CMP:.+]] = cmpi sgt, %[[ARG2]], %[[ARG3]] +// CHECK: linalg_ext.yield %[[CMP]] +// CHECK: return %[[SORT]]#0 + +// ----- + +func @scatter_update_scalar_1D(%arg0: tensor<8xi32>, %arg1: tensor<4x1xi32>, + %arg2: tensor<4xi32>) -> tensor<8xi32> { + %0 = "mhlo.scatter"(%arg0, %arg1, %arg2) ( { + ^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors + "mhlo.return"(%arg4) : (tensor<i32>) -> () + }) { + indices_are_sorted = false, + scatter_dimension_numbers = { + index_vector_dim = 1 : i64, + inserted_window_dims = dense<0> : tensor<1xi64>, + scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, + update_window_dims = dense<> : tensor<0xi64> + }, + unique_indices = false + } : (tensor<8xi32>, tensor<4x1xi32>, tensor<4xi32>) -> tensor<8xi32> + return %0 : tensor<8xi32> +} +// CHECK-LABEL: func @scatter_update_scalar_1D +// CHECK: %[[ARG0:[a-zA-Z0-9]+]] +// CHECK: %[[ARG1:[a-zA-Z0-9]+]] +// CHECK: %[[ARG2:[a-zA-Z0-9]+]] +// CHECK: %[[SCATTER:.+]] = linalg_ext.scatter +// CHECK-SAME: ins(%[[ARG2]], %[[ARG1]] : tensor<4xi32>, tensor<4x1xi32>) +// CHECK-SAME: outs(%[[ARG0]] : tensor<8xi32>) +// CHECK: ^bb0(%[[V1:.+]]: i32, %[[V2:.+]]: i32): // no predecessors +// CEECK: linalg.yield %[[V1]] +// CHECK: return %[[SCATTER]] + +// ----- + +func @scatter_update_scalar_2D(%arg0: tensor<4x3xi32>, %arg1: tensor<3x2xi32>, + %arg2: tensor<3xi32>) -> tensor<4x3xi32> { + %0 = "mhlo.scatter"(%arg0, %arg1, %arg2) ( { + ^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors + "mhlo.return"(%arg4) : (tensor<i32>) -> () + }) {indices_are_sorted = false, + scatter_dimension_numbers = { + index_vector_dim = 1 : i64, + inserted_window_dims = dense<[0, 1]> : tensor<2xi64>, + scatter_dims_to_operand_dims = dense<[0, 1]> : tensor<2xi64>, + update_window_dims = dense<> : tensor<0xi64> + }, + unique_indices = false + } : (tensor<4x3xi32>, tensor<3x2xi32>, tensor<3xi32>) -> tensor<4x3xi32> + return %0 : tensor<4x3xi32> +} +// CHECK-LABEL: func @scatter_update_scalar_2D +// CHECK: %[[ARG0:[a-zA-Z0-9]+]] +// CHECK: %[[ARG1:[a-zA-Z0-9]+]] +// CHECK: %[[ARG2:[a-zA-Z0-9]+]] +// CHECK: %[[SCATTER:.+]] = linalg_ext.scatter +// CHECK-SAME: ins(%[[ARG2]], %[[ARG1]] : tensor<3xi32>, tensor<3x2xi32>) +// CHECK-SAME: outs(%[[ARG0]] : tensor<4x3xi32>) +// CHECK: ^bb0(%[[V1:.+]]: i32, %[[V2:.+]]: i32): // no predecessors +// CEECK: linalg.yield %[[V1]] +// CHECK: return %[[SCATTER]] + +// ----- + +func @scatter_update_slice_2D(%arg0: tensor<6x3xi32>, %arg1: tensor<2x1xi32>, + %arg2: tensor<2x3xi32>) -> tensor<6x3xi32> { + %0 = "mhlo.scatter"(%arg0, %arg1, %arg2) ( { + ^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors + "mhlo.return"(%arg4) : (tensor<i32>) -> () + }) { + indices_are_sorted = false, + scatter_dimension_numbers = { + index_vector_dim = 1 : i64, + inserted_window_dims = dense<0> : tensor<1xi64>, + scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, + update_window_dims = dense<1> : tensor<1xi64> + }, + unique_indices = false + } : (tensor<6x3xi32>, tensor<2x1xi32>, tensor<2x3xi32>) -> tensor<6x3xi32> + return %0 : tensor<6x3xi32> +} +// CHECK-LABEL: func @scatter_update_slice_2D +// CHECK: %[[ARG0:[a-zA-Z0-9]+]] +// CHECK: %[[ARG1:[a-zA-Z0-9]+]] +// CHECK: %[[ARG2:[a-zA-Z0-9]+]] +// CHECK: %[[SCATTER:.+]] = linalg_ext.scatter +// CHECK-SAME: ins(%[[ARG2]], %[[ARG1]] : tensor<2x3xi32>, tensor<2x1xi32>) +// CHECK-SAME: outs(%[[ARG0]] : tensor<6x3xi32>) +// CHECK: ^bb0(%[[V1:.+]]: i32, %[[V2:.+]]: i32): // no predecessors +// CEECK: linalg.yield %[[V1]] +// CHECK: return %[[SCATTER]] + +// ----- + +func @scatter_add_slice_2D(%arg0: tensor<6x3xi32>, %arg1: tensor<2x1xi32>, + %arg2: tensor<2x3xi32>) -> tensor<6x3xi32> { + %0 = "mhlo.scatter"(%arg0, %arg1, %arg2) ( { + ^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors + %1 = mhlo.add %arg3, %arg4 : tensor<i32> + "mhlo.return"(%1) : (tensor<i32>) -> () + }) { + indices_are_sorted = false, + scatter_dimension_numbers = { + index_vector_dim = 1 : i64, + inserted_window_dims = dense<0> : tensor<1xi64>, + scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, + update_window_dims = dense<1> : tensor<1xi64> + }, + unique_indices = false + } : (tensor<6x3xi32>, tensor<2x1xi32>, tensor<2x3xi32>) -> tensor<6x3xi32> + return %0 : tensor<6x3xi32> +} +// CHECK-LABEL: func @scatter_add_slice_2D +// CHECK: %[[ARG0:[a-zA-Z0-9]+]] +// CHECK: %[[ARG1:[a-zA-Z0-9]+]] +// CHECK: %[[ARG2:[a-zA-Z0-9]+]] +// CHECK: %[[SCATTER:.+]] = linalg_ext.scatter +// CHECK-SAME: ins(%[[ARG2]], %[[ARG1]] : tensor<2x3xi32>, tensor<2x1xi32>) +// CHECK-SAME: outs(%[[ARG0]] : tensor<6x3xi32>) +// CHECK: ^bb0(%[[V1:.+]]: i32, %[[V2:.+]]: i32): // no predecessors +// +// The order is reverse. +// CHECK: %[[V3:.+]] = addi %[[V2]], %[[V1]] +// CEECK: linalg.yield %[[V3]] +// CHECK: return %[[SCATTER]] + +// ----- + +func @scatter_update_batch_scalar_1D(%arg0: tensor<8xi32>, + %arg1: tensor<3x4x1xi32>, %arg2: tensor<3x4xi32>) -> tensor<8xi32> { + %0 = "mhlo.scatter"(%arg0, %arg1, %arg2) ( { + ^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors + "mhlo.return"(%arg4) : (tensor<i32>) -> () + }) { + indices_are_sorted = false, + scatter_dimension_numbers = { + index_vector_dim = 2 : i64, + inserted_window_dims = dense<0> : tensor<i64>, + scatter_dims_to_operand_dims = dense<0> : tensor<1xi64>, + update_window_dims = dense<> : tensor<0xi64> + }, + unique_indices = false + } : (tensor<8xi32>, tensor<3x4x1xi32>, tensor<3x4xi32>) -> tensor<8xi32> + return %0 : tensor<8xi32> +} +// CHECK-LABEL: func @scatter_update_batch_scalar_1D +// CHECK: %[[ARG0:[a-zA-Z0-9]+]] +// CHECK: %[[ARG1:[a-zA-Z0-9]+]] +// CHECK: %[[ARG2:[a-zA-Z0-9]+]] +// CHECK: %[[COLLAPSED_INDICES:.+]] = linalg.tensor_collapse_shape +// CHECK-SAME: %[[ARG1]] {{\[}}[0, 1], [2]] : tensor<3x4x1xi32> into tensor<12x1xi32> +// CHECK: %[[COLLAPSED_UPDATES:.+]] = linalg.tensor_collapse_shape +// CHECK-SAME: %[[ARG2]] {{\[}}[0, 1]] : tensor<3x4xi32> into tensor<12xi32> +// CHECK: %[[SCATTER:.+]] = linalg_ext.scatter +// CHECK-SAME: ins(%[[COLLAPSED_UPDATES]], %[[COLLAPSED_INDICES]] : tensor<12xi32>, tensor<12x1xi32>) +// CHECK-SAME: outs(%[[ARG0]] : tensor<8xi32>) +// CHECK: ^bb0(%[[V1:.+]]: i32, %[[V2:.+]]: i32): // no predecessors +// CEECK: linalg.yield %[[V1]] +// CHECK: return %[[SCATTER]] + +// ----- + +func @scatter_update_batch_slice_3D_dynamic(%arg0: tensor<1x24x512xi32>, + %arg1: tensor<?x3x2xi32>, %arg2: tensor<?x3x512xi32>) -> tensor<1x24x512xi32> { + %0 = "mhlo.scatter"(%arg0, %arg1, %arg2) ( { + ^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors + "mhlo.return"(%arg4) : (tensor<i32>) -> () + }) {indices_are_sorted = false, + scatter_dimension_numbers = { + index_vector_dim = 2 : i64, + inserted_window_dims = dense<[0, 1]> : tensor<2xi64>, + scatter_dims_to_operand_dims = dense<[0, 1]> : tensor<2xi64>, + update_window_dims = dense<2> : tensor<1xi64> + }, + unique_indices = false + } : (tensor<1x24x512xi32>, tensor<?x3x2xi32>, tensor<?x3x512xi32>) -> tensor<1x24x512xi32> + return %0 : tensor<1x24x512xi32> +} +// CHECK-LABEL: func @scatter_update_batch_slice_3D_dynamic +// CHECK: %[[ARG0:[a-zA-Z0-9]+]] +// CHECK: %[[ARG1:[a-zA-Z0-9]+]] +// CHECK: %[[ARG2:[a-zA-Z0-9]+]] +// CHECK: %[[COLLAPSED_INDICES:.+]] = linalg.tensor_collapse_shape +// CHECK-SAME: %[[ARG1]] {{\[}}[0, 1], [2]] : tensor<?x3x2xi32> into tensor<?x2xi32> +// CHECK: %[[COLLAPSED_UPDATES:.+]] = linalg.tensor_collapse_shape +// CHECK-SAME: %[[ARG2]] {{\[}}[0, 1], [2]] : tensor<?x3x512xi32> into tensor<?x512xi32> +// CHECK: %[[SCATTER:.+]] = linalg_ext.scatter +// CHECK-SAME: ins(%[[COLLAPSED_UPDATES]], %[[COLLAPSED_INDICES]] : tensor<?x512xi32>, tensor<?x2xi32>) +// CHECK-SAME: outs(%[[ARG0]] : tensor<1x24x512xi32>) +// CHECK: ^bb0(%[[V1:.+]]: i32, %[[V2:.+]]: i32): // no predecessors +// CEECK: linalg.yield %[[V1]] +// CHECK: return %[[SCATTER]]
diff --git a/iree/hal/allocator_heap.c b/iree/hal/allocator_heap.c index 7951398..21b7f3e 100644 --- a/iree/hal/allocator_heap.c +++ b/iree/hal/allocator_heap.c
@@ -28,7 +28,8 @@ IREE_TRACE_ZONE_BEGIN(z0); iree_hal_heap_allocator_t* allocator = NULL; - iree_host_size_t total_size = sizeof(*allocator) + identifier.size; + iree_host_size_t total_size = + iree_sizeof_struct(*allocator) + identifier.size; iree_status_t status = iree_allocator_malloc(host_allocator, total_size, (void**)&allocator); if (iree_status_is_ok(status)) { @@ -37,7 +38,7 @@ allocator->host_allocator = host_allocator; iree_string_view_append_to_buffer( identifier, &allocator->identifier, - (char*)allocator + total_size - identifier.size); + (char*)allocator + iree_sizeof_struct(*allocator)); *out_allocator = (iree_hal_allocator_t*)allocator; }
diff --git a/iree/hal/buffer_heap.c b/iree/hal/buffer_heap.c index 7621667..9d74cc1 100644 --- a/iree/hal/buffer_heap.c +++ b/iree/hal/buffer_heap.c
@@ -32,8 +32,10 @@ IREE_ASSERT_ARGUMENT(out_buffer); IREE_TRACE_ZONE_BEGIN(z0); + // NOTE: we want the buffer data to always be 16-byte aligned. iree_hal_heap_buffer_t* buffer = NULL; - iree_host_size_t header_size = iree_host_align(sizeof(*buffer), 16); + iree_host_size_t header_size = + iree_host_align(iree_sizeof_struct(*buffer), 16); iree_host_size_t total_size = header_size + allocation_size; iree_status_t status = iree_allocator_malloc(host_allocator, total_size, (void**)&buffer);
diff --git a/iree/hal/cuda/BUILD b/iree/hal/cuda/BUILD index 78685b7..6ebe4c9 100644 --- a/iree/hal/cuda/BUILD +++ b/iree/hal/cuda/BUILD
@@ -48,6 +48,8 @@ "nop_executable_cache.h", "status_util.c", "status_util.h", + "stream_command_buffer.c", + "stream_command_buffer.h", ], hdrs = [ "api.h", @@ -59,9 +61,11 @@ "//iree/base:core_headers", "//iree/base:tracing", "//iree/base/internal", + "//iree/base/internal:arena", "//iree/base/internal:flatcc", "//iree/base/internal:synchronization", "//iree/hal", + "//iree/hal/utils:deferred_command_buffer", "//iree/schemas:cuda_executable_def_c_fbs", ], )
diff --git a/iree/hal/cuda/CMakeLists.txt b/iree/hal/cuda/CMakeLists.txt index 997a146..829c381 100644 --- a/iree/hal/cuda/CMakeLists.txt +++ b/iree/hal/cuda/CMakeLists.txt
@@ -45,15 +45,19 @@ "nop_executable_cache.h" "status_util.c" "status_util.h" + "stream_command_buffer.c" + "stream_command_buffer.h" DEPS ::dynamic_symbols iree::base iree::base::core_headers iree::base::internal + iree::base::internal::arena iree::base::internal::flatcc iree::base::internal::synchronization iree::base::tracing iree::hal + iree::hal::utils::deferred_command_buffer iree::schemas::cuda_executable_def_c_fbs PUBLIC )
diff --git a/iree/hal/cuda/api.h b/iree/hal/cuda/api.h index 0a3cd1d..0a7ec62 100644 --- a/iree/hal/cuda/api.h +++ b/iree/hal/cuda/api.h
@@ -16,6 +16,27 @@ extern "C" { #endif // __cplusplus +// Parameters configuring an iree_hal_cuda_device_t. +// Must be initialized with iree_hal_cuda_device_params_initialize prior to use. +typedef struct iree_hal_cuda_device_params_t { + // Number of queues exposed on the device. + // Each queue acts as a separate synchronization scope where all work executes + // concurrently unless prohibited by semaphores. + iree_host_size_t queue_count; + + // Total size of each block in the device shared block pool. + // Larger sizes will lower overhead and ensure the heap isn't hit for + // transient allocations while also increasing memory consumption. + iree_host_size_t arena_block_size; + + // Switch for using deferred command buffer or default graph command buffer + bool use_deferred_submission; +} iree_hal_cuda_device_params_t; + +// Initializes |out_params| to default values. +void iree_hal_cuda_device_params_initialize( + iree_hal_cuda_device_params_t* out_params); + //===----------------------------------------------------------------------===// // iree_hal_cuda_driver_t //===----------------------------------------------------------------------===// @@ -35,6 +56,7 @@ // |out_driver| must be released by the caller (see |iree_hal_driver_release|). IREE_API_EXPORT iree_status_t iree_hal_cuda_driver_create( iree_string_view_t identifier, + const iree_hal_cuda_device_params_t* default_params, const iree_hal_cuda_driver_options_t* options, iree_allocator_t host_allocator, iree_hal_driver_t** out_driver);
diff --git a/iree/hal/cuda/cuda_buffer.c b/iree/hal/cuda/cuda_buffer.c index a87eb24..63334a5 100644 --- a/iree/hal/cuda/cuda_buffer.c +++ b/iree/hal/cuda/cuda_buffer.c
@@ -59,7 +59,7 @@ } IREE_TRACE_ZONE_END(z0); - return iree_ok_status(); + return status; } static void iree_hal_cuda_buffer_destroy(iree_hal_buffer_t* base_buffer) { @@ -95,7 +95,7 @@ // would only work if the entire buffer was discarded. #ifndef NDEBUG if (iree_any_bit_set(memory_access, IREE_HAL_MEMORY_ACCESS_DISCARD)) { - memset(data_ptr + local_byte_offset, 0xCD, local_byte_length); + memset(data_ptr, 0xCD, local_byte_length); } #endif // !NDEBUG *out_data_ptr = data_ptr;
diff --git a/iree/hal/cuda/cuda_device.c b/iree/hal/cuda/cuda_device.c index d78d21f..0239b96 100644 --- a/iree/hal/cuda/cuda_device.c +++ b/iree/hal/cuda/cuda_device.c
@@ -10,6 +10,7 @@ #include <stdint.h> #include <string.h> +#include "iree/base/internal/arena.h" #include "iree/base/tracing.h" #include "iree/hal/cuda/context_wrapper.h" #include "iree/hal/cuda/cuda_allocator.h" @@ -21,6 +22,8 @@ #include "iree/hal/cuda/graph_command_buffer.h" #include "iree/hal/cuda/nop_executable_cache.h" #include "iree/hal/cuda/status_util.h" +#include "iree/hal/cuda/stream_command_buffer.h" +#include "iree/hal/utils/deferred_command_buffer.h" //===----------------------------------------------------------------------===// // iree_hal_cuda_device_t @@ -30,6 +33,10 @@ iree_hal_resource_t resource; iree_string_view_t identifier; + // Block pool used for command buffers with a larger block size (as command + // buffers can contain inlined data uploads). + iree_arena_block_pool_t block_pool; + // Optional driver that owns the CUDA symbols. We retain it for our lifetime // to ensure the symbols remains valid. iree_hal_driver_t* driver; @@ -41,6 +48,8 @@ iree_hal_cuda_context_wrapper_t context_wrapper; iree_hal_allocator_t* device_allocator; + // Switch for using deferred command buffer or default graph command buffer + bool use_deferred_submission; } iree_hal_cuda_device_t; extern const iree_hal_device_vtable_t iree_hal_cuda_device_vtable; @@ -51,6 +60,26 @@ return (iree_hal_cuda_device_t*)base_value; } +void iree_hal_cuda_device_params_initialize( + iree_hal_cuda_device_params_t* out_params) { + out_params->arena_block_size = 32 * 1024; + out_params->queue_count = 8; + out_params->use_deferred_submission = false; +} + +static iree_status_t iree_hal_cuda_device_check_params( + const iree_hal_cuda_device_params_t* params) { + if (params->arena_block_size < 4096) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "arena block size too small (< 4096 bytes)"); + } + if (params->queue_count == 0) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "at least one queue is required"); + } + return iree_ok_status(); +} + static void iree_hal_cuda_device_destroy(iree_hal_device_t* base_device) { iree_hal_cuda_device_t* device = iree_hal_cuda_device_cast(base_device); iree_allocator_t host_allocator = iree_hal_device_host_allocator(base_device); @@ -61,6 +90,7 @@ CUDA_IGNORE_ERROR(device->context_wrapper.syms, cuStreamDestroy(device->stream)); + iree_arena_block_pool_deinitialize(&device->block_pool); // Finally, destroy the device. iree_hal_driver_release(device->driver); @@ -71,25 +101,28 @@ static iree_status_t iree_hal_cuda_device_create_internal( iree_hal_driver_t* driver, iree_string_view_t identifier, - CUdevice cu_device, CUstream stream, CUcontext context, - iree_hal_cuda_dynamic_symbols_t* syms, iree_allocator_t host_allocator, - iree_hal_device_t** out_device) { + const iree_hal_cuda_device_params_t* params, CUdevice cu_device, + CUstream stream, CUcontext context, iree_hal_cuda_dynamic_symbols_t* syms, + iree_allocator_t host_allocator, iree_hal_device_t** out_device) { iree_hal_cuda_device_t* device = NULL; - iree_host_size_t total_size = sizeof(*device) + identifier.size; + iree_host_size_t total_size = iree_sizeof_struct(*device) + identifier.size; IREE_RETURN_IF_ERROR( iree_allocator_malloc(host_allocator, total_size, (void**)&device)); memset(device, 0, total_size); iree_hal_resource_initialize(&iree_hal_cuda_device_vtable, &device->resource); device->driver = driver; iree_hal_driver_retain(device->driver); - uint8_t* buffer_ptr = (uint8_t*)device + sizeof(*device); - buffer_ptr += iree_string_view_append_to_buffer( - identifier, &device->identifier, (char*)buffer_ptr); + iree_string_view_append_to_buffer( + identifier, &device->identifier, + (char*)device + iree_sizeof_struct(*device)); device->device = cu_device; device->stream = stream; device->context_wrapper.cu_context = context; device->context_wrapper.host_allocator = host_allocator; + iree_arena_block_pool_initialize(params->arena_block_size, host_allocator, + &device->block_pool); device->context_wrapper.syms = syms; + device->use_deferred_submission = params->use_deferred_submission; iree_status_t status = iree_hal_cuda_allocator_create( &device->context_wrapper, &device->device_allocator); if (iree_status_is_ok(status)) { @@ -100,13 +133,15 @@ return status; } -iree_status_t iree_hal_cuda_device_create(iree_hal_driver_t* driver, - iree_string_view_t identifier, - iree_hal_cuda_dynamic_symbols_t* syms, - CUdevice device, - iree_allocator_t host_allocator, - iree_hal_device_t** out_device) { +iree_status_t iree_hal_cuda_device_create( + iree_hal_driver_t* driver, iree_string_view_t identifier, + const iree_hal_cuda_device_params_t* params, + iree_hal_cuda_dynamic_symbols_t* syms, CUdevice device, + iree_allocator_t host_allocator, iree_hal_device_t** out_device) { + IREE_ASSERT_ARGUMENT(params); IREE_TRACE_ZONE_BEGIN(z0); + IREE_RETURN_AND_END_ZONE_IF_ERROR(z0, + iree_hal_cuda_device_check_params(params)); CUcontext context; IREE_RETURN_AND_END_ZONE_IF_ERROR( z0, CU_RESULT_TO_STATUS(syms, cuCtxCreate(&context, 0, device))); @@ -115,8 +150,8 @@ syms, cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING)); if (iree_status_is_ok(status)) { - status = iree_hal_cuda_device_create_internal(driver, identifier, device, - stream, context, syms, + status = iree_hal_cuda_device_create_internal(driver, identifier, params, + device, stream, context, syms, host_allocator, out_device); } if (!iree_status_is_ok(status)) { @@ -174,7 +209,12 @@ iree_hal_queue_affinity_t queue_affinity, iree_hal_command_buffer_t** out_command_buffer) { iree_hal_cuda_device_t* device = iree_hal_cuda_device_cast(base_device); - return iree_hal_cuda_graph_command_buffer_allocate( + if (device->use_deferred_submission) { + return iree_hal_deferred_command_buffer_create( + mode, command_categories, &device->block_pool, + iree_hal_device_host_allocator(base_device), out_command_buffer); + } + return iree_hal_cuda_graph_command_buffer_create( &device->context_wrapper, mode, command_categories, queue_affinity, out_command_buffer); } @@ -240,13 +280,30 @@ iree_hal_queue_affinity_t queue_affinity, iree_host_size_t batch_count, const iree_hal_submission_batch_t* batches) { iree_hal_cuda_device_t* device = iree_hal_cuda_device_cast(base_device); - for (int i = 0; i < batch_count; i++) { - for (int j = 0; j < batches[i].command_buffer_count; j++) { - CUgraphExec exec = iree_hal_cuda_graph_command_buffer_exec( - batches[i].command_buffers[j]); - CUDA_RETURN_IF_ERROR(device->context_wrapper.syms, - cuGraphLaunch(exec, device->stream), - "cuGraphLaunch"); + if (device->use_deferred_submission) { + iree_hal_command_buffer_t* stream_command_buffer; + iree_status_t status = iree_hal_cuda_stream_command_buffer_create( + &device->context_wrapper, + IREE_HAL_COMMAND_BUFFER_MODE_ALLOW_INLINE_EXECUTION, command_categories, + device->stream, &stream_command_buffer); + if (iree_status_is_ok(status)) { + for (int i = 0; i < batch_count; i++) { + for (int j = 0; j < batches[i].command_buffer_count; j++) { + iree_hal_deferred_command_buffer_apply(batches[i].command_buffers[j], + stream_command_buffer); + } + } + } + iree_hal_command_buffer_release(stream_command_buffer); + } else { + for (int i = 0; i < batch_count; i++) { + for (int j = 0; j < batches[i].command_buffer_count; j++) { + CUgraphExec exec = iree_hal_cuda_graph_command_buffer_exec( + batches[i].command_buffers[j]); + CUDA_RETURN_IF_ERROR(device->context_wrapper.syms, + cuGraphLaunch(exec, device->stream), + "cuGraphLaunch"); + } } } // TODO(thomasraoux): Conservatively syncronize after every submit until we
diff --git a/iree/hal/cuda/cuda_device.h b/iree/hal/cuda/cuda_device.h index b6b47af..d7b5790 100644 --- a/iree/hal/cuda/cuda_device.h +++ b/iree/hal/cuda/cuda_device.h
@@ -17,12 +17,11 @@ #endif // __cplusplus // Creates a device that owns and manages its own CUcontext. -iree_status_t iree_hal_cuda_device_create(iree_hal_driver_t* driver, - iree_string_view_t identifier, - iree_hal_cuda_dynamic_symbols_t* syms, - CUdevice device, - iree_allocator_t host_allocator, - iree_hal_device_t** out_device); +iree_status_t iree_hal_cuda_device_create( + iree_hal_driver_t* driver, iree_string_view_t identifier, + const iree_hal_cuda_device_params_t* params, + iree_hal_cuda_dynamic_symbols_t* syms, CUdevice device, + iree_allocator_t host_allocator, iree_hal_device_t** out_device); #ifdef __cplusplus } // extern "C"
diff --git a/iree/hal/cuda/cuda_driver.c b/iree/hal/cuda/cuda_driver.c index 6e5617d..8903775 100644 --- a/iree/hal/cuda/cuda_driver.c +++ b/iree/hal/cuda/cuda_driver.c
@@ -22,6 +22,7 @@ // We allow overriding so that multiple CUDA versions can be exposed in the // same process. iree_string_view_t identifier; + iree_hal_cuda_device_params_t default_params; int default_device_index; // CUDA symbols. iree_hal_cuda_dynamic_symbols_t syms; @@ -46,18 +47,23 @@ static iree_status_t iree_hal_cuda_driver_create_internal( iree_string_view_t identifier, + const iree_hal_cuda_device_params_t* default_params, const iree_hal_cuda_driver_options_t* options, iree_allocator_t host_allocator, iree_hal_driver_t** out_driver) { iree_hal_cuda_driver_t* driver = NULL; - iree_host_size_t total_size = sizeof(*driver) + identifier.size; + iree_host_size_t total_size = iree_sizeof_struct(*driver) + identifier.size; IREE_RETURN_IF_ERROR( iree_allocator_malloc(host_allocator, total_size, (void**)&driver)); + iree_hal_resource_initialize(&iree_hal_cuda_driver_vtable, &driver->resource); driver->host_allocator = host_allocator; iree_string_view_append_to_buffer( identifier, &driver->identifier, - (char*)driver + total_size - identifier.size); + (char*)driver + iree_sizeof_struct(*driver)); + memcpy(&driver->default_params, default_params, + sizeof(driver->default_params)); driver->default_device_index = options->default_device_index; + iree_status_t status = iree_hal_cuda_dynamic_symbols_initialize(host_allocator, &driver->syms); if (iree_status_is_ok(status)) { @@ -81,14 +87,16 @@ IREE_API_EXPORT iree_status_t iree_hal_cuda_driver_create( iree_string_view_t identifier, + const iree_hal_cuda_device_params_t* default_params, const iree_hal_cuda_driver_options_t* options, iree_allocator_t host_allocator, iree_hal_driver_t** out_driver) { + IREE_ASSERT_ARGUMENT(default_params); IREE_ASSERT_ARGUMENT(options); IREE_ASSERT_ARGUMENT(out_driver); IREE_TRACE_ZONE_BEGIN(z0); iree_status_t status = iree_hal_cuda_driver_create_internal( - identifier, options, host_allocator, out_driver); + identifier, default_params, options, host_allocator, out_driver); IREE_TRACE_ZONE_END(z0); return status; @@ -193,9 +201,9 @@ iree_string_view_t device_name = iree_make_cstring_view("cuda"); // Attempt to create the device. - iree_status_t status = - iree_hal_cuda_device_create(base_driver, device_name, &driver->syms, - device, host_allocator, out_device); + iree_status_t status = iree_hal_cuda_device_create( + base_driver, device_name, &driver->default_params, &driver->syms, device, + host_allocator, out_device); IREE_TRACE_ZONE_END(z0); return status;
diff --git a/iree/hal/cuda/dynamic_symbol_tables.h b/iree/hal/cuda/dynamic_symbol_tables.h index f131e14..6222219 100644 --- a/iree/hal/cuda/dynamic_symbol_tables.h +++ b/iree/hal/cuda/dynamic_symbol_tables.h
@@ -39,3 +39,8 @@ CU_PFN_DECL(cuStreamDestroy, CUstream) CU_PFN_DECL(cuStreamSynchronize, CUstream) CU_PFN_DECL(cuStreamWaitEvent, CUstream, CUevent, unsigned int) +CU_PFN_DECL(cuMemsetD32Async, unsigned long long, int, size_t, CUstream) +CU_PFN_DECL(cuMemcpyAsync, CUdeviceptr, CUdeviceptr, size_t, CUstream) +CU_PFN_DECL(cuLaunchKernel, CUfunction, unsigned int, unsigned int, + unsigned int, unsigned int, unsigned int, unsigned int, + unsigned int, CUstream, void **, void **)
diff --git a/iree/hal/cuda/graph_command_buffer.c b/iree/hal/cuda/graph_command_buffer.c index 90126e8..f5667fa 100644 --- a/iree/hal/cuda/graph_command_buffer.c +++ b/iree/hal/cuda/graph_command_buffer.c
@@ -47,7 +47,7 @@ return (iree_hal_cuda_graph_command_buffer_t*)base_value; } -iree_status_t iree_hal_cuda_graph_command_buffer_allocate( +iree_status_t iree_hal_cuda_graph_command_buffer_create( iree_hal_cuda_context_wrapper_t* context, iree_hal_command_buffer_mode_t mode, iree_hal_command_category_t command_categories, @@ -136,7 +136,7 @@ static iree_status_t iree_hal_cuda_graph_command_buffer_begin( iree_hal_command_buffer_t* base_command_buffer) { - // nothing to do. + // Nothing to do. return iree_ok_status(); }
diff --git a/iree/hal/cuda/graph_command_buffer.h b/iree/hal/cuda/graph_command_buffer.h index b38b251..c50ccf9 100644 --- a/iree/hal/cuda/graph_command_buffer.h +++ b/iree/hal/cuda/graph_command_buffer.h
@@ -18,7 +18,7 @@ #endif // __cplusplus // Creates a cuda graph. -iree_status_t iree_hal_cuda_graph_command_buffer_allocate( +iree_status_t iree_hal_cuda_graph_command_buffer_create( iree_hal_cuda_context_wrapper_t* context, iree_hal_command_buffer_mode_t mode, iree_hal_command_category_t command_categories,
diff --git a/iree/hal/cuda/registration/driver_module.c b/iree/hal/cuda/registration/driver_module.c index f8fb736..0a51ed9 100644 --- a/iree/hal/cuda/registration/driver_module.c +++ b/iree/hal/cuda/registration/driver_module.c
@@ -41,6 +41,9 @@ driver_id); } IREE_TRACE_ZONE_BEGIN(z0); + iree_hal_cuda_device_params_t default_params; + iree_hal_cuda_device_params_initialize(&default_params); + // TODO(jinchen62): set up default_params.use_deferred_submission by flag // When we expose more than one driver (different cuda versions, etc) we // can name them here: iree_string_view_t identifier = iree_make_cstring_view("cuda"); @@ -48,7 +51,7 @@ iree_hal_cuda_driver_options_t driver_options; iree_hal_cuda_driver_options_initialize(&driver_options); iree_status_t status = iree_hal_cuda_driver_create( - identifier, &driver_options, allocator, out_driver); + identifier, &default_params, &driver_options, allocator, out_driver); IREE_TRACE_ZONE_END(z0); return status; }
diff --git a/iree/hal/cuda/stream_command_buffer.c b/iree/hal/cuda/stream_command_buffer.c new file mode 100644 index 0000000..9ab153d --- /dev/null +++ b/iree/hal/cuda/stream_command_buffer.c
@@ -0,0 +1,353 @@ +// Copyright 2021 The IREE Authors +// +// 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 + +#include "iree/hal/cuda/stream_command_buffer.h" + +#include "iree/base/tracing.h" +#include "iree/hal/cuda/cuda_buffer.h" +#include "iree/hal/cuda/cuda_event.h" +#include "iree/hal/cuda/executable_layout.h" +#include "iree/hal/cuda/native_executable.h" +#include "iree/hal/cuda/status_util.h" + +#define IREE_HAL_CUDA_MAX_BINDING_COUNT 64 + +// This records the commands on the calling thread without additional threading +// indirection. + +typedef struct { + iree_hal_resource_t resource; + iree_hal_cuda_context_wrapper_t* context; + iree_hal_command_buffer_mode_t mode; + iree_hal_command_category_t allowed_categories; + CUstream stream; + // Keep track of the current set of kernel arguments. + void* current_descriptor[IREE_HAL_CUDA_MAX_BINDING_COUNT]; + CUdeviceptr* device_ptrs[IREE_HAL_CUDA_MAX_BINDING_COUNT]; +} iree_hal_cuda_stream_command_buffer_t; + +extern const iree_hal_command_buffer_vtable_t + iree_hal_cuda_stream_command_buffer_vtable; + +static iree_hal_cuda_stream_command_buffer_t* +iree_hal_cuda_stream_command_buffer_cast( + iree_hal_command_buffer_t* base_value) { + IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_cuda_stream_command_buffer_vtable); + return (iree_hal_cuda_stream_command_buffer_t*)base_value; +} + +iree_status_t iree_hal_cuda_stream_command_buffer_create( + iree_hal_cuda_context_wrapper_t* context, + iree_hal_command_buffer_mode_t mode, + iree_hal_command_category_t command_categories, CUstream stream, + iree_hal_command_buffer_t** out_command_buffer) { + IREE_ASSERT_ARGUMENT(context); + IREE_ASSERT_ARGUMENT(out_command_buffer); + *out_command_buffer = NULL; + IREE_TRACE_ZONE_BEGIN(z0); + + iree_hal_cuda_stream_command_buffer_t* command_buffer = NULL; + size_t total_size = sizeof(*command_buffer) + + IREE_HAL_CUDA_MAX_BINDING_COUNT * sizeof(void*) + + IREE_HAL_CUDA_MAX_BINDING_COUNT * sizeof(CUdeviceptr); + iree_status_t status = iree_allocator_malloc( + context->host_allocator, total_size, (void**)&command_buffer); + if (iree_status_is_ok(status)) { + iree_hal_resource_initialize(&iree_hal_cuda_stream_command_buffer_vtable, + &command_buffer->resource); + command_buffer->context = context; + command_buffer->mode = mode; + command_buffer->allowed_categories = command_categories; + command_buffer->stream = stream; + for (size_t i = 0; i < IREE_HAL_CUDA_MAX_BINDING_COUNT; i++) { + command_buffer->current_descriptor[i] = &command_buffer->device_ptrs[i]; + } + } + + *out_command_buffer = (iree_hal_command_buffer_t*)command_buffer; + IREE_TRACE_ZONE_END(z0); + return status; +} + +static void iree_hal_cuda_stream_command_buffer_destroy( + iree_hal_command_buffer_t* base_command_buffer) { + iree_hal_cuda_stream_command_buffer_t* command_buffer = + iree_hal_cuda_stream_command_buffer_cast(base_command_buffer); + IREE_TRACE_ZONE_BEGIN(z0); + + iree_allocator_free(command_buffer->context->host_allocator, command_buffer); + + IREE_TRACE_ZONE_END(z0); +} + +static iree_hal_command_buffer_mode_t iree_hal_cuda_stream_command_buffer_mode( + const iree_hal_command_buffer_t* base_command_buffer) { + const iree_hal_cuda_stream_command_buffer_t* command_buffer = + (const iree_hal_cuda_stream_command_buffer_t*)(base_command_buffer); + return command_buffer->mode; +} + +static iree_hal_command_category_t +iree_hal_cuda_stream_command_buffer_allowed_categories( + const iree_hal_command_buffer_t* base_command_buffer) { + const iree_hal_cuda_stream_command_buffer_t* command_buffer = + (const iree_hal_cuda_stream_command_buffer_t*)(base_command_buffer); + return command_buffer->allowed_categories; +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_begin( + iree_hal_command_buffer_t* base_command_buffer) { + return iree_ok_status(); +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_end( + iree_hal_command_buffer_t* base_command_buffer) { + return iree_ok_status(); +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_execution_barrier( + iree_hal_command_buffer_t* base_command_buffer, + iree_hal_execution_stage_t source_stage_mask, + iree_hal_execution_stage_t target_stage_mask, + iree_hal_execution_barrier_flags_t flags, + iree_host_size_t memory_barrier_count, + const iree_hal_memory_barrier_t* memory_barriers, + iree_host_size_t buffer_barrier_count, + const iree_hal_buffer_barrier_t* buffer_barriers) { + // TODO(jinchen62): implement CUDA barrier + return iree_ok_status(); +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_signal_event( + iree_hal_command_buffer_t* base_command_buffer, iree_hal_event_t* event, + iree_hal_execution_stage_t source_stage_mask) { + // TODO(jinchen62): implement CUDA barrier + return iree_ok_status(); +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_reset_event( + iree_hal_command_buffer_t* base_command_buffer, iree_hal_event_t* event, + iree_hal_execution_stage_t source_stage_mask) { + // TODO(jinchen62): implement CUDA barrier + return iree_ok_status(); +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_wait_events( + iree_hal_command_buffer_t* base_command_buffer, + iree_host_size_t event_count, const iree_hal_event_t** events, + iree_hal_execution_stage_t source_stage_mask, + iree_hal_execution_stage_t target_stage_mask, + iree_host_size_t memory_barrier_count, + const iree_hal_memory_barrier_t* memory_barriers, + iree_host_size_t buffer_barrier_count, + const iree_hal_buffer_barrier_t* buffer_barriers) { + // TODO(jinchen62): implement CUDA barrier + return iree_ok_status(); +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_discard_buffer( + iree_hal_command_buffer_t* base_command_buffer, iree_hal_buffer_t* buffer) { + // nothing to do. + return iree_ok_status(); +} + +// Splats a pattern value of 1, 2, or 4 bytes out to a 4 byte value. +static uint32_t iree_hal_cuda_splat_pattern(const void* pattern, + size_t pattern_length) { + switch (pattern_length) { + case 1: { + uint32_t pattern_value = *(const uint8_t*)(pattern); + return (pattern_value << 24) | (pattern_value << 16) | + (pattern_value << 8) | pattern_value; + } + case 2: { + uint32_t pattern_value = *(const uint16_t*)(pattern); + return (pattern_value << 16) | pattern_value; + } + case 4: { + uint32_t pattern_value = *(const uint32_t*)(pattern); + return pattern_value; + } + default: + return 0; // Already verified that this should not be possible. + } +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_fill_buffer( + iree_hal_command_buffer_t* base_command_buffer, + iree_hal_buffer_t* target_buffer, iree_device_size_t target_offset, + iree_device_size_t length, const void* pattern, + iree_host_size_t pattern_length) { + iree_hal_cuda_stream_command_buffer_t* command_buffer = + iree_hal_cuda_stream_command_buffer_cast(base_command_buffer); + + CUdeviceptr target_device_buffer = iree_hal_cuda_buffer_device_pointer( + iree_hal_buffer_allocated_buffer(target_buffer)); + target_offset += iree_hal_buffer_byte_offset(target_buffer); + uint32_t dword_pattern = iree_hal_cuda_splat_pattern(pattern, pattern_length); + CUdeviceptr dst = target_device_buffer + target_offset; + int value = dword_pattern; + size_t sizeBytes = length; + CUDA_RETURN_IF_ERROR( + command_buffer->context->syms, + cuMemsetD32Async(dst, value, sizeBytes, command_buffer->stream), + "cuMemsetD32Async"); + return iree_ok_status(); +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_update_buffer( + iree_hal_command_buffer_t* base_command_buffer, const void* source_buffer, + iree_host_size_t source_offset, iree_hal_buffer_t* target_buffer, + iree_device_size_t target_offset, iree_device_size_t length) { + return iree_make_status(IREE_STATUS_UNIMPLEMENTED, + "need cuda implementation of update buffer"); +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_copy_buffer( + iree_hal_command_buffer_t* base_command_buffer, + iree_hal_buffer_t* source_buffer, iree_device_size_t source_offset, + iree_hal_buffer_t* target_buffer, iree_device_size_t target_offset, + iree_device_size_t length) { + iree_hal_cuda_stream_command_buffer_t* command_buffer = + iree_hal_cuda_stream_command_buffer_cast(base_command_buffer); + + CUdeviceptr target_device_buffer = iree_hal_cuda_buffer_device_pointer( + iree_hal_buffer_allocated_buffer(target_buffer)); + target_offset += iree_hal_buffer_byte_offset(target_buffer); + CUdeviceptr source_device_buffer = iree_hal_cuda_buffer_device_pointer( + iree_hal_buffer_allocated_buffer(source_buffer)); + source_offset += iree_hal_buffer_byte_offset(source_buffer); + CUDA_RETURN_IF_ERROR(command_buffer->context->syms, + cuMemcpyAsync(target_device_buffer, source_device_buffer, + length, command_buffer->stream), + "cuMemcpyAsync"); + return iree_ok_status(); +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_push_constants( + iree_hal_command_buffer_t* base_command_buffer, + iree_hal_executable_layout_t* executable_layout, iree_host_size_t offset, + const void* values, iree_host_size_t values_length) { + return iree_make_status(IREE_STATUS_UNIMPLEMENTED, + "need cuda implementation of push constants"); +} + +// Tie together the binding index and its index in |bindings| array. +typedef struct { + uint32_t index; + uint32_t binding; +} iree_hal_cuda_binding_mapping_t; + +// Helper to sort the binding based on their binding index. +static int compare_binding_index(const void* a, const void* b) { + const iree_hal_cuda_binding_mapping_t buffer_a = + *(const iree_hal_cuda_binding_mapping_t*)a; + const iree_hal_cuda_binding_mapping_t buffer_b = + *(const iree_hal_cuda_binding_mapping_t*)b; + return buffer_a.binding < buffer_b.binding ? -1 : 1; +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_push_descriptor_set( + iree_hal_command_buffer_t* base_command_buffer, + iree_hal_executable_layout_t* executable_layout, uint32_t set, + iree_host_size_t binding_count, + const iree_hal_descriptor_set_binding_t* bindings) { + iree_hal_cuda_stream_command_buffer_t* command_buffer = + iree_hal_cuda_stream_command_buffer_cast(base_command_buffer); + iree_host_size_t base_binding = + iree_hal_cuda_base_binding_index(executable_layout, set); + // Convention with the compiler side. We map bindings to kernel argument. + // We compact the bindings to get a dense set of arguments and keep them order + // based on the binding index. + // Sort the binding based on the binding index and map the array index to the + // argument index. + iree_hal_cuda_binding_mapping_t binding_used[IREE_HAL_CUDA_MAX_BINDING_COUNT]; + for (iree_host_size_t i = 0; i < binding_count; i++) { + iree_hal_cuda_binding_mapping_t buffer = {i, bindings[i].binding}; + binding_used[i] = buffer; + } + qsort(binding_used, binding_count, sizeof(iree_hal_cuda_binding_mapping_t), + compare_binding_index); + assert(binding_count < IREE_HAL_CUDA_MAX_BINDING_COUNT && + "binding count larger than the max expected."); + for (iree_host_size_t i = 0; i < binding_count; i++) { + iree_hal_descriptor_set_binding_t binding = bindings[binding_used[i].index]; + CUdeviceptr device_ptr = + iree_hal_cuda_buffer_device_pointer( + iree_hal_buffer_allocated_buffer(binding.buffer)) + + iree_hal_buffer_byte_offset(binding.buffer) + binding.offset; + *((CUdeviceptr*)command_buffer->current_descriptor[i + base_binding]) = + device_ptr; + } + return iree_ok_status(); +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_bind_descriptor_set( + iree_hal_command_buffer_t* base_command_buffer, + iree_hal_executable_layout_t* executable_layout, uint32_t set, + iree_hal_descriptor_set_t* descriptor_set, + iree_host_size_t dynamic_offset_count, + const iree_device_size_t* dynamic_offsets) { + return iree_make_status(IREE_STATUS_UNIMPLEMENTED, + "need cuda implementation of bind descriptor set"); +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_dispatch( + iree_hal_command_buffer_t* base_command_buffer, + iree_hal_executable_t* executable, int32_t entry_point, + uint32_t workgroup_x, uint32_t workgroup_y, uint32_t workgroup_z) { + iree_hal_cuda_stream_command_buffer_t* command_buffer = + iree_hal_cuda_stream_command_buffer_cast(base_command_buffer); + + int32_t block_size_x, block_size_y, block_size_z; + IREE_RETURN_IF_ERROR(iree_hal_cuda_native_executable_block_size( + executable, entry_point, &block_size_x, &block_size_y, &block_size_z)); + CUfunction func = + iree_hal_cuda_native_executable_for_entry_point(executable, entry_point); + CUDA_RETURN_IF_ERROR( + command_buffer->context->syms, + cuLaunchKernel(func, workgroup_x, workgroup_y, workgroup_z, block_size_x, + block_size_y, block_size_z, 0, command_buffer->stream, + command_buffer->current_descriptor, NULL), + "cuLaunchKernel"); + return iree_ok_status(); +} + +static iree_status_t iree_hal_cuda_stream_command_buffer_dispatch_indirect( + iree_hal_command_buffer_t* base_command_buffer, + iree_hal_executable_t* executable, int32_t entry_point, + iree_hal_buffer_t* workgroups_buffer, + iree_device_size_t workgroups_offset) { + return iree_make_status(IREE_STATUS_UNIMPLEMENTED, + "need cuda implementation of dispatch indirect"); +} + +const iree_hal_command_buffer_vtable_t + iree_hal_cuda_stream_command_buffer_vtable = { + .destroy = iree_hal_cuda_stream_command_buffer_destroy, + .mode = iree_hal_cuda_stream_command_buffer_mode, + .allowed_categories = + iree_hal_cuda_stream_command_buffer_allowed_categories, + .begin = iree_hal_cuda_stream_command_buffer_begin, + .end = iree_hal_cuda_stream_command_buffer_end, + .execution_barrier = + iree_hal_cuda_stream_command_buffer_execution_barrier, + .signal_event = iree_hal_cuda_stream_command_buffer_signal_event, + .reset_event = iree_hal_cuda_stream_command_buffer_reset_event, + .wait_events = iree_hal_cuda_stream_command_buffer_wait_events, + .discard_buffer = iree_hal_cuda_stream_command_buffer_discard_buffer, + .fill_buffer = iree_hal_cuda_stream_command_buffer_fill_buffer, + .update_buffer = iree_hal_cuda_stream_command_buffer_update_buffer, + .copy_buffer = iree_hal_cuda_stream_command_buffer_copy_buffer, + .push_constants = iree_hal_cuda_stream_command_buffer_push_constants, + .push_descriptor_set = + iree_hal_cuda_stream_command_buffer_push_descriptor_set, + .bind_descriptor_set = + iree_hal_cuda_stream_command_buffer_bind_descriptor_set, + .dispatch = iree_hal_cuda_stream_command_buffer_dispatch, + .dispatch_indirect = + iree_hal_cuda_stream_command_buffer_dispatch_indirect, +};
diff --git a/iree/hal/cuda/stream_command_buffer.h b/iree/hal/cuda/stream_command_buffer.h new file mode 100644 index 0000000..b4b901a --- /dev/null +++ b/iree/hal/cuda/stream_command_buffer.h
@@ -0,0 +1,35 @@ +// Copyright 2021 The IREE Authors +// +// 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 + +#ifndef IREE_HAL_CUDA_STREAM_COMMAND_BUFFER_H_ +#define IREE_HAL_CUDA_STREAM_COMMAND_BUFFER_H_ + +#include "iree/base/internal/arena.h" +#include "iree/hal/api.h" +#include "iree/hal/cuda/context_wrapper.h" +#include "iree/hal/cuda/cuda_headers.h" +#include "iree/hal/cuda/dynamic_symbols.h" + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +// Creates a cuda stream command buffer that immediately +// issues commands against the given |stream|. +// Access to |stream| must be synchronized by the user. +// Used for replaying commands in special situations and +// never returned to a user from the device_create_command_buffer +iree_status_t iree_hal_cuda_stream_command_buffer_create( + iree_hal_cuda_context_wrapper_t *context, + iree_hal_command_buffer_mode_t mode, + iree_hal_command_category_t command_categories, CUstream stream, + iree_hal_command_buffer_t **out_command_buffer); + +#ifdef __cplusplus +} // extern "C" +#endif // __cplusplus + +#endif // IREE_HAL_CUDA_STREAM_COMMAND_BUFFER_H_
diff --git a/iree/hal/driver_registry.c b/iree/hal/driver_registry.c index 972bff9..8594023 100644 --- a/iree/hal/driver_registry.c +++ b/iree/hal/driver_registry.c
@@ -188,8 +188,8 @@ // Allocate the required memory for both the driver infos and the string // storage in a single block. - iree_host_size_t total_driver_infos_size = iree_host_align( - total_driver_info_count * sizeof(iree_hal_driver_info_t), 8); + iree_host_size_t total_driver_infos_size = + total_driver_info_count * sizeof(iree_hal_driver_info_t); if (iree_status_is_ok(status)) { status = iree_allocator_malloc(allocator, total_driver_infos_size + total_storage_size,
diff --git a/iree/hal/local/elf/arch/x86_32.c b/iree/hal/local/elf/arch/x86_32.c index 6a7c103..05d08d7 100644 --- a/iree/hal/local/elf/arch/x86_32.c +++ b/iree/hal/local/elf/arch/x86_32.c
@@ -120,10 +120,10 @@ // Cross-ABI function calls //============================================================================== -// System V 386 ABI (used in IREE): +// System V i386 ABI (used in IREE): // https://uclibc.org/docs/psABI-i386.pdf // Arguments: -// +// (reverse order on the stack; last arg furthest from stack pointer) // // Results: // EAX @@ -136,7 +136,7 @@ #if defined(IREE_PLATFORM_WINDOWS) -#error "TODO" +#error "TODO(#6554): need cdecl -> sysv ABI shims in x86_32_msvc.asm" #else
diff --git a/iree/hal/local/loaders/embedded_library_loader.c b/iree/hal/local/loaders/embedded_library_loader.c index ddaafe2..74d48fe 100644 --- a/iree/hal/local/loaders/embedded_library_loader.c +++ b/iree/hal/local/loaders/embedded_library_loader.c
@@ -35,6 +35,8 @@ const iree_hal_executable_library_header_t** header; const iree_hal_executable_library_v0_t* v0; } library; + + iree_hal_local_executable_layout_t* layouts[]; } iree_hal_elf_executable_t; extern const iree_hal_local_executable_vtable_t iree_hal_elf_executable_vtable; @@ -140,16 +142,13 @@ iree_hal_elf_executable_t* executable = NULL; iree_host_size_t total_size = sizeof(*executable) + - executable_layout_count * sizeof(iree_hal_local_executable_layout_t); + executable_layout_count * sizeof(*executable->layouts); iree_status_t status = iree_allocator_malloc(host_allocator, total_size, (void**)&executable); if (iree_status_is_ok(status)) { - iree_hal_local_executable_layout_t** executable_layouts_ptr = - (iree_hal_local_executable_layout_t**)(((uint8_t*)executable) + - sizeof(*executable)); iree_hal_local_executable_initialize( &iree_hal_elf_executable_vtable, executable_layout_count, - executable_layouts, executable_layouts_ptr, host_allocator, + executable_layouts, &executable->layouts[0], host_allocator, &executable->base); } if (iree_status_is_ok(status)) {
diff --git a/iree/hal/local/loaders/legacy_library_loader.c b/iree/hal/local/loaders/legacy_library_loader.c index f370b94..fa9fa1f 100644 --- a/iree/hal/local/loaders/legacy_library_loader.c +++ b/iree/hal/local/loaders/legacy_library_loader.c
@@ -81,6 +81,8 @@ const iree_hal_executable_library_header_t** header; const iree_hal_executable_library_v0_t* v0; } library; + + iree_hal_local_executable_layout_t* layouts[]; } iree_hal_legacy_executable_t; extern const iree_hal_local_executable_vtable_t @@ -224,16 +226,13 @@ iree_hal_legacy_executable_t* executable = NULL; iree_host_size_t total_size = sizeof(*executable) + - executable_layout_count * sizeof(iree_hal_local_executable_layout_t); + executable_layout_count * sizeof(*executable->layouts); iree_status_t status = iree_allocator_malloc(host_allocator, total_size, (void**)&executable); if (iree_status_is_ok(status)) { - iree_hal_local_executable_layout_t** executable_layouts_ptr = - (iree_hal_local_executable_layout_t**)(((uint8_t*)executable) + - sizeof(*executable)); iree_hal_local_executable_initialize( &iree_hal_legacy_executable_vtable, executable_layout_count, - executable_layouts, executable_layouts_ptr, host_allocator, + executable_layouts, &executable->layouts[0], host_allocator, &executable->base); executable->def = executable_def; }
diff --git a/iree/hal/local/loaders/static_library_loader.c b/iree/hal/local/loaders/static_library_loader.c index a7ff63b..1e7819c 100644 --- a/iree/hal/local/loaders/static_library_loader.c +++ b/iree/hal/local/loaders/static_library_loader.c
@@ -30,6 +30,8 @@ const iree_hal_executable_library_header_t** header; const iree_hal_executable_library_v0_t* v0; } library; + + iree_hal_local_executable_layout_t* layouts[]; } iree_hal_static_executable_t; static const iree_hal_local_executable_vtable_t @@ -50,16 +52,13 @@ iree_hal_static_executable_t* executable = NULL; iree_host_size_t total_size = sizeof(*executable) + - executable_layout_count * sizeof(iree_hal_local_executable_layout_t); + executable_layout_count * sizeof(*executable->layouts); iree_status_t status = iree_allocator_malloc(host_allocator, total_size, (void**)&executable); if (iree_status_is_ok(status)) { - iree_hal_local_executable_layout_t** executable_layouts_ptr = - (iree_hal_local_executable_layout_t**)(((uint8_t*)executable) + - sizeof(*executable)); iree_hal_local_executable_initialize( &iree_hal_static_executable_vtable, executable_layout_count, - executable_layouts, executable_layouts_ptr, host_allocator, + executable_layouts, &executable->layouts[0], host_allocator, &executable->base); executable->library.header = library_header; executable->identifier = iree_make_cstring_view((*library_header)->name);
diff --git a/iree/hal/local/sync_device.c b/iree/hal/local/sync_device.c index b6083ab..924da83 100644 --- a/iree/hal/local/sync_device.c +++ b/iree/hal/local/sync_device.c
@@ -23,13 +23,13 @@ iree_hal_resource_t resource; iree_string_view_t identifier; - iree_host_size_t loader_count; - iree_hal_executable_loader_t** loaders; - iree_allocator_t host_allocator; iree_hal_allocator_t* device_allocator; iree_hal_sync_semaphore_state_t semaphore_state; + + iree_host_size_t loader_count; + iree_hal_executable_loader_t* loaders[]; } iree_hal_sync_device_t; static const iree_hal_device_vtable_t iree_hal_sync_device_vtable; @@ -64,8 +64,9 @@ iree_hal_sync_device_check_params(params)); iree_hal_sync_device_t* device = NULL; - iree_host_size_t total_size = sizeof(*device) + identifier.size + - loader_count * sizeof(*device->loaders); + iree_host_size_t struct_size = + sizeof(*device) + loader_count * sizeof(*device->loaders); + iree_host_size_t total_size = struct_size + identifier.size; iree_status_t status = iree_allocator_malloc(host_allocator, total_size, (void**)&device); if (iree_status_is_ok(status)) { @@ -73,13 +74,10 @@ iree_hal_resource_initialize(&iree_hal_sync_device_vtable, &device->resource); iree_string_view_append_to_buffer(identifier, &device->identifier, - (char*)device + sizeof(*device)); + (char*)device + struct_size); device->host_allocator = host_allocator; device->loader_count = loader_count; - device->loaders = - (iree_hal_executable_loader_t**)((uint8_t*)device->identifier.data + - identifier.size); for (iree_host_size_t i = 0; i < device->loader_count; ++i) { device->loaders[i] = loaders[i]; iree_hal_executable_loader_retain(device->loaders[i]);
diff --git a/iree/hal/local/task_device.c b/iree/hal/local/task_device.c index 158b9e4..236c83e 100644 --- a/iree/hal/local/task_device.c +++ b/iree/hal/local/task_device.c
@@ -92,19 +92,18 @@ iree_hal_task_device_check_params(params)); iree_hal_task_device_t* device = NULL; - iree_host_size_t total_size = - sizeof(*device) + params->queue_count * sizeof(*device->queues) + - identifier.size + loader_count * sizeof(*device->loaders); + iree_host_size_t struct_size = sizeof(*device) + + params->queue_count * sizeof(*device->queues) + + loader_count * sizeof(*device->loaders); + iree_host_size_t total_size = struct_size + identifier.size; iree_status_t status = iree_allocator_malloc(host_allocator, total_size, (void**)&device); if (iree_status_is_ok(status)) { memset(device, 0, total_size); iree_hal_resource_initialize(&iree_hal_task_device_vtable, &device->resource); - iree_string_view_append_to_buffer( - identifier, &device->identifier, - (char*)device + sizeof(*device) + - params->queue_count * sizeof(*device->queues)); + iree_string_view_append_to_buffer(identifier, &device->identifier, + (char*)device + struct_size); device->host_allocator = host_allocator; iree_arena_block_pool_initialize(4096, host_allocator, &device->small_block_pool); @@ -117,8 +116,9 @@ device->loader_count = loader_count; device->loaders = - (iree_hal_executable_loader_t**)((uint8_t*)device->identifier.data + - identifier.size); + (iree_hal_executable_loader_t**)((uint8_t*)device + sizeof(*device) + + params->queue_count * + sizeof(*device->queues)); for (iree_host_size_t i = 0; i < device->loader_count; ++i) { device->loaders[i] = loaders[i]; iree_hal_executable_loader_retain(device->loaders[i]);
diff --git a/iree/hal/local/task_driver.c b/iree/hal/local/task_driver.c index e3f10ed..7711520 100644 --- a/iree/hal/local/task_driver.c +++ b/iree/hal/local/task_driver.c
@@ -47,9 +47,9 @@ IREE_TRACE_ZONE_BEGIN(z0); iree_hal_task_driver_t* driver = NULL; - iree_host_size_t total_size = sizeof(*driver) + - loader_count * sizeof(*driver->loaders) + - identifier.size; + iree_host_size_t struct_size = + sizeof(*driver) + loader_count * sizeof(*driver->loaders); + iree_host_size_t total_size = struct_size + identifier.size; iree_status_t status = iree_allocator_malloc(host_allocator, total_size, (void**)&driver); if (iree_status_is_ok(status)) { @@ -57,9 +57,8 @@ &driver->resource); driver->host_allocator = host_allocator; - iree_string_view_append_to_buffer( - identifier, &driver->identifier, - (char*)driver + total_size - identifier.size); + iree_string_view_append_to_buffer(identifier, &driver->identifier, + (char*)driver + struct_size); memcpy(&driver->default_params, default_params, sizeof(driver->default_params));
diff --git a/iree/hal/vulkan/native_descriptor_set_layout.cc b/iree/hal/vulkan/native_descriptor_set_layout.cc index 1d3c4a3..dff1661 100644 --- a/iree/hal/vulkan/native_descriptor_set_layout.cc +++ b/iree/hal/vulkan/native_descriptor_set_layout.cc
@@ -56,6 +56,7 @@ VkDescriptorSetLayoutBinding* native_bindings = NULL; if (binding_count > 0) { + // TODO(benvanik): avoid this allocation if possible (inline_array). IREE_RETURN_IF_ERROR(iree_allocator_malloc( logical_device->host_allocator(), binding_count * sizeof(VkDescriptorSetLayoutBinding),
diff --git a/iree/hal/vulkan/vma_buffer.cc b/iree/hal/vulkan/vma_buffer.cc index 9030775..ffdb8ff 100644 --- a/iree/hal/vulkan/vma_buffer.cc +++ b/iree/hal/vulkan/vma_buffer.cc
@@ -123,7 +123,7 @@ // would only work if the entire buffer was discarded. #ifndef NDEBUG if (iree_any_bit_set(memory_access, IREE_HAL_MEMORY_ACCESS_DISCARD)) { - memset(data_ptr + local_byte_offset, 0xCD, local_byte_length); + memset(*out_data_ptr, 0xCD, local_byte_length); } #endif // !NDEBUG
diff --git a/iree/samples/custom_modules/module.cc b/iree/samples/custom_modules/module.cc index 0477ef5..1759d82 100644 --- a/iree/samples/custom_modules/module.cc +++ b/iree/samples/custom_modules/module.cc
@@ -56,7 +56,7 @@ // Note that we allocate the message and the string value together. iree_custom_message_t* message = NULL; IREE_RETURN_IF_ERROR(iree_allocator_malloc( - allocator, sizeof(iree_custom_message_t) + value.size, (void**)&message)); + allocator, sizeof(*message) + value.size, (void**)&message)); message->ref_object.counter = IREE_ATOMIC_VAR_INIT(1); message->allocator = allocator; message->value.data = ((const char*)message) + sizeof(iree_custom_message_t); @@ -71,8 +71,8 @@ iree_custom_message_t** out_message) { IREE_ASSERT_ARGUMENT(out_message); iree_custom_message_t* message = NULL; - IREE_RETURN_IF_ERROR(iree_allocator_malloc( - allocator, sizeof(iree_custom_message_t), (void**)&message)); + IREE_RETURN_IF_ERROR( + iree_allocator_malloc(allocator, sizeof(*message), (void**)&message)); message->ref_object.counter = IREE_ATOMIC_VAR_INIT(1); message->allocator = allocator; message->value = value; // Unowned.
diff --git a/iree/samples/simple_embedding/BUILD b/iree/samples/simple_embedding/BUILD index 94c29c2..d2fd399 100644 --- a/iree/samples/simple_embedding/BUILD +++ b/iree/samples/simple_embedding/BUILD
@@ -265,6 +265,53 @@ iree_cmake_extra_content( content = """ endif() + +if(${IREE_HAL_DRIVER_CUDA} AND (${IREE_TARGET_BACKEND_CUDA} OR DEFINED IREE_HOST_BINARY_ROOT)) +""", + inline = True, +) + +cc_binary( + name = "simple_embedding_cuda", + srcs = [ + "device_cuda.c", + "simple_embedding.c", + ], + deps = [ + ":simple_embedding_test_bytecode_module_cuda_c", + "//iree/base", + "//iree/hal", + "//iree/hal/cuda/registration", + "//iree/modules/hal", + "//iree/vm", + "//iree/vm:bytecode_module", + ], +) + +iree_bytecode_module( + name = "simple_embedding_test_bytecode_module_cuda", + src = "simple_embedding_test.mlir", + c_identifier = "iree_samples_simple_embedding_test_module_cuda", + flags = [ + "-iree-input-type=mhlo", + "-iree-mlir-to-vm-bytecode-module", + "-iree-hal-target-backends=cuda", + "-iree-llvm-debug-symbols=false", + ], +) + +# Simple embedding is failing in the CI. +# run_binary_test( +# name = "simple_embedding_cuda_test", +# tags = [ +# "driver=cuda", +# ], +# test_binary = ":simple_embedding_cuda", +# ) + +iree_cmake_extra_content( + content = """ +endif() """, inline = True, )
diff --git a/iree/samples/simple_embedding/CMakeLists.txt b/iree/samples/simple_embedding/CMakeLists.txt index 4623869..aaee677 100644 --- a/iree/samples/simple_embedding/CMakeLists.txt +++ b/iree/samples/simple_embedding/CMakeLists.txt
@@ -245,4 +245,48 @@ endif() +if(${IREE_HAL_DRIVER_CUDA} AND (${IREE_TARGET_BACKEND_CUDA} OR DEFINED IREE_HOST_BINARY_ROOT)) + +iree_cc_binary( + NAME + simple_embedding_cuda + SRCS + "device_cuda.c" + "simple_embedding.c" + DEPS + ::simple_embedding_test_bytecode_module_cuda_c + iree::base + iree::hal + iree::hal::cuda::registration + iree::modules::hal + iree::vm + iree::vm::bytecode_module +) + +iree_bytecode_module( + NAME + simple_embedding_test_bytecode_module_cuda + SRC + "simple_embedding_test.mlir" + C_IDENTIFIER + "iree_samples_simple_embedding_test_module_cuda" + FLAGS + "-iree-input-type=mhlo" + "-iree-mlir-to-vm-bytecode-module" + "-iree-hal-target-backends=cuda" + "-iree-llvm-debug-symbols=false" + PUBLIC +) + +iree_run_binary_test( + NAME + "simple_embedding_cuda_test" + LABELS + "driver=cuda" + TEST_BINARY + ::simple_embedding_cuda +) + +endif() + ### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/samples/simple_embedding/device_cuda.c b/iree/samples/simple_embedding/device_cuda.c new file mode 100644 index 0000000..e8cf2f4 --- /dev/null +++ b/iree/samples/simple_embedding/device_cuda.c
@@ -0,0 +1,39 @@ +// Copyright 2021 The IREE Authors +// +// 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 + +// A example of setting up the the cuda driver. + +#include <stddef.h> + +#include "iree/base/api.h" +#include "iree/hal/api.h" +#include "iree/hal/cuda/registration/driver_module.h" + +// Compiled module embedded here to avoid file IO: +#include "iree/samples/simple_embedding/simple_embedding_test_bytecode_module_cuda_c.h" + +iree_status_t create_sample_device(iree_hal_device_t** device) { + // Only register the cuda HAL driver. + IREE_RETURN_IF_ERROR( + iree_hal_cuda_driver_module_register(iree_hal_driver_registry_default())); + // Create the hal driver from the name. + iree_hal_driver_t* driver = NULL; + iree_string_view_t identifier = iree_make_cstring_view("cuda"); + IREE_RETURN_IF_ERROR(iree_hal_driver_registry_try_create_by_name( + iree_hal_driver_registry_default(), identifier, iree_allocator_system(), + &driver)); + IREE_RETURN_IF_ERROR(iree_hal_driver_create_default_device( + driver, iree_allocator_system(), device)); + iree_hal_driver_release(driver); + return iree_ok_status(); +} + +const iree_const_byte_span_t load_bytecode_module_data() { + const struct iree_file_toc_t* module_file_toc = + iree_samples_simple_embedding_test_module_cuda_create(); + return iree_make_const_byte_span(module_file_toc->data, + module_file_toc->size); +}
diff --git a/iree/tools/run_lit.ps1 b/iree/tools/run_lit.ps1 index 7f19601..e88dcb4 100644 --- a/iree/tools/run_lit.ps1 +++ b/iree/tools/run_lit.ps1
@@ -63,6 +63,7 @@ } $test_line = $test_line.Substring("// RUN: ".Length) $test_line = $test_line -replace "%s", $test_file + $test_line = $test_line -replace "`"", "\`"" Write-Host -ForegroundColor Blue "Running test command:" Write-Host -ForegroundColor Yellow "$test_line" & $bashExe -c $test_line | Out-Default
diff --git a/iree/vm/buffer.c b/iree/vm/buffer.c index b6b2117..d433a9f 100644 --- a/iree/vm/buffer.c +++ b/iree/vm/buffer.c
@@ -95,8 +95,7 @@ // The actual buffer payload is prefixed with the buffer type so we need only // a single allocation. - iree_host_size_t prefix_size = - iree_host_align(sizeof(iree_vm_buffer_t), iree_max_align_t); + iree_host_size_t prefix_size = iree_sizeof_struct(**out_buffer); iree_host_size_t total_size = prefix_size + length; // Allocate combined [prefix | buffer] memory.
diff --git a/iree/vm/bytecode_dispatch_util.h b/iree/vm/bytecode_dispatch_util.h index 2e692be..596968e 100644 --- a/iree/vm/bytecode_dispatch_util.h +++ b/iree/vm/bytecode_dispatch_util.h
@@ -153,34 +153,12 @@ // Bytecode data access macros for reading values of a given type from a byte // offset within the current function. -#if defined(IREE_ENDIANNESS_LITTLE) -#define OP_I8(i) bytecode_data[pc + (i)] -#define OP_I16(i) *((uint16_t*)&bytecode_data[pc + (i)]) -#define OP_I32(i) *((uint32_t*)&bytecode_data[pc + (i)]) -#define OP_I64(i) *((uint64_t*)&bytecode_data[pc + (i)]) -#define OP_F32(i) *((float*)&bytecode_data[pc + (i)]) -#define OP_F64(i) *((double*)&bytecode_data[pc + (i)]) -#else -#define OP_I8(i) bytecode_data[pc + (i)] -#define OP_I16(i) \ - ((uint16_t)bytecode_data[pc + 0 + (i)]) | \ - ((uint16_t)bytecode_data[pc + 1 + (i)] << 8) -#define OP_I32(i) \ - ((uint32_t)bytecode_data[pc + 0 + (i)]) | \ - ((uint32_t)bytecode_data[pc + 1 + (i)] << 8) | \ - ((uint32_t)bytecode_data[pc + 2 + (i)] << 16) | \ - ((uint32_t)bytecode_data[pc + 3 + (i)] << 24) -#define OP_I64(i) \ - ((uint64_t)bytecode_data[pc + 0 + (i)]) | \ - ((uint64_t)bytecode_data[pc + 1 + (i)] << 8) | \ - ((uint64_t)bytecode_data[pc + 2 + (i)] << 16) | \ - ((uint64_t)bytecode_data[pc + 3 + (i)] << 24) | \ - ((uint64_t)bytecode_data[pc + 4 + (i)] << 32) | \ - ((uint64_t)bytecode_data[pc + 5 + (i)] << 40) | \ - ((uint64_t)bytecode_data[pc + 6 + (i)] << 48) | \ - ((uint64_t)bytecode_data[pc + 7 + (i)] << 56) -#error "TODO: OP_F32 and OP_F64 for big endian systems" -#endif // IREE_ENDIANNESS_LITTLE +#define OP_I8(i) iree_unaligned_load_le((uint8_t*)&bytecode_data[pc + (i)]) +#define OP_I16(i) iree_unaligned_load_le((uint16_t*)&bytecode_data[pc + (i)]) +#define OP_I32(i) iree_unaligned_load_le((uint32_t*)&bytecode_data[pc + (i)]) +#define OP_I64(i) iree_unaligned_load_le((uint64_t*)&bytecode_data[pc + (i)]) +#define OP_F32(i) iree_unaligned_load_le((float*)&bytecode_data[pc + (i)]) +#define OP_F64(i) iree_unaligned_load_le((double*)&bytecode_data[pc + (i)]) //===----------------------------------------------------------------------===// // Utilities matching the tablegen op encoding scheme
diff --git a/iree/vm/bytecode_module.c b/iree/vm/bytecode_module.c index 0feedd8..e88972d 100644 --- a/iree/vm/bytecode_module.c +++ b/iree/vm/bytecode_module.c
@@ -761,9 +761,8 @@ iree_vm_bytecode_module_t* module = NULL; IREE_RETURN_AND_END_ZONE_IF_ERROR( - z0, iree_allocator_malloc( - allocator, sizeof(iree_vm_bytecode_module_t) + type_table_size, - (void**)&module)); + z0, iree_allocator_malloc(allocator, sizeof(*module) + type_table_size, + (void**)&module)); module->allocator = allocator; iree_vm_FunctionDescriptor_vec_t function_descriptors = @@ -782,8 +781,6 @@ module->def = module_def; module->type_count = iree_vm_TypeDef_vec_len(type_defs); - module->type_table = (iree_vm_type_def_t*)((uint8_t*)module + - sizeof(iree_vm_bytecode_module_t)); iree_status_t resolve_status = iree_vm_bytecode_module_resolve_types(type_defs, module->type_table); if (!iree_status_is_ok(resolve_status)) {
diff --git a/iree/vm/bytecode_module_impl.h b/iree/vm/bytecode_module_impl.h index d602194..a63cbf8 100644 --- a/iree/vm/bytecode_module_impl.h +++ b/iree/vm/bytecode_module_impl.h
@@ -67,7 +67,7 @@ // Type table mapping module type IDs to registered VM types. iree_host_size_t type_count; - iree_vm_type_def_t* type_table; + iree_vm_type_def_t type_table[]; } iree_vm_bytecode_module_t; // A resolved and split import in the module state table.
diff --git a/iree/vm/context.c b/iree/vm/context.c index fa56039..2652fbf 100644 --- a/iree/vm/context.c +++ b/iree/vm/context.c
@@ -321,12 +321,12 @@ // TODO(benvanik): tune list growth for module count >> 4. new_capacity = context->list.capacity * 2; } - iree_vm_module_t** new_module_list; + iree_vm_module_t** new_module_list = NULL; IREE_RETURN_AND_END_ZONE_IF_ERROR( z0, iree_allocator_malloc(context->allocator, sizeof(iree_vm_module_t*) * new_capacity, (void**)&new_module_list)); - iree_vm_module_state_t** new_module_state_list; + iree_vm_module_state_t** new_module_state_list = NULL; IREE_RETURN_AND_END_ZONE_IF_ERROR( z0, iree_allocator_malloc(context->allocator,
diff --git a/iree/vm/instance.c b/iree/vm/instance.c index 34fee0c..9d1f4c6 100644 --- a/iree/vm/instance.c +++ b/iree/vm/instance.c
@@ -27,8 +27,8 @@ iree_vm_instance_t* instance = NULL; IREE_RETURN_AND_END_ZONE_IF_ERROR( - z0, iree_allocator_malloc(allocator, sizeof(iree_vm_instance_t), - (void**)&instance)); + z0, + iree_allocator_malloc(allocator, sizeof(*instance), (void**)&instance)); instance->allocator = allocator; iree_atomic_ref_count_init(&instance->ref_count);
diff --git a/iree/vm/list.c b/iree/vm/list.c index e6e784d..d668969 100644 --- a/iree/vm/list.c +++ b/iree/vm/list.c
@@ -166,7 +166,7 @@ iree_allocator_t allocator, iree_vm_list_t** out_list) { iree_vm_list_t* list = NULL; IREE_RETURN_IF_ERROR( - iree_allocator_malloc(allocator, sizeof(iree_vm_list_t), (void**)&list)); + iree_allocator_malloc(allocator, sizeof(*list), (void**)&list)); memset(list, 0, sizeof(*list)); iree_atomic_ref_count_init(&list->ref_object.counter); list->allocator = allocator;
diff --git a/iree/vm/native_module.c b/iree/vm/native_module.c index 1926ab2..7e90813 100644 --- a/iree/vm/native_module.c +++ b/iree/vm/native_module.c
@@ -360,8 +360,8 @@ // to expose this via a query_size function so that we could adjust the size // of our storage independent of the definition of the user module. iree_vm_native_module_t* module = NULL; - IREE_RETURN_IF_ERROR(iree_allocator_malloc( - allocator, sizeof(iree_vm_native_module_t), (void**)&module)); + IREE_RETURN_IF_ERROR( + iree_allocator_malloc(allocator, sizeof(*module), (void**)&module)); iree_status_t status = iree_vm_native_module_initialize( interface, module_descriptor, allocator, (iree_vm_module_t*)module);
diff --git a/scripts/git/check_submodule_init.py b/scripts/git/check_submodule_init.py new file mode 100644 index 0000000..0523c9d --- /dev/null +++ b/scripts/git/check_submodule_init.py
@@ -0,0 +1,33 @@ +# Copyright 2021 The IREE Authors +# +# 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 + +import os +import subprocess +import sys + + +def run(): + # No-op if we're not in a git repository. + try: + subprocess.check_call(['git', 'rev-parse', '--is-inside-work-tree'], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL) + except: + return + + output = os.popen("git submodule status") + submodules = output.readlines() + + for submodule in submodules: + if (submodule.strip()[0] == "-"): + print( + "The git submodule '%s' is not initialized. Please run `git submodule update --init`" + % (submodule.split()[1])) + sys.exit(1) + + +if __name__ == "__main__": + run()
