Adding `--iree-hal-substitute-executable=` flag. (#12240)
This allows for specifying one or more `executable_name=file.xxx` pairs
that each replace a `hal.executable` op with the `executable_name` with
the contents of `file.xxx`. .mlir/.mlirbc files are loaded and a
`hal.executable` with the matching name is used as a replacement while
any other file type will cause the original executable to be
externalized and linked with the specified file (.ptx/.spv/etc).
The additional `--iree-hal-substitute-executable-*-from=` flags allow
for scanning a directory for executables by name to build the
substitution mapping. Only files in the executable_name or
module_executable_name form will be substituted but we can extend this
in the future to support variant naming.
Because of phase ordering constraints around where codegen is able to
mutate host-related code such as workgroup count calculations there are
two flag sets:
`--iree-hal-substitute-executable-source=name=file.xxx`
`--iree-hal-substitute-executable-sources-from=path/`
and
`--iree-hal-substitute-executable-object=name=file.xxx`
`--iree-hal-substitute-executable-objects-from=path/`
Sources are substituted immediately prior to benchmark generation and
when substituting target objects (.ptx, .spv, etc) require that it's ok
to skip codegen (workgroup count calculation is not dependent on root op
detection, etc). Objects are substituted immediately after codegen for
use in cases where codegen is to generate the host code. There are uses
for both depending on what the input IR is and what the developer wants
to modify (host code, device code, or both).
The primary developer workflows this covers are:
1. dump executable sources via `--iree-hal-dump-executable-sources-to=`
and modify them, potentially running any number of iree-opt passes,
before linking them back in to the original program they came from
2. author custom implementations ala the custom_dispatch sample in
target toolchains (.cu -> .ptx, .glsl -> .spv, etc) and use those in
full programs without needing to modify the compiler
3. do either of the above and use the substituted executable for
microbenchmarking via `--iree-hal-dump-executable-benchmarks-to=` (so
one can easily microbenchmark handwritten kernels)
Example usage:
```sh
# dump sources for a program
iree-compile ... \
--iree-hal-dump-executable-sources-to=~/sources/
# <modify some of the sources>
# recompile with the new changes and substitute 2 executables
iree-compile ... \
--iree-hal-substitute-executable-source=_main_dispatch_0=~/sources/modified_dispatch_0.mlir \
--iree-hal-substitute-executable-source=_main_dispatch_1=~/sources/modified_dispatch_1.mlir
# same thing with search paths
iree-compile ... \
--iree-hal-executable-object-search-path=~/sources/ \
--iree-hal-substitute-executable-source=_main_dispatch_0=modified_dispatch_0.mlir
# same thing but matching all files by name as from dump-sources-to:
iree-compile ... \
--iree-hal-substitute-executable-sources-from=~/sources/
```
This works with ptx/spv as well:
```sh
# dump ptx binaries
iree-compile ... \
--iree-hal-dump-executable-binaries-to=~/binaries/
# <modify dispatch ptx>
# replace @_main_dispatch_0 with the external ptx file
iree-compile ... \
--iree-hal-substitute-executable-object=_main_dispatch_0=~/binaries/modified_dispatch_0.ptx
```
It's also possible to iterate on microbenchmarks using the custom
sources/objects:
```sh
# dump all benchmarks with the substitution active
iree-compile ... \
--iree-hal-substitute-executable-source=_main_dispatch_0=~/sources/modified_dispatch_0.mlir \
--iree-hal-dump-executable-benchmarks-to=~/benchmarks/
# inspect benchmark for the dispatch and see the hello.world attr
# can use iree-compile to build the benchmark and then iree-benchmark-module
```
Progress on #12222 (the rest for linking alternative formats is
orthogonal).diff --git a/.github/workflows/benchmark_execution.yml b/.github/workflows/benchmark_execution.yml
index eff04ac..de0b04d 100644
--- a/.github/workflows/benchmark_execution.yml
+++ b/.github/workflows/benchmark_execution.yml
@@ -33,10 +33,14 @@
type: string
outputs:
benchmark-results-dir:
- description: "Local path that stores all benchmark results."
+ description: |
+ Local path that stores all benchmark results.
+ Empty if no benchmark runs.
value: ${{ jobs.run_benchmarks.outputs.benchmark-results-dir }}
benchmark-results-gcs-artifact-dir:
- description: "GCS path that stores all benchmark results."
+ description: |
+ GCS path that stores all benchmark results.
+ Empty if no benchmark runs.
value: ${{ jobs.run_benchmarks.outputs.benchmark-results-gcs-artifact-dir }}
env:
@@ -81,6 +85,7 @@
>> "${GITHUB_OUTPUT}"
- name: "Uploading benchmark config"
id: upload
+ if: steps.export.outputs.benchmark-matrix != '[]'
env:
BENCHMARK_CONFIG: ${{ steps.export.outputs.benchmark-config }}
BENCHMARK_CONFIG_GCS_ARTIFACT: ${{ env.GCS_DIR }}/${{ steps.export.outputs.benchmark-config }}
@@ -92,6 +97,7 @@
run_benchmarks:
needs: [export_benchmark_config]
+ if: needs.export_benchmark_config.outputs.benchmark-matrix != '[]'
strategy:
# Matrix is dynamically generated by the job export_benchmark_config. So
# we only runs the benchmarks specified in inputs.benchmark-presets.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 47420a5..b59ef9b 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -889,7 +889,7 @@
compilation_benchmarks:
needs: [setup, build_e2e_test_artifacts]
- if: needs.setup.outputs.should-run == 'true'
+ if: needs.setup.outputs.should-run == 'true' && needs.setup.outputs.benchmark-presets != ''
runs-on:
- self-hosted # must come first
- runner-group=${{ needs.setup.outputs.runner-group }}
@@ -977,13 +977,7 @@
process_benchmark_results:
needs: [setup, compilation_benchmarks, execution_benchmarks]
- # execution_benchmarks is the optional dependency and skipped in presubmit
- # if no benchmark is specified to run.
- if: |
- always() &&
- needs.setup.outputs.should-run == 'true' &&
- needs.compilation_benchmarks.result == 'success' &&
- contains(fromJSON('["success", "skipped"]'), needs.execution_benchmarks.result)
+ if: needs.setup.outputs.should-run == 'true' && needs.setup.outputs.benchmark-presets != ''
runs-on:
- self-hosted # must come first
- runner-group=${{ needs.setup.outputs.runner-group }}
@@ -993,7 +987,9 @@
env:
COMPILE_STATS_RESULTS: ${{ needs.compilation_benchmarks.outputs.compile-stats-results }}
COMPILE_STATS_RESULTS_GCS_ARTIFACT: ${{ needs.compilation_benchmarks.outputs.compile-stats-results-gcs-artifact }}
+ # Empty if no execution benchmark runs.
EXECUTION_BENCHMARK_RESULTS_DIR: ${{ needs.execution_benchmarks.outputs.benchmark-results-dir }}
+ # Empty if no execution benchmark runs.
EXECUTION_BENCHMARK_RESULTS_GCS_ARTIFACT_DIR: ${{ needs.execution_benchmarks.outputs.benchmark-results-gcs-artifact-dir }}
steps:
- name: "Checking out repository"
@@ -1008,7 +1004,10 @@
"${COMPILE_STATS_RESULTS}"
- name: Downloading execution benchmark results
id: download-execution-results
- if: needs.execution_benchmarks.result == 'success'
+ # Skip the download if there is no execution benchmark results (e.g. no
+ # benchmark matches the preset/filter). In such case, no benchmark job
+ # is run in benchmark_execution.yml and the output variables are empty.
+ if: env.EXECUTION_BENCHMARK_RESULTS_GCS_ARTIFACT_DIR != ''
run: |
gcloud storage cp -r \
"${EXECUTION_BENCHMARK_RESULTS_GCS_ARTIFACT_DIR}/benchmark-results-*.json" \
@@ -1057,6 +1056,7 @@
IREE_DASHBOARD_API_TOKEN: ${{ secrets.IREE_DASHBOARD_API_TOKEN }}
run: |
build_tools/github_actions/docker_run.sh \
+ --env "IREE_DASHBOARD_API_TOKEN=${IREE_DASHBOARD_API_TOKEN}" \
gcr.io/iree-oss/benchmark-report@sha256:7498c6f32f63f13faf085463cc38656d4297519c824e63e1c99c8c258147f6ff \
./build_tools/benchmarks/upload_benchmarks_to_dashboard.py \
--verbose \
diff --git a/.github/workflows/post_benchmark_comment.yaml b/.github/workflows/post_benchmark_comment.yaml
index 9a0e61d..c3389d0 100644
--- a/.github/workflows/post_benchmark_comment.yaml
+++ b/.github/workflows/post_benchmark_comment.yaml
@@ -80,7 +80,7 @@
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GIST_BOT_TOKEN: ${{ secrets.GIST_BOT_TOKEN }}
- COMMENT_BOT_USER: iree-github-actions-bot
+ COMMENT_BOT_USER: github-actions[bot]
# Get PR number from the event instead of the untrusted comment
# artifact, to make sure we won't incorrectly update other PRs.
BENCHMARK_COMMENT_ARTIFACT: ${{ steps.download.outputs.benchmark-comment-artifact }}
diff --git a/build_tools/benchmarks/common/benchmark_definition.py b/build_tools/benchmarks/common/benchmark_definition.py
index f38ad5b..b3c91f0 100644
--- a/build_tools/benchmarks/common/benchmark_definition.py
+++ b/build_tools/benchmarks/common/benchmark_definition.py
@@ -138,7 +138,6 @@
repetitions = 10
cmd = [
- "--device_allocator=caching",
"--time_unit=ns",
"--benchmark_format=json",
"--benchmark_out_format=json",
diff --git a/build_tools/benchmarks/export_benchmark_config.py b/build_tools/benchmarks/export_benchmark_config.py
index 0a153a0..315dbdf 100755
--- a/build_tools/benchmarks/export_benchmark_config.py
+++ b/build_tools/benchmarks/export_benchmark_config.py
@@ -63,6 +63,9 @@
(config.target_device_spec.architecture.type == common_definitions.
ArchitectureType.GPU and config.target_device_spec.host_environment.
platform == "android"),
+ # Not a preset for execution benchmarks.
+ "comp-stats":
+ lambda _config: False,
}
diff --git a/build_tools/cmake/build_riscv.sh b/build_tools/cmake/build_riscv.sh
index 042932f..d3f4978 100755
--- a/build_tools/cmake/build_riscv.sh
+++ b/build_tools/cmake/build_riscv.sh
@@ -22,6 +22,7 @@
set -xeuo pipefail
BUILD_DIR="${1:-${IREE_TARGET_BUILD_DIR:-build-riscv}}"
+BUILD_TYPE="${IREE_BUILD_TYPE:-RelWithDebInfo}"
RISCV_PLATFORM="${IREE_TARGET_PLATFORM:-linux}"
RISCV_ARCH="${IREE_TARGET_ARCH:-riscv_64}"
RISCV_COMPILER_FLAGS="${RISCV_COMPILER_FLAGS:--O3}"
@@ -38,6 +39,7 @@
args=(
"-G" "Ninja"
"-B" "${BUILD_DIR}"
+ "-DCMAKE_BUILD_TYPE=${BUILD_TYPE}"
"-DPython3_EXECUTABLE=${IREE_PYTHON3_EXECUTABLE}"
"-DPYTHON_EXECUTABLE=${IREE_PYTHON3_EXECUTABLE}"
"-DCMAKE_TOOLCHAIN_FILE=${TOOLCHAIN_FILE}"
diff --git a/build_tools/github_actions/configure_ci.py b/build_tools/github_actions/configure_ci.py
index f125401..762ea6e 100755
--- a/build_tools/github_actions/configure_ci.py
+++ b/build_tools/github_actions/configure_ci.py
@@ -61,7 +61,7 @@
RUNNER_ENV_DEFAULT = "prod"
RUNNER_ENV_OPTIONS = [RUNNER_ENV_DEFAULT, "testing"]
-BENCHMARK_PRESET_OPTIONS = ["all", "cuda", "x86_64"]
+BENCHMARK_PRESET_OPTIONS = ["all", "cuda", "x86_64", "comp-stats"]
def skip_path(path: str) -> bool:
diff --git a/build_tools/python/benchmark_suites/iree/module_execution_configs.py b/build_tools/python/benchmark_suites/iree/module_execution_configs.py
index d1b5dbb..fe81e97 100644
--- a/build_tools/python/benchmark_suites/iree/module_execution_configs.py
+++ b/build_tools/python/benchmark_suites/iree/module_execution_configs.py
@@ -5,35 +5,54 @@
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Defines ModuleExecutionConfig for benchmarks."""
+from typing import List, Optional, Sequence
+
from e2e_test_framework.definitions import iree_definitions
from e2e_test_framework import unique_ids
-ELF_LOCAL_SYNC_CONFIG = iree_definitions.ModuleExecutionConfig(
+
+def _with_caching_allocator(
+ id: str,
+ tags: List[str],
+ loader: iree_definitions.RuntimeLoader,
+ driver: iree_definitions.RuntimeDriver,
+ extra_flags: Optional[Sequence[str]] = None
+) -> iree_definitions.ModuleExecutionConfig:
+ extra_flags = [] if extra_flags is None else list(extra_flags)
+ return iree_definitions.ModuleExecutionConfig(
+ id=id,
+ tags=tags,
+ loader=loader,
+ driver=driver,
+ extra_flags=["--device_allocator=caching"] + extra_flags)
+
+
+ELF_LOCAL_SYNC_CONFIG = _with_caching_allocator(
id=unique_ids.IREE_MODULE_EXECUTION_CONFIG_LOCAL_SYNC,
tags=["full-inference", "default-flags"],
loader=iree_definitions.RuntimeLoader.EMBEDDED_ELF,
driver=iree_definitions.RuntimeDriver.LOCAL_SYNC)
-CUDA_CONFIG = iree_definitions.ModuleExecutionConfig(
+CUDA_CONFIG = _with_caching_allocator(
id=unique_ids.IREE_MODULE_EXECUTION_CONFIG_CUDA,
tags=["full-inference", "default-flags"],
loader=iree_definitions.RuntimeLoader.NONE,
driver=iree_definitions.RuntimeDriver.CUDA)
-VULKAN_CONFIG = iree_definitions.ModuleExecutionConfig(
+VULKAN_CONFIG = _with_caching_allocator(
id=unique_ids.IREE_MODULE_EXECUTION_CONFIG_VULKAN,
tags=["full-inference", "default-flags"],
loader=iree_definitions.RuntimeLoader.NONE,
driver=iree_definitions.RuntimeDriver.VULKAN)
-VULKAN_BATCH_SIZE_16_CONFIG = iree_definitions.ModuleExecutionConfig(
+VULKAN_BATCH_SIZE_16_CONFIG = _with_caching_allocator(
id=unique_ids.IREE_MODULE_EXECUTION_CONFIG_VULKAN_BATCH_SIZE_16,
tags=["full-inference", "experimental-flags"],
loader=iree_definitions.RuntimeLoader.NONE,
driver=iree_definitions.RuntimeDriver.VULKAN,
extra_flags=["--batch_size=16"])
-VULKAN_BATCH_SIZE_32_CONFIG = iree_definitions.ModuleExecutionConfig(
+VULKAN_BATCH_SIZE_32_CONFIG = _with_caching_allocator(
id=unique_ids.IREE_MODULE_EXECUTION_CONFIG_VULKAN_BATCH_SIZE_32,
tags=["full-inference", "experimental-flags"],
loader=iree_definitions.RuntimeLoader.NONE,
@@ -43,7 +62,7 @@
def get_elf_local_task_config(thread_num: int):
config_id = f"{unique_ids.IREE_MODULE_EXECUTION_CONFIG_LOCAL_TASK_BASE}-{thread_num}"
- return iree_definitions.ModuleExecutionConfig(
+ return _with_caching_allocator(
id=config_id,
tags=[f"{thread_num}-thread", "full-inference", "default-flags"],
loader=iree_definitions.RuntimeLoader.EMBEDDED_ELF,
@@ -53,7 +72,7 @@
def get_vmvx_local_task_config(thread_num: int):
config_id = f"{unique_ids.IREE_MODULE_EXECUTION_CONFIG_VMVX_LOCAL_TASK_BASE}-{thread_num}"
- return iree_definitions.ModuleExecutionConfig(
+ return _with_caching_allocator(
id=config_id,
tags=[f"{thread_num}-thread", "full-inference", "default-flags"],
loader=iree_definitions.RuntimeLoader.VMVX_MODULE,
diff --git a/build_tools/python/e2e_model_tests/cmake_generator.py b/build_tools/python/e2e_model_tests/cmake_generator.py
index baf5e6c..1345c92 100644
--- a/build_tools/python/e2e_model_tests/cmake_generator.py
+++ b/build_tools/python/e2e_model_tests/cmake_generator.py
@@ -26,6 +26,11 @@
runner_args = run_module_utils.build_run_flags_for_model(
model=model,
model_input_data=test_config.input_data) + test_config.extra_test_flags
+ # TODO(#11136): Currently the DRIVER is a separate field in the CMake rule (
+ # and has effect on test labels). Rules should be generated in another way
+ # to avoid that. Generates the flags without the driver for now.
+ runner_args += run_module_utils.build_run_flags_for_execution_config(
+ test_config.execution_config, with_driver=False)
cmake_rule = cmake_builder.rules.build_iree_benchmark_suite_module_test(
target_name=test_config.name,
model=f"{model.id}_{model.name}",
diff --git a/build_tools/python/e2e_model_tests/run_module_utils.py b/build_tools/python/e2e_model_tests/run_module_utils.py
index 7e3827e..bd88c08 100644
--- a/build_tools/python/e2e_model_tests/run_module_utils.py
+++ b/build_tools/python/e2e_model_tests/run_module_utils.py
@@ -26,15 +26,26 @@
def build_run_flags_for_execution_config(
module_execution_config: ModuleExecutionConfig,
- gpu_id: str = "0") -> List[str]:
- """Returns the IREE run module flags of the execution config."""
+ gpu_id: str = "0",
+ with_driver: bool = True) -> List[str]:
+ """Returns the IREE run module flags of the execution config.
- run_flags = list(module_execution_config.extra_flags)
- driver = module_execution_config.driver
- if driver == RuntimeDriver.CUDA:
- run_flags.append(f"--device=cuda://{gpu_id}")
- else:
- run_flags.append(f"--device={driver.value}")
+ Args:
+ module_execution_config: execution config.
+ gpu_id: target gpu id, if runs on GPUs.
+ with_driver: populate the driver flags if true. False can be used for
+ generating flags for some CMake rules with a separate DRIVER arg.
+ Returns:
+ List of flags.
+ """
+
+ run_flags = module_execution_config.extra_flags.copy()
+ if with_driver:
+ driver = module_execution_config.driver
+ if driver == RuntimeDriver.CUDA:
+ run_flags.append(f"--device=cuda://{gpu_id}")
+ else:
+ run_flags.append(f"--device={driver.value}")
return run_flags
diff --git a/build_tools/python/e2e_model_tests/run_module_utils_test.py b/build_tools/python/e2e_model_tests/run_module_utils_test.py
index dccb98f..5a255da 100644
--- a/build_tools/python/e2e_model_tests/run_module_utils_test.py
+++ b/build_tools/python/e2e_model_tests/run_module_utils_test.py
@@ -55,6 +55,19 @@
self.assertEqual(flags, ["--device=cuda://3"])
+ def test_build_run_flags_for_execution_config_without_driver(self):
+ execution_config = iree_definitions.ModuleExecutionConfig(
+ id="123",
+ tags=["test"],
+ loader=iree_definitions.RuntimeLoader.EMBEDDED_ELF,
+ driver=iree_definitions.RuntimeDriver.LOCAL_TASK,
+ extra_flags=["--task=10"])
+
+ flags = run_module_utils.build_run_flags_for_execution_config(
+ execution_config, with_driver=False)
+
+ self.assertEqual(flags, ["--task=10"])
+
def test_build_linux_wrapper_cmds_for_device_spec(self):
device_spec = common_definitions.DeviceSpec(
id="abc",
diff --git a/compiler/src/iree/compiler/API2/Internal/Embed.cpp b/compiler/src/iree/compiler/API2/Internal/Embed.cpp
index 1184e88..dea1149 100644
--- a/compiler/src/iree/compiler/API2/Internal/Embed.cpp
+++ b/compiler/src/iree/compiler/API2/Internal/Embed.cpp
@@ -692,7 +692,7 @@
};
unwrap(inv)->passManager.enableCrashReproducerGeneration(
- [&](std::string &errorMessage)
+ [=](std::string &errorMessage)
-> std::unique_ptr<mlir::PassManager::ReproducerStream> {
iree_compiler_output_t *output = nullptr;
auto error = onCrashCallback(&output, userData);
diff --git a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp
index ca7a346..46a6474 100644
--- a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp
@@ -87,7 +87,8 @@
ADD_PATTERN(swapPaddingElideConditional,
getSwapPaddingElideConditionalAttrName)
ADD_PATTERN(swappingPatterns, getSwappingPatternsAttrName)
- ADD_PATTERN(unrollVectorsGpuMma, getUnrollVectorsGpuMmaAttrName)
+ ADD_PATTERN(unrollVectorsGpuMmaSync, getUnrollVectorsGpuMmaSyncAttrName)
+ ADD_PATTERN(unrollVectorsGpuWmma, getUnrollVectorsGpuWmmaAttrName)
#undef ADD_PATTERN
result.addTypes({pdl::OperationType::get(ctx)});
}
@@ -204,12 +205,12 @@
});
}
-static Optional<SmallVector<int64_t>> getGPUTensorCoreNativeVectorSize(
+static Optional<SmallVector<int64_t>> getGPUTensorCoreNativeMmaSyncVectorSize(
Operation *op) {
- return getWmmaNativeVectorSize(op);
+ return getMmaNativeVectorSize(op);
}
-static void addUnrollVectorsGpuMmaPatterns(RewritePatternSet &patterns) {
+static void addUnrollVectorsGpuMmaSyncPatterns(RewritePatternSet &patterns) {
auto unrollOrder = [](Operation *op) -> Optional<SmallVector<int64_t>> {
auto contract = dyn_cast<vector::ContractionOp>(op);
if (!contract) return std::nullopt;
@@ -217,7 +218,24 @@
};
vector::populateVectorUnrollPatterns(
patterns, vector::UnrollVectorOptions()
- .setNativeShapeFn(getGPUTensorCoreNativeVectorSize)
+ .setNativeShapeFn(getGPUTensorCoreNativeMmaSyncVectorSize)
+ .setUnrollTraversalOrderFn(unrollOrder));
+}
+
+static Optional<SmallVector<int64_t>> getGPUTensorCoreNativeWmmaVectorSize(
+ Operation *op) {
+ return getWmmaNativeVectorSize(op);
+}
+
+static void addUnrollVectorsGpuWmmaPatterns(RewritePatternSet &patterns) {
+ auto unrollOrder = [](Operation *op) -> Optional<SmallVector<int64_t>> {
+ auto contract = dyn_cast<vector::ContractionOp>(op);
+ if (!contract) return std::nullopt;
+ return mlir::iree_compiler::gpuMmaUnrollOrder(contract);
+ };
+ vector::populateVectorUnrollPatterns(
+ patterns, vector::UnrollVectorOptions()
+ .setNativeShapeFn(getGPUTensorCoreNativeWmmaVectorSize)
.setUnrollTraversalOrderFn(unrollOrder));
}
@@ -265,7 +283,9 @@
linalg::populateFoldReshapeOpsByExpansionPatterns(
patterns, [](OpOperand *) { return true; });
}
- if (getUnrollVectorsGpuMma()) addUnrollVectorsGpuMmaPatterns(patterns);
+ if (getUnrollVectorsGpuMmaSync())
+ addUnrollVectorsGpuMmaSyncPatterns(patterns);
+ if (getUnrollVectorsGpuWmma()) addUnrollVectorsGpuWmmaPatterns(patterns);
TrackingListener listener(state);
GreedyRewriteConfig config;
diff --git a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h
index e268c4e..d9241fc 100644
--- a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h
+++ b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h
@@ -50,7 +50,8 @@
bool rewritePackOps = false;
bool swapPaddingElideConditional = false;
bool swappingPatterns = false;
- bool unrollVectorsGpuMma = false;
+ bool unrollVectorsGpuMmaSync = false;
+ bool unrollVectorsGpuWmma = false;
};
} // namespace transform_dialect
} // namespace IREE
diff --git a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensionsOps.td b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensionsOps.td
index 52f535d..cbd948b 100644
--- a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensionsOps.td
+++ b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensionsOps.td
@@ -68,9 +68,12 @@
tensor.extract_slice swapping pattern. This injects static information
that guarantees padding is smaller than the window size which guarantees
we never see a tile comprised of padding-only.
- - unroll_vectors_gpu_mma: adds patterns that unroll vectors to a native tile
+ - unroll_vectors_gpu_mma_sync: adds patterns that unroll vectors to a native tile
size for GPUs with mma operations. The size is currently hardcoded but
should be refactored upstream and made pluggable.
+ - unroll_vectors_gpu_wmma: adds patterns that unroll vectors to a native tile
+ size for GPUs with wmma operations. The size is currently hardcoded but
+ should be refactored upstream and made pluggable.
#### Return modes:
@@ -101,7 +104,8 @@
UnitAttr:$rewrite_pack_ops,
UnitAttr:$swap_padding_elide_conditional,
UnitAttr:$swapping_patterns,
- UnitAttr:$unroll_vectors_gpu_mma);
+ UnitAttr:$unroll_vectors_gpu_mma_sync,
+ UnitAttr:$unroll_vectors_gpu_wmma);
let results = (outs PDL_Operation:$result);
let assemblyFormat = "$target attr-dict";
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUTensorCoreVectorization.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUTensorCoreVectorization.cpp
index 73b167e..3b9ad7e 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUTensorCoreVectorization.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUTensorCoreVectorization.cpp
@@ -46,154 +46,9 @@
vector::populateVectorReductionToContractPatterns(patterns);
}
-/// Returns vector::ContractionOp operand's index where the result is used.
-static Optional<int> getVectorContractOpOperandId(
- vector::ContractionOp contractOp, OpResult result) {
- if (contractOp.getLhs() == result) return 0;
- if (contractOp.getRhs() == result) return 1;
- if (contractOp.getAcc() == result) return 2;
- return std::nullopt;
-}
-
-/// Returns vector::ContractionOp operand's index where the
-/// vector::TransferReadOp is consumed either consumed directly or via
-/// vector::ExtractStridedSliceOp.
-static Optional<int> getVectorContractOpOperandIdForVectorReadOp(
- Operation *op) {
- vector::ContractionOp contractOp;
-
- Operation *firstLevelUser = *((op->getUsers()).begin());
- if (auto contractOp = dyn_cast<vector::ContractionOp>(firstLevelUser))
- return getVectorContractOpOperandId(contractOp, op->getResult(0));
- Operation *secondLevelUser = *((firstLevelUser->getUsers()).begin());
- if (auto contractOp = dyn_cast<vector::ContractionOp>(secondLevelUser))
- return getVectorContractOpOperandId(contractOp,
- firstLevelUser->getResult(0));
- return std::nullopt;
-}
-
-/// Helper function to return native size for MMA.SYNC-based operations.
-static Optional<SmallVector<int64_t>> getMmaNativeVectorSize(Operation *op) {
- // Shape of native Tensor Core GPU mma.sync operations.
- int64_t mmaShapeM = 16;
- int64_t mmaShapeN = 8;
- int64_t mmaShapeK;
-
- // Shape the mma.sync warp-level operation.
- if (auto contract = dyn_cast<vector::ContractionOp>(op)) {
- Type sourceType = contract.getLhsType().getElementType();
-
- // Set mmaShapeK based on sourceType.
- if (sourceType.isInteger(4))
- mmaShapeK = 64;
- else if (sourceType.isInteger(8))
- mmaShapeK = 32;
- else if (sourceType.isF16() || sourceType.isBF16())
- mmaShapeK = 16;
- else if (sourceType.isF32())
- mmaShapeK = 8;
- else
- return std::nullopt;
-
- // Initialize/set the starting dims of the ranked shape, such as batch,
- // to 1.
- SmallVector<int64_t> mmaShape(contract.getIteratorTypes().size() - 3, 1);
- mmaShape.append({mmaShapeM, mmaShapeN, mmaShapeK});
- return mmaShape;
- }
-
- // Shape of warp-level vector write operation.
- if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op)) {
- SmallVector<int64_t> outputShape(writeOp.getVectorType().getRank() - 2, 1);
- outputShape.append({mmaShapeM, mmaShapeN});
- return outputShape;
- }
-
- // Shape of warp-level vector read (load) operation.
- if (auto readOp = dyn_cast<vector::TransferReadOp>(op)) {
- auto resultVectorType = readOp.getVector().getType().cast<VectorType>();
- Type resultElementType = resultVectorType.getElementType();
-
- Optional<int> operandId = getVectorContractOpOperandIdForVectorReadOp(op);
- if (!operandId) {
- op->emitError() << "Cannot determine operandId this "
- "vector::TransferReadOp is used as in the "
- "vector::TransferContractOp";
- return std::nullopt;
- }
-
- // Loading F16 values from Shared Memory to Registers.
- if (resultElementType.isF16() || resultElementType.isBF16()) {
- // For matrixC.
- if (*operandId == 2) {
- SmallVector<int64_t> readShape;
- readShape.append({mmaShapeM, mmaShapeN});
- return readShape;
- }
-
- // For matrixA and matrixB.
- if (*operandId == 0 || *operandId == 1) {
- // MmaSyncOp input operands: matrixA and matrixB.
- // LDSMx1, x2, x4:
- // - LDSMx1 loads a 1 tile of 8x8.
- // - LDSMx2 loads a 2 tiles of 8x8.
- // - LDSMx4 loads a 4 tiles of 8x8. (in use)
- // IREE uses the largest tiled load, i.e., LDSMx4.
-
- // MmaSyncOp source operand: matrixC.
- // matrixC is also read/written in tiled block of 16x16. In the pass
- // OptimizeVectorTransfer, matrixC reads are moved above the mainloop
- // and writes are moved below the mainloop. Thus, mma.sync read/write
- // accumulator inplace.
-
- SmallVector<int64_t> readShape;
- readShape.append({16, 16});
- return readShape;
- }
- }
-
- // Loading F32 values from Shared Memory to Registers.
- if (resultElementType.isF32()) {
- // Set mmaShapeK for F32 datatype mma.sync.f32.tf32.m16n8k8.
- mmaShapeK = 8;
-
- // For matrixC.
- if (*operandId == 2) {
- SmallVector<int64_t> readShape;
- readShape.append({mmaShapeM, mmaShapeN});
- return readShape;
- }
- // For matrixA.
- if (*operandId == 0) {
- SmallVector<int64_t> readShape;
- readShape.append({mmaShapeM, mmaShapeK});
- return readShape;
- }
- // For matrixB.
- if (*operandId == 1) {
- // Do not use ldmatrix for matrixB.
- // Transfer read ops may need different shapes based on how they are
- // being used. For simplicity just match the shape used by the extract
- // strided op.
- VectorType sliceType;
- for (Operation *users : op->getUsers()) {
- auto extract = dyn_cast<vector::ExtractStridedSliceOp>(users);
- if (!extract) return std::nullopt;
- auto vecType = extract.getResult().getType().cast<VectorType>();
- if (sliceType && sliceType != vecType) return std::nullopt;
- sliceType = vecType;
- }
- return llvm::to_vector<>(sliceType.getShape());
- }
- }
- }
- return std::nullopt;
-}
-
static Optional<SmallVector<int64_t>> getGPUTensorCoreNativeVectorSize(
Operation *op) {
if (llvmgpuUseMMASync) return getMmaNativeVectorSize(op);
-
return getWmmaNativeVectorSize(op);
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/BUILD b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/BUILD
index eea5e4e..ab3492a 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/BUILD
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/BUILD
@@ -75,6 +75,7 @@
"@llvm-project//mlir:LinalgUtils",
"@llvm-project//mlir:MemRefDialect",
"@llvm-project//mlir:NVGPUDialect",
+ "@llvm-project//mlir:NVGPUTransforms",
"@llvm-project//mlir:PDLDialect",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:SCFDialect",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/CMakeLists.txt
index 98d3539..a1be50e 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/CMakeLists.txt
@@ -43,6 +43,7 @@
MLIRLinalgUtils
MLIRMemRefDialect
MLIRNVGPUDialect
+ MLIRNVGPUTransforms
MLIRPDLDialect
MLIRPass
MLIRSCFDialect
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp
index 5aa0b5a..4883ca3 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp
@@ -19,6 +19,7 @@
#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
+#include "mlir/Dialect/NVGPU/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/IR/DeviceMappingInterface.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Transform/IR/TransformUtils.h"
@@ -635,35 +636,55 @@
"patterns greedily");
return emitDefaultDefiniteFailure(target);
}
- MLIRContext *ctx = target->getContext();
- // Step 1. Unroll vectors to native size.
- RewritePatternSet unrollPatterns(ctx);
- vector::populateVectorUnrollPatterns(
- unrollPatterns,
- vector::UnrollVectorOptions().setNativeShapeFn(getWmmaNativeVectorSize));
- if (failed(applyPatternsAndFoldGreedily(target, std::move(unrollPatterns)))) {
- target->emitOpError(
- "failed to break up vector operations into mma native size");
+ auto funcOp = dyn_cast<func::FuncOp>(target);
+ if (!funcOp) {
+ target->emitOpError("Must apply to a func op");
return emitDefaultDefiniteFailure(target);
}
- // TODO: Step 2. add pattern to propagate the extract through the scf.for ops.
+ if (!(getUseMmaSync() ^ getUseWmma())) {
+ target->emitOpError(
+ "Exactly one of use_mma_sync or use_wmma must be specified");
+ return emitDefaultDefiniteFailure(target);
+ }
- // Step 3. Convert slice of contract operations to wmma ops.
+ MLIRContext *ctx = target->getContext();
+
+ // Unrolling to native vector size must have previously occurred.
+ // TODO: Add pattern to propagate the extract through the scf.for ops.
+ // Convert slice of contract operations to mma_sync/wmma ops.
RewritePatternSet patterns(ctx);
mlir::vector::populateCastAwayVectorLeadingOneDimPatterns(patterns);
- populatePrepareVectorToMMAPatterns(patterns, /*llvmgpuUseMMASync=*/false);
+ populatePrepareVectorToMMAPatterns(patterns, getUseMmaSync());
if (failed(applyPatternsAndFoldGreedily(target, std::move(patterns)))) {
target->emitOpError("vector to mma preparation patterns failed to apply");
return emitDefaultDefiniteFailure(target);
}
+
IRRewriter rewriter(getContext());
- if (failed(convertVectorToMMAOps(rewriter, target))) {
+ if (getUseWmma()) {
+ if (failed(convertVectorToMMAOps(rewriter, target))) {
+ target->emitOpError("vector to wmma patterns failed to apply");
+ return emitDefaultDefiniteFailure(target);
+ }
+ results.push_back(target);
+ return DiagnosedSilenceableFailure::success();
+ }
+
+ if (failed(convertVectorToNVVMCompatibleMMASync(rewriter, funcOp))) {
target->emitOpError("vector to mma patterns failed to apply");
return emitDefaultDefiniteFailure(target);
}
-
+ // Using TF32 for Float.
+ RewritePatternSet f32ToTF32patterns(funcOp.getContext());
+ nvgpu::populateMmaSyncF32ToTF32Patterns(f32ToTF32patterns,
+ nvgpu::MmaSyncF32Lowering::TF32);
+ if (failed(applyPatternsAndFoldGreedily(getOperation(),
+ std::move(f32ToTF32patterns)))) {
+ target->emitOpError("vector to mma F32ToTF32 patterns failed to apply");
+ return emitDefaultDefiniteFailure(target);
+ }
results.push_back(target);
return DiagnosedSilenceableFailure::success();
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensionsOps.td b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensionsOps.td
index a13519b..656fee7 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensionsOps.td
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensionsOps.td
@@ -328,12 +328,16 @@
operations are bigger than the native mma size it will first split up those
vector operations.
+ Exactly one of use_wmma or use_mma_sync must be specified.
+
#### Return modes
This transform consumes the target handle and produces a result handle.
}];
- let arguments = (ins PDL_Operation:$target);
+ let arguments = (ins PDL_Operation:$target,
+ UnitAttr:$use_mma_sync,
+ UnitAttr:$use_wmma);
let results = (outs PDL_Operation:$result);
let assemblyFormat = "$target attr-dict";
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/transform_vector_to_mma.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/transform_vector_to_mma.mlir
index f182946..c89d080 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/transform_vector_to_mma.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/transform_vector_to_mma.mlir
@@ -50,6 +50,7 @@
transform.structured.canonicalized_sequence failures(propagate) {
^bb1(%variant_op: !pdl.operation):
%func = transform.structured.match ops{["func.func"]} in %variant_op : (!pdl.operation) -> !pdl.operation
- transform.iree.vector.vector_to_mma_conversion %func
+ %func_2 = transform.iree.apply_patterns %func { unroll_vectors_gpu_wmma }
+ transform.iree.vector.vector_to_mma_conversion %func_2 { use_wmma }
}
}
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp b/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
index 2a59748..16e75c6 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
@@ -541,8 +541,9 @@
auto lhsType = lhs->get().getType().cast<ShapedType>();
auto rhsType = rhs->get().getType().cast<ShapedType>();
- auto elementBits = lhsType.getElementType().getIntOrFloatBitWidth();
- if (elementBits != 16 && elementBits != 32) return success();
+ auto elementBits =
+ static_cast<int>(lhsType.getElementType().getIntOrFloatBitWidth());
+ if (!llvm::is_contained({8, 16, 32}, elementBits)) return success();
ArrayRef<int64_t> lhsShape = lhsType.getShape();
ArrayRef<int64_t> rhsShape = rhsType.getShape();
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/test/config_default_matmul.mlir b/compiler/src/iree/compiler/Codegen/SPIRV/test/config_default_matmul.mlir
index 8fbf54e..b71b594 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/test/config_default_matmul.mlir
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/test/config_default_matmul.mlir
@@ -55,7 +55,7 @@
// -----
-// Non-16 / non-32 bit types cannot be vectorized right now.
+// 8-bit integers can be vectorized.
#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
#hal.descriptor_set.layout<0, bindings = [
@@ -64,7 +64,7 @@
#hal.descriptor_set.binding<2, storage_buffer>
]>
]>
-hal.executable private @matmul_64x16 {
+hal.executable private @matmul_64x16xi8 {
hal.executable.variant public @vulkan_spirv_fb, target = <"vulkan", "vulkan-spirv-fb", {
spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader], []>, Unknown:IntegratedGPU, #spirv.resource_limits<
max_compute_shared_memory_size = 16384,
@@ -72,9 +72,9 @@
max_compute_workgroup_size = [128, 128, 64],
subgroup_size = 64>>
}> {
- hal.executable.export public @matmul_64x16 layout(#pipeline_layout)
+ hal.executable.export public @matmul_64x16xi8 layout(#pipeline_layout)
builtin.module {
- func.func @matmul_64x16() {
+ func.func @matmul_64x16xi8() {
%c0 = arith.constant 0 : index
%c16 = arith.constant 16 : index
%c64 = arith.constant 64 : index
@@ -98,12 +98,66 @@
}
}
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 16], [2, 8], [0, 0, 8]{{\]}}>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<SPIRVBaseVectorize>
+// CHECK: hal.executable.export public @matmul_64x16xi8
+// CHECK-SAME: translation_info = #[[TRANSLATION]]
+// CHECK-SAME: workgroup_size = [2 : index, 32 : index, 1 : index]
+// CHECK: func.func @matmul_64x16xi8()
+// CHECK: linalg.matmul
+// CHECK-SAME: lowering_config = #[[CONFIG]]
+
+// -----
+
+// Non-16 / non-32 bit types cannot be vectorized right now.
+
+#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
+ #hal.descriptor_set.layout<0, bindings = [
+ #hal.descriptor_set.binding<0, storage_buffer>,
+ #hal.descriptor_set.binding<1, storage_buffer>,
+ #hal.descriptor_set.binding<2, storage_buffer>
+ ]>
+]>
+hal.executable private @matmul_64x16xi64 {
+ hal.executable.variant public @vulkan_spirv_fb, target = <"vulkan", "vulkan-spirv-fb", {
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, Int64], []>, Unknown:IntegratedGPU, #spirv.resource_limits<
+ max_compute_shared_memory_size = 16384,
+ max_compute_workgroup_invocations = 128,
+ max_compute_workgroup_size = [128, 128, 64],
+ subgroup_size = 64>>
+ }> {
+ hal.executable.export public @matmul_64x16xi64 layout(#pipeline_layout)
+ builtin.module {
+ func.func @matmul_64x16xi64() {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c64 = arith.constant 64 : index
+ %c0_i32 = arith.constant 0 : i32
+ %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) : !flow.dispatch.tensor<readonly:tensor<64x32xi64>>
+ %1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) : !flow.dispatch.tensor<readonly:tensor<32x16xi64>>
+ %2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) : !flow.dispatch.tensor<writeonly:tensor<64x16xi64>>
+ %8 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [64, 32], strides = [1, 1]
+ : !flow.dispatch.tensor<readonly:tensor<64x32xi64>> -> tensor<64x32xi64>
+ %10 = flow.dispatch.tensor.load %1, offsets = [0, 0], sizes = [32, 16], strides = [1, 1]
+ : !flow.dispatch.tensor<readonly:tensor<32x16xi64>> -> tensor<32x16xi64>
+ %15 = tensor.empty() : tensor<64x16xi64>
+ %16 = linalg.fill ins(%c0_i32 : i32) outs(%15 : tensor<64x16xi64>) -> tensor<64x16xi64>
+ %17 = linalg.matmul {__internal_linalg_transform__ = "workgroup"}
+ ins(%8, %10 : tensor<64x32xi64>, tensor<32x16xi64>) outs(%16 : tensor<64x16xi64>) -> tensor<64x16xi64>
+ flow.dispatch.tensor.store %17, %2, offsets = [0, 0], sizes = [64, 16], strides = [1, 1]
+ : tensor<64x16xi64> -> !flow.dispatch.tensor<writeonly:tensor<64x16xi64>>
+ return
+ }
+ }
+ }
+}
+
// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[4, 16], [1, 1]{{\]}}>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<SPIRVBaseDistribute>
-// CHECK: hal.executable.export public @matmul_64x16
+// CHECK: hal.executable.export public @matmul_64x16xi64
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK-SAME: workgroup_size = [16 : index, 4 : index, 1 : index]
-// CHECK: func.func @matmul_64x16()
+// CHECK: func.func @matmul_64x16xi64()
// CHECK: linalg.matmul
// CHECK-SAME: lowering_config = #[[CONFIG]]
diff --git a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp
index eb87340..e261db4 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp
@@ -596,5 +596,152 @@
return std::nullopt;
}
+//===----------------------------------------------------------------------===//
+// getMmaNativeVectorSize
+//===----------------------------------------------------------------------===//
+/// Returns vector::ContractionOp operand's index where the result is used.
+static Optional<int> getVectorContractOpOperandId(
+ vector::ContractionOp contractOp, OpResult result) {
+ if (contractOp.getLhs() == result) return 0;
+ if (contractOp.getRhs() == result) return 1;
+ if (contractOp.getAcc() == result) return 2;
+ return std::nullopt;
+}
+
+/// Returns vector::ContractionOp operand's index where the
+/// vector::TransferReadOp is consumed either consumed directly or via
+/// vector::ExtractStridedSliceOp.
+static Optional<int> getVectorContractOpOperandIdForVectorReadOp(
+ Operation *op) {
+ vector::ContractionOp contractOp;
+
+ Operation *firstLevelUser = *((op->getUsers()).begin());
+ if (auto contractOp = dyn_cast<vector::ContractionOp>(firstLevelUser))
+ return getVectorContractOpOperandId(contractOp, op->getResult(0));
+ Operation *secondLevelUser = *((firstLevelUser->getUsers()).begin());
+ if (auto contractOp = dyn_cast<vector::ContractionOp>(secondLevelUser))
+ return getVectorContractOpOperandId(contractOp,
+ firstLevelUser->getResult(0));
+ return std::nullopt;
+}
+
+/// Helper function to return native size for MMA.SYNC-based operations.
+Optional<SmallVector<int64_t>> getMmaNativeVectorSize(Operation *op) {
+ // Shape of native Tensor Core GPU mma.sync operations.
+ int64_t mmaShapeM = 16;
+ int64_t mmaShapeN = 8;
+ int64_t mmaShapeK;
+
+ // Shape the mma.sync warp-level operation.
+ if (auto contract = dyn_cast<vector::ContractionOp>(op)) {
+ Type sourceType = contract.getLhsType().getElementType();
+
+ // Set mmaShapeK based on sourceType.
+ if (sourceType.isInteger(4))
+ mmaShapeK = 64;
+ else if (sourceType.isInteger(8))
+ mmaShapeK = 32;
+ else if (sourceType.isF16() || sourceType.isBF16())
+ mmaShapeK = 16;
+ else if (sourceType.isF32())
+ mmaShapeK = 8;
+ else
+ return std::nullopt;
+
+ // Initialize/set the starting dims of the ranked shape, such as batch,
+ // to 1.
+ SmallVector<int64_t> mmaShape(contract.getIteratorTypes().size() - 3, 1);
+ mmaShape.append({mmaShapeM, mmaShapeN, mmaShapeK});
+ return mmaShape;
+ }
+
+ // Shape of warp-level vector write operation.
+ if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op)) {
+ SmallVector<int64_t> outputShape(writeOp.getVectorType().getRank() - 2, 1);
+ outputShape.append({mmaShapeM, mmaShapeN});
+ return outputShape;
+ }
+
+ // Shape of warp-level vector read (load) operation.
+ if (auto readOp = dyn_cast<vector::TransferReadOp>(op)) {
+ auto resultVectorType = readOp.getVector().getType().cast<VectorType>();
+ Type resultElementType = resultVectorType.getElementType();
+
+ Optional<int> operandId = getVectorContractOpOperandIdForVectorReadOp(op);
+ if (!operandId) {
+ op->emitError() << "Cannot determine operandId this "
+ "vector::TransferReadOp is used as in the "
+ "vector::TransferContractOp";
+ return std::nullopt;
+ }
+
+ // Loading F16 values from Shared Memory to Registers.
+ if (resultElementType.isF16() || resultElementType.isBF16()) {
+ // For matrixC.
+ if (*operandId == 2) {
+ SmallVector<int64_t> readShape;
+ readShape.append({mmaShapeM, mmaShapeN});
+ return readShape;
+ }
+
+ // For matrixA and matrixB.
+ if (*operandId == 0 || *operandId == 1) {
+ // MmaSyncOp input operands: matrixA and matrixB.
+ // LDSMx1, x2, x4:
+ // - LDSMx1 loads a 1 tile of 8x8.
+ // - LDSMx2 loads a 2 tiles of 8x8.
+ // - LDSMx4 loads a 4 tiles of 8x8. (in use)
+ // IREE uses the largest tiled load, i.e., LDSMx4.
+
+ // MmaSyncOp source operand: matrixC.
+ // matrixC is also read/written in tiled block of 16x16. In the pass
+ // OptimizeVectorTransfer, matrixC reads are moved above the mainloop
+ // and writes are moved below the mainloop. Thus, mma.sync read/write
+ // accumulator inplace.
+
+ SmallVector<int64_t> readShape;
+ readShape.append({16, 16});
+ return readShape;
+ }
+ }
+
+ // Loading F32 values from Shared Memory to Registers.
+ if (resultElementType.isF32()) {
+ // Set mmaShapeK for F32 datatype mma.sync.f32.tf32.m16n8k8.
+ mmaShapeK = 8;
+
+ // For matrixC.
+ if (*operandId == 2) {
+ SmallVector<int64_t> readShape;
+ readShape.append({mmaShapeM, mmaShapeN});
+ return readShape;
+ }
+ // For matrixA.
+ if (*operandId == 0) {
+ SmallVector<int64_t> readShape;
+ readShape.append({mmaShapeM, mmaShapeK});
+ return readShape;
+ }
+ // For matrixB.
+ if (*operandId == 1) {
+ // Do not use ldmatrix for matrixB.
+ // Transfer read ops may need different shapes based on how they are
+ // being used. For simplicity just match the shape used by the extract
+ // strided op.
+ VectorType sliceType;
+ for (Operation *users : op->getUsers()) {
+ auto extract = dyn_cast<vector::ExtractStridedSliceOp>(users);
+ if (!extract) return std::nullopt;
+ auto vecType = extract.getResult().getType().cast<VectorType>();
+ if (sliceType && sliceType != vecType) return std::nullopt;
+ sliceType = vecType;
+ }
+ return llvm::to_vector<>(sliceType.getShape());
+ }
+ }
+ }
+ return std::nullopt;
+}
+
} // namespace iree_compiler
} // namespace mlir
diff --git a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h
index ace232c..e113466 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h
+++ b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h
@@ -86,6 +86,9 @@
// TODO: Make this take HW specific sizes.
Optional<SmallVector<int64_t>> getWmmaNativeVectorSize(Operation *op);
+/// Helper function to return native size for MMA.SYNC-based operations.
+Optional<SmallVector<int64_t>> getMmaNativeVectorSize(Operation *op);
+
} // namespace iree_compiler
} // namespace mlir
diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowBase.td b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowBase.td
index 456e6c1..1f6e1b6 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowBase.td
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowBase.td
@@ -156,4 +156,82 @@
let mnemonic = "dummy";
}
+//===----------------------------------------------------------------------===//
+// Flow enums
+//===----------------------------------------------------------------------===//
+
+def FLOW_CollectiveElementType_Sint8 : I32EnumAttrCase<"Sint8", 0, "si8">;
+def FLOW_CollectiveElementType_Uint8 : I32EnumAttrCase<"Uint8", 1, "ui8">;
+def FLOW_CollectiveElementType_Sint16 : I32EnumAttrCase<"Sint16", 2, "si16">;
+def FLOW_CollectiveElementType_Uint16 : I32EnumAttrCase<"Uint16", 3, "ui16">;
+def FLOW_CollectiveElementType_Sint32 : I32EnumAttrCase<"Sint32", 4, "si32">;
+def FLOW_CollectiveElementType_Uint32 : I32EnumAttrCase<"Uint32", 5, "ui32">;
+def FLOW_CollectiveElementType_Sint64 : I32EnumAttrCase<"Sint64", 6, "si64">;
+def FLOW_CollectiveElementType_Uint64 : I32EnumAttrCase<"Uint64", 7, "ui64">;
+def FLOW_CollectiveElementType_Float16 : I32EnumAttrCase<"Float16", 8, "f16">;
+def FLOW_CollectiveElementType_Float32 : I32EnumAttrCase<"Float32", 9, "f32">;
+def FLOW_CollectiveElementType_Float64 : I32EnumAttrCase<"Float64", 10, "f64">;
+def FLOW_CollectiveElementType_BFloat16 : I32EnumAttrCase<"BFloat16", 11, "bf16">;
+def FLOW_CollectiveElementTypeAttr :
+ I32EnumAttr<"CollectiveElementType", "valid CollectiveElementType", [
+ FLOW_CollectiveElementType_Sint8,
+ FLOW_CollectiveElementType_Uint8,
+ FLOW_CollectiveElementType_Sint16,
+ FLOW_CollectiveElementType_Uint16,
+ FLOW_CollectiveElementType_Sint32,
+ FLOW_CollectiveElementType_Uint32,
+ FLOW_CollectiveElementType_Sint64,
+ FLOW_CollectiveElementType_Uint64,
+ FLOW_CollectiveElementType_Float16,
+ FLOW_CollectiveElementType_Float32,
+ FLOW_CollectiveElementType_Float64,
+ FLOW_CollectiveElementType_BFloat16,
+ ]> {
+ let cppNamespace = "::mlir::iree_compiler::IREE::Flow";
+}
+
+//===----------------------------------------------------------------------===//
+// Flow channel type
+//===----------------------------------------------------------------------===//
+
+def FLOW_Channel : TypeDef<Flow_Dialect, "Channel", []> {
+ let mnemonic = "channel";
+ let summary = [{a collecive communication channel}];
+ let description = [{
+ Represents a single participant in a collective clique. Multiple channels
+ may exist within the same program to allow for partial operations or
+ hierarchical operations.
+
+ In programs that have already been partitioned prior to being compiled there
+ will often exist only one channel and `flow.channel.default` can be used
+ to reference it. In programs that model SPMD behavior internally channels
+ can be created or provided by hosting applications.
+ }];
+}
+
+//===----------------------------------------------------------------------===//
+// Flow collective reduction op
+//===----------------------------------------------------------------------===//
+
+// NOTE: the enum values must exactly match with the corresponding enum values
+// of the Stream reduction op.
+
+def FLOW_CollectiveReductionOp_None : I32EnumAttrCase<"None", 0, "none">;
+def FLOW_CollectiveReductionOp_ReductionSum : I32EnumAttrCase<"ReductionSum", 1, "sum">;
+def FLOW_CollectiveReductionOp_ReductionProduct : I32EnumAttrCase<"ReductionProduct", 2, "product">;
+def FLOW_CollectiveReductionOp_ReductionMinimum : I32EnumAttrCase<"ReductionMinimum", 3, "minimum">;
+def FLOW_CollectiveReductionOp_ReductionMaximum : I32EnumAttrCase<"ReductionMaximum", 4, "maximum">;
+def FLOW_CollectiveReductionOp_ReductionAverage : I32EnumAttrCase<"ReductionAverage", 5, "average">;
+def FLOW_CollectiveReductionOpAttr :
+ I32EnumAttr<"CollectiveReductionOp", "valid CollectiveReductionOp", [
+ FLOW_CollectiveReductionOp_None,
+ FLOW_CollectiveReductionOp_ReductionSum,
+ FLOW_CollectiveReductionOp_ReductionProduct,
+ FLOW_CollectiveReductionOp_ReductionMinimum,
+ FLOW_CollectiveReductionOp_ReductionMaximum,
+ FLOW_CollectiveReductionOp_ReductionAverage,
+ ]> {
+ let cppNamespace = "mlir::iree_compiler::IREE::Flow";
+}
+
#endif // IREE_DIALECT_FLOW_BASE
diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp
index 244dbd4..46596ff 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp
@@ -1545,6 +1545,118 @@
context);
}
+//===----------------------------------------------------------------------===//
+// flow.channel.count
+//===----------------------------------------------------------------------===//
+
+void ChannelCountOp::getAsmResultNames(
+ function_ref<void(Value, StringRef)> setNameFn) {
+ setNameFn(getResult(), "channel_count");
+}
+
+//===----------------------------------------------------------------------===//
+// flow.channel.default
+//===----------------------------------------------------------------------===//
+
+void ChannelDefaultOp::getAsmResultNames(
+ function_ref<void(Value, StringRef)> setNameFn) {
+ setNameFn(getResult(), "channel_default");
+}
+
+//===----------------------------------------------------------------------===//
+// flow.channel.rank
+//===----------------------------------------------------------------------===//
+
+void ChannelRankOp::getAsmResultNames(
+ function_ref<void(Value, StringRef)> setNameFn) {
+ setNameFn(getResult(), "channel_rank");
+}
+
+//===----------------------------------------------------------------------===//
+// flow.collective.all_gather
+//===----------------------------------------------------------------------===//
+
+Value CollectiveAllGatherOp::getTiedResult(unsigned resultIndex) {
+ return IREE::Util::TiedOpInterface::findTiedBaseValue(getTarget());
+}
+
+::llvm::Optional<unsigned> CollectiveAllGatherOp::getTiedResultOperandIndex(
+ unsigned resultIndex) {
+ return {0}; // target
+}
+
+SmallVector<int64_t, 4> CollectiveAllGatherOp::getTiedResultOperandIndices() {
+ return {0}; // target
+}
+
+void CollectiveAllGatherOp::build(OpBuilder &builder, OperationState &state,
+ CollectiveElementTypeAttr elementType,
+ Value target, Value source, Value channel) {
+ auto targetDims =
+ IREE::Util::buildDynamicDimsForValue(state.location, target, builder);
+
+ build(builder, state, elementType, target, targetDims, source, channel,
+ builder.getIndexArrayAttr({0}));
+}
+
+//===----------------------------------------------------------------------===//
+// flow.collective.all_reduce
+//===----------------------------------------------------------------------===//
+
+Value CollectiveAllReduceOp::getTiedResult(unsigned resultIndex) {
+ return IREE::Util::TiedOpInterface::findTiedBaseValue(getTarget());
+}
+
+::llvm::Optional<unsigned> CollectiveAllReduceOp::getTiedResultOperandIndex(
+ unsigned resultIndex) {
+ return {0}; // target
+}
+
+SmallVector<int64_t, 4> CollectiveAllReduceOp::getTiedResultOperandIndices() {
+ return {0}; // target
+}
+
+void CollectiveAllReduceOp::build(OpBuilder &builder, OperationState &state,
+ CollectiveReductionOpAttr reductionOp,
+ CollectiveElementTypeAttr elementType,
+ Value target, Value source, Value channel) {
+ auto targetDims =
+ IREE::Util::buildDynamicDimsForValue(state.location, target, builder);
+
+ build(builder, state, reductionOp, elementType, target, targetDims, source,
+ channel, builder.getIndexArrayAttr({0}));
+}
+
+//===----------------------------------------------------------------------===//
+// flow.collective.reduce_scatter
+//===----------------------------------------------------------------------===//
+
+Value CollectiveReduceScatterOp::getTiedResult(unsigned resultIndex) {
+ return IREE::Util::TiedOpInterface::findTiedBaseValue(getTarget());
+}
+
+::llvm::Optional<unsigned> CollectiveReduceScatterOp::getTiedResultOperandIndex(
+ unsigned resultIndex) {
+ return {0}; // target
+}
+
+SmallVector<int64_t, 4>
+CollectiveReduceScatterOp::getTiedResultOperandIndices() {
+ return {0}; // target
+}
+
+void CollectiveReduceScatterOp::build(OpBuilder &builder, OperationState &state,
+ CollectiveReductionOpAttr reductionOp,
+ CollectiveElementTypeAttr elementType,
+ Value target, Value source,
+ Value channel) {
+ auto targetDims =
+ IREE::Util::buildDynamicDimsForValue(state.location, target, builder);
+
+ build(builder, state, reductionOp, elementType, target, targetDims, source,
+ channel, builder.getIndexArrayAttr({0}));
+}
+
} // namespace Flow
} // namespace IREE
} // namespace iree_compiler
diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td
index 5e1acab..e6f8c29 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td
@@ -1304,4 +1304,185 @@
}];
}
+//===----------------------------------------------------------------------===//
+// Collective communication ops
+//===----------------------------------------------------------------------===//
+
+def FLOW_ChannelDefaultOp : FLOW_Op<"channel.default", [
+ DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>
+]> {
+ let summary = [{returns a default collective communication channel}];
+ let description = [{
+ Returns a channel initialized using the runtime environment.
+ }];
+
+ let results = (outs
+ FLOW_Channel:$result
+ );
+
+ let assemblyFormat = [{
+ `:` type($result)
+ attr-dict-with-keyword
+ }];
+}
+
+def FLOW_ChannelRankOp : FLOW_Op<"channel.rank", [
+ DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
+]> {
+ let summary = [{returns the rank of the local participant in the group}];
+ let description = [{
+ Returns the rank the channel represents as a participant in a collective
+ group in `[0, count)`.
+ }];
+
+ let arguments = (ins
+ FLOW_Channel:$channel
+ );
+ let results = (outs
+ Index:$result
+ );
+
+ let assemblyFormat = [{
+ $channel `:` type($result)
+ attr-dict-with-keyword
+ }];
+}
+
+def FLOW_ChannelCountOp : FLOW_Op<"channel.count", [
+ DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
+]> {
+ let summary = [{returns the total number of participants in the group}];
+ let description = [{
+ Returns the total participant count in the collective communicator group.
+ }];
+
+ let arguments = (ins
+ FLOW_Channel:$channel
+ );
+ let results = (outs
+ Index:$result
+ );
+
+ let assemblyFormat = [{
+ $channel `:` type($result)
+ attr-dict-with-keyword
+ }];
+}
+
+def FLOW_CollectiveAllGatherOp : FLOW_Op<"collective.all_gather", [
+ AllTypesMatch<["target", "result"]>,
+ DeclareOpInterfaceMethods<Util_TiedOpInterface, [
+ "getTiedResult",
+ "getTiedResultOperandIndex",
+ "getTiedResultOperandIndices",
+ ]>,
+]> {
+ let summary = [{performs all-gather operation}];
+ let description = [{It gathers data from all ranks and concatenates them on the 0-th dimension.}];
+
+ let arguments = (ins
+ FLOW_CollectiveElementTypeAttr:$element_type,
+ FLOW_Tensor:$target,
+ FLOW_ShapeDynamicDims:$target_dims,
+ FLOW_Tensor:$source,
+ FLOW_Channel:$channel,
+ OptionalAttr<Util_TiedOpStorageAttr>:$tied_operands
+ );
+ let results = (outs
+ FLOW_Tensor:$result
+ );
+ let assemblyFormat = [{
+ $element_type `,` $target `,` $source `,` $channel `:`
+ `(` type($target) `,` type($source) `,` type($channel) `)` `->`
+ custom<ShapedTiedResult>(type($result), $target_dims, $tied_operands)
+ attr-dict-with-keyword
+ }];
+ let builders = [
+ OpBuilder<(ins
+ "CollectiveElementTypeAttr":$element_type,
+ "Value":$target,
+ "Value":$source,
+ "Value":$channel)>,
+ ];
+}
+
+def FLOW_CollectiveAllReduceOp : FLOW_Op<"collective.all_reduce", [
+ AllTypesMatch<["source", "target", "result"]>,
+ DeclareOpInterfaceMethods<Util_TiedOpInterface, [
+ "getTiedResult",
+ "getTiedResultOperandIndex",
+ "getTiedResultOperandIndices",
+ ]>,
+]> {
+ let summary = [{performs all-reduce operation}];
+ let description = [{The operation reduces data across all the ranks in the channel.}];
+
+ let arguments = (ins
+ FLOW_CollectiveReductionOpAttr:$reduction_op,
+ FLOW_CollectiveElementTypeAttr:$element_type,
+ FLOW_Tensor:$target,
+ FLOW_ShapeDynamicDims:$target_dims,
+ FLOW_Tensor:$source,
+ FLOW_Channel:$channel,
+ OptionalAttr<Util_TiedOpStorageAttr>:$tied_operands
+ );
+ let results = (outs
+ FLOW_Tensor:$result
+ );
+ let assemblyFormat = [{
+ $reduction_op `,` $element_type `,` $target `,` $source `,` $channel `:`
+ `(` type($target) `,` type($source) `,` type($channel) `)` `->`
+ custom<ShapedTiedResult>(type($result), $target_dims, $tied_operands)
+ attr-dict-with-keyword
+ }];
+ let builders = [
+ OpBuilder<(ins
+ "CollectiveReductionOpAttr":$reduction_op,
+ "CollectiveElementTypeAttr":$element_type,
+ "Value":$target,
+ "Value":$source,
+ "Value":$channel)>,
+ ];
+}
+
+def FLOW_CollectiveReduceScatterOp : FLOW_Op<"collective.reduce_scatter", [
+ AllTypesMatch<["target", "result"]>,
+ DeclareOpInterfaceMethods<Util_TiedOpInterface, [
+ "getTiedResult",
+ "getTiedResultOperandIndex",
+ "getTiedResultOperandIndices",
+ ]>,
+]> {
+ let summary = [{performs reduce and scatter operations}];
+ let description = [{The operation reduces data across all the ranks in the channel and
+ scatters the result to each rank.}];
+
+ let arguments = (ins
+ FLOW_CollectiveReductionOpAttr:$reduction_op,
+ FLOW_CollectiveElementTypeAttr:$element_type,
+ FLOW_Tensor:$target,
+ FLOW_ShapeDynamicDims:$target_dims,
+ FLOW_Tensor:$source,
+ FLOW_Channel:$channel,
+ OptionalAttr<Util_TiedOpStorageAttr>:$tied_operands
+ );
+ let results = (outs
+ FLOW_Tensor:$result
+ );
+ let assemblyFormat = [{
+ $reduction_op `,` $element_type `,` $target `,` $source `,` $channel `:`
+ `(` type($target) `,` type($source) `,` type($channel) `)` `->`
+ custom<ShapedTiedResult>(type($result), $target_dims, $tied_operands)
+ attr-dict-with-keyword
+ }];
+ let builders = [
+ OpBuilder<(ins
+ "CollectiveReductionOpAttr":$reduction_op,
+ "CollectiveElementTypeAttr":$element_type,
+ "Value":$target,
+ "Value":$source,
+ "Value":$channel)>,
+ ];
+}
+
#endif // IREE_DIALECT_FLOW_OPS
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/CollapseReductionDims.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/CollapseReductionDims.cpp
index 863b855..7161888 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/CollapseReductionDims.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/CollapseReductionDims.cpp
@@ -31,7 +31,8 @@
}
// Check that the following dimensions are match the order of `dims`
for (unsigned j = 1, numDims = dims.size(); j < numDims; j++) {
- if (map.getDimPosition(i + j) != dims[j]) {
+ unsigned pos = i + j;
+ if (pos >= map.getNumResults() || map.getDimPosition(pos) != dims[j]) {
return false;
}
}
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/collapse_reduction.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/collapse_reduction.mlir
index 2631e4e..25484ec 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/collapse_reduction.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/collapse_reduction.mlir
@@ -18,3 +18,22 @@
// Check that we collapse dimensions.
// CHECK: @multi_reduce_dim
// CHECK: linalg.generic {{.*}} iterator_types = ["parallel", "parallel", "reduction"]
+
+// -----
+
+// Collapsing is not supported when an input is broadcasted; we can't collapse
+// the input from tensor<4xf32> to tensor<32xf32> for example.
+
+func.func @input_broadcast(%arg0: tensor<4x8xf32>, %arg1: tensor<4xf32>) -> tensor<f32> {
+ %empty = tensor.empty() : tensor<f32>
+ %reduce = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> ()>], iterator_types = ["reduction", "reduction"]} ins(%arg0, %arg1 : tensor<4x8xf32>, tensor<4xf32>) outs(%empty : tensor<f32>) {
+ ^bb0(%arg2: f32, %arg3: f32, %out: f32):
+ %div = arith.divf %arg2, %arg3 : f32
+ %add = arith.addf %out, %div : f32
+ linalg.yield %add : f32
+ } -> tensor<f32>
+ return %reduce : tensor<f32>
+}
+
+// CHECK: @input_broadcast
+// CHECK-NOT: tensor.collapse_shape
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
index 3a9fb2f..5565bdf 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
@@ -482,6 +482,152 @@
}
};
+// -----------------------------------------------------------------------------
+// Collective Ops
+// -----------------------------------------------------------------------------
+
+struct ConvertAllGatherOp
+ : public OpConversionPattern<IREE::Flow::CollectiveAllGatherOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::CollectiveAllGatherOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto shape = op.getSource().getType().cast<ShapedType>();
+ auto collectiveAttr = IREE::Stream::CollectiveAttr::get(
+ op.getContext(), IREE::Stream::CollectiveKind::AllGather,
+ /*reduction=*/std::nullopt,
+ static_cast<IREE::Stream::CollectiveElementType>(op.getElementType()));
+
+ auto zeroOffset = rewriter.create<arith::ConstantIndexOp>(op.getLoc(), 0);
+ auto elementCount = rewriter.create<arith::ConstantIndexOp>(
+ op.getLoc(), shape.getNumElements());
+ auto newTargetCast =
+ consumeTensorOperand(op.getLoc(), adaptor.getTarget(), rewriter);
+ auto newSourceCast =
+ consumeTensorOperand(op.getLoc(), adaptor.getSource(), rewriter);
+
+ rewriter.replaceOpWithNewOp<IREE::Stream::AsyncCollectiveOp>(
+ op, collectiveAttr, adaptor.getTarget(),
+ /*target_size=*/newTargetCast.resourceSize,
+ /*target_offset=*/zeroOffset,
+ /*target_end=*/newTargetCast.resourceSize,
+ /*target_length=*/newTargetCast.resourceSize, adaptor.getSource(),
+ /*source_size=*/newSourceCast.resourceSize,
+ /*source_offset=*/zeroOffset, /*source_end=*/newSourceCast.resourceSize,
+ /*source_length=*/newSourceCast.resourceSize, elementCount,
+ adaptor.getChannel(),
+ /*param=*/mlir::Value(), getAffinityFor(op));
+ return success();
+ }
+};
+
+struct ConvertAllReduceOp
+ : public OpConversionPattern<IREE::Flow::CollectiveAllReduceOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::CollectiveAllReduceOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto shape = op.getType().cast<ShapedType>();
+ auto collectiveAttr = IREE::Stream::CollectiveAttr::get(
+ op.getContext(), IREE::Stream::CollectiveKind::AllReduce,
+ static_cast<IREE::Stream::CollectiveReductionOp>(op.getReductionOp()),
+ static_cast<IREE::Stream::CollectiveElementType>(op.getElementType()));
+
+ auto zeroOffset = rewriter.create<arith::ConstantIndexOp>(op.getLoc(), 0);
+ auto elementCount = rewriter.create<arith::ConstantIndexOp>(
+ op.getLoc(), shape.getNumElements());
+ auto newTargetCast =
+ consumeTensorOperand(op.getLoc(), adaptor.getTarget(), rewriter);
+ auto newSourceCast =
+ consumeTensorOperand(op.getLoc(), adaptor.getSource(), rewriter);
+
+ rewriter.replaceOpWithNewOp<IREE::Stream::AsyncCollectiveOp>(
+ op, collectiveAttr, adaptor.getTarget(),
+ /*target_size=*/newTargetCast.resourceSize,
+ /*target_offset=*/zeroOffset,
+ /*target_end=*/newTargetCast.resourceSize,
+ /*target_length=*/newTargetCast.resourceSize, adaptor.getSource(),
+ /*source_size=*/newSourceCast.resourceSize,
+ /*source_offset=*/zeroOffset, /*source_end=*/newSourceCast.resourceSize,
+ /*source_length=*/newSourceCast.resourceSize, elementCount,
+ adaptor.getChannel(),
+ /*param=*/mlir::Value(), getAffinityFor(op));
+ return success();
+ }
+};
+
+struct ConvertReduceScatterOp
+ : public OpConversionPattern<IREE::Flow::CollectiveReduceScatterOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::CollectiveReduceScatterOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto shape = op.getType().cast<ShapedType>();
+ auto collectiveAttr = IREE::Stream::CollectiveAttr::get(
+ op.getContext(), IREE::Stream::CollectiveKind::ReduceScatter,
+ static_cast<IREE::Stream::CollectiveReductionOp>(op.getReductionOp()),
+ static_cast<IREE::Stream::CollectiveElementType>(op.getElementType()));
+
+ auto zeroOffset = rewriter.create<arith::ConstantIndexOp>(op.getLoc(), 0);
+ auto elementCount = rewriter.create<arith::ConstantIndexOp>(
+ op.getLoc(), shape.getNumElements());
+ auto newTargetCast =
+ consumeTensorOperand(op.getLoc(), adaptor.getTarget(), rewriter);
+ auto newSourceCast =
+ consumeTensorOperand(op.getLoc(), adaptor.getSource(), rewriter);
+
+ rewriter.replaceOpWithNewOp<IREE::Stream::AsyncCollectiveOp>(
+ op, collectiveAttr, adaptor.getTarget(),
+ /*target_size=*/newTargetCast.resourceSize,
+ /*target_offset=*/zeroOffset,
+ /*target_end=*/newTargetCast.resourceSize,
+ /*target_length=*/newTargetCast.resourceSize, adaptor.getSource(),
+ /*source_size=*/newSourceCast.resourceSize,
+ /*source_offset=*/zeroOffset, /*source_end=*/newSourceCast.resourceSize,
+ /*source_length=*/newSourceCast.resourceSize, elementCount,
+ adaptor.getChannel(),
+ /*param=*/mlir::Value(), getAffinityFor(op));
+ return success();
+ }
+};
+
+struct ConvertChannelDefaultOp
+ : public OpConversionPattern<IREE::Flow::ChannelDefaultOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::ChannelDefaultOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ IREE::Stream::AffinityAttr affinityAttr;
+ rewriter.replaceOpWithNewOp<IREE::Stream::ChannelDefaultOp>(op,
+ affinityAttr);
+ return success();
+ }
+};
+
+struct ConvertChannelCountOp
+ : public OpConversionPattern<IREE::Flow::ChannelCountOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::ChannelCountOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ rewriter.replaceOpWithNewOp<IREE::Stream::ChannelCountOp>(
+ op, adaptor.getOperands());
+ return success();
+ }
+};
+
+struct ConvertChannelRankOp
+ : public OpConversionPattern<IREE::Flow::ChannelRankOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::ChannelRankOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ rewriter.replaceOpWithNewOp<IREE::Stream::ChannelRankOp>(
+ op, adaptor.getOperands());
+ return success();
+ }
+};
+
} // namespace
void populateFlowToStreamConversionPatterns(MLIRContext *context,
@@ -495,6 +641,13 @@
patterns.insert<ConvertDispatchOp>(typeConverter, context);
patterns.insert<ConvertExecutableOp>(typeConverter, context);
patterns.insert<ConvertReturnOp>(typeConverter, context);
+ // collective ops
+ patterns.insert<ConvertAllGatherOp>(typeConverter, context);
+ patterns.insert<ConvertAllReduceOp>(typeConverter, context);
+ patterns.insert<ConvertChannelCountOp>(typeConverter, context);
+ patterns.insert<ConvertChannelDefaultOp>(typeConverter, context);
+ patterns.insert<ConvertChannelRankOp>(typeConverter, context);
+ patterns.insert<ConvertReduceScatterOp>(typeConverter, context);
}
void populateFlowToStreamConversionPatterns(MLIRContext *context,
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/BUILD b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/BUILD
index c027a47..b6a902b 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/BUILD
+++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/BUILD
@@ -16,6 +16,7 @@
name = "lit",
srcs = enforce_glob(
[
+ "collective_ops.mlir",
"dispatch_ops.mlir",
"executable_ops.mlir",
"tensor_ops.mlir",
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/CMakeLists.txt
index 82ea148..c038d00 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/CMakeLists.txt
@@ -14,6 +14,7 @@
NAME
lit
SRCS
+ "collective_ops.mlir"
"dispatch_ops.mlir"
"executable_ops.mlir"
"tensor_ops.mlir"
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/collective_ops.mlir b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/collective_ops.mlir
new file mode 100644
index 0000000..9e800f6
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/collective_ops.mlir
@@ -0,0 +1,68 @@
+// RUN: iree-opt --split-input-file --iree-stream-conversion %s | FileCheck %s
+
+// CHECK-LABEL: @channel_count
+func.func @channel_count() -> index {
+ // CHECK: [[CHANNEL:%.+]] = stream.channel.default : !stream.channel
+ // CHECK: [[COUNT:%.+]] = stream.channel.count [[CHANNEL]] : index
+ // CHECK: return [[COUNT]] : index
+ %channel_default = flow.channel.default : !flow.channel
+ %count = flow.channel.count %channel_default : index
+ return %count : index
+}
+
+//-----
+
+// CHECK-LABEL: @channel_rank
+func.func @channel_rank() -> index {
+ // CHECK: [[CHANNEL:%.+]] = stream.channel.default : !stream.channel
+ // CHECK: [[RANK:%.+]] = stream.channel.rank [[CHANNEL]] : index
+ // CHECK: return [[RANK]] : index
+ %channel_default = flow.channel.default : !flow.channel
+ %rank = flow.channel.rank %channel_default : index
+ return %rank : index
+}
+
+//-----
+
+// CHECK-LABEL: @all_reduce_sum
+func.func @all_reduce_sum(%arg0: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub} {
+ // CHECK: stream.channel.default
+ // CHECK: stream.tensor.empty : tensor<2304xf32>
+ // CHECK: stream.async.collective<all_reduce with sum : f32>
+ %0 = hal.tensor.import %arg0 : !hal.buffer_view -> tensor<2304xf32>
+ %channel_default = flow.channel.default : !flow.channel
+ %1 = flow.tensor.empty : tensor<2304xf32>
+ %2 = flow.collective.all_reduce sum, f32, %1, %0, %channel_default : (tensor<2304xf32>, tensor<2304xf32>, !flow.channel) -> tensor<2304xf32>
+ %3 = hal.tensor.export %2 : tensor<2304xf32> -> !hal.buffer_view
+ return %3 : !hal.buffer_view
+}
+
+//-----
+
+// CHECK-LABEL: @allgather
+func.func @allgather(%arg0: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub} {
+ // CHECK: stream.channel.default
+ // CHECK: stream.tensor.empty : tensor<1024xf32>
+ // CHECK: stream.async.collective<all_gather : f32>
+ %0 = hal.tensor.import %arg0 : !hal.buffer_view -> tensor<512xf32>
+ %channel_default = flow.channel.default : !flow.channel
+ %1 = flow.tensor.empty : tensor<1024xf32>
+ %2 = flow.collective.all_gather f32, %1, %0, %channel_default : (tensor<1024xf32>, tensor<512xf32>, !flow.channel) -> tensor<1024xf32>
+ %3 = hal.tensor.export %2 : tensor<1024xf32> -> !hal.buffer_view
+ return %3 : !hal.buffer_view
+}
+
+//-----
+
+// CHECK-LABEL: @reduce_scatter
+func.func @reduce_scatter(%arg0: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub} {
+ // CHECK: stream.channel.default
+ // CHECK: stream.tensor.empty : tensor<2x2xf32>
+ // CHECK: stream.async.collective<reduce_scatter with sum : f32>
+ %0 = hal.tensor.import %arg0 : !hal.buffer_view -> tensor<4x2xf32>
+ %channel_default = flow.channel.default : !flow.channel
+ %1 = flow.tensor.empty : tensor<2x2xf32>
+ %2 = flow.collective.reduce_scatter sum, f32, %1, %0, %channel_default : (tensor<2x2xf32>, tensor<4x2xf32>, !flow.channel) -> tensor<2x2xf32>
+ %3 = hal.tensor.export %2 : tensor<2x2xf32> -> !hal.buffer_view
+ return %3 : !hal.buffer_view
+}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp
index baa36a5..0b0acc8 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp
@@ -5,6 +5,7 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
+#include "iree/compiler/Dialect/Flow/IR/FlowTypes.h"
#include "iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.h"
#include "iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.h"
#include "iree/compiler/Dialect/Stream/Conversion/PatternUtils.h"
@@ -202,8 +203,13 @@
// Allow unknown types to pass through; these come from custom dialects that
// may be mixed into the IR we are converting.
- typeConverter.addConversion(
- [](Type type) { return !type.isa<TensorType>() ? type : Type{}; });
+ typeConverter.addConversion([=](Type type) -> Type {
+ // convert flow.channel into stream.channel
+ if (type.isa<IREE::Flow::ChannelType>())
+ return IREE::Stream::ChannelType::get(context);
+
+ return !type.isa<TensorType>() ? type : Type{};
+ });
// Disallow tensor dialects; the goal here is to remove all tensors and
// turn them into stream resource ops.
diff --git a/compiler/src/iree/compiler/InputConversion/MHLO/BUILD b/compiler/src/iree/compiler/InputConversion/MHLO/BUILD
index 4127a84..2de86f4 100644
--- a/compiler/src/iree/compiler/InputConversion/MHLO/BUILD
+++ b/compiler/src/iree/compiler/InputConversion/MHLO/BUILD
@@ -47,6 +47,7 @@
name = "MHLO",
srcs = [
"BroadcastingToLinalgPatterns.cpp",
+ "ConvertCollectiveOps.cpp",
"ConvertComplexToReal.cpp",
"ConvertMHLOToFlow.cpp",
"ConvertMHLOToFlow.h",
diff --git a/compiler/src/iree/compiler/InputConversion/MHLO/CMakeLists.txt b/compiler/src/iree/compiler/InputConversion/MHLO/CMakeLists.txt
index c713b78..26d37ce 100644
--- a/compiler/src/iree/compiler/InputConversion/MHLO/CMakeLists.txt
+++ b/compiler/src/iree/compiler/InputConversion/MHLO/CMakeLists.txt
@@ -41,6 +41,7 @@
"Passes.h"
SRCS
"BroadcastingToLinalgPatterns.cpp"
+ "ConvertCollectiveOps.cpp"
"ConvertComplexToReal.cpp"
"ConvertMHLOToFlow.cpp"
"ConvertMHLOToFlow.h"
diff --git a/compiler/src/iree/compiler/InputConversion/MHLO/ConvertCollectiveOps.cpp b/compiler/src/iree/compiler/InputConversion/MHLO/ConvertCollectiveOps.cpp
new file mode 100644
index 0000000..94607d9
--- /dev/null
+++ b/compiler/src/iree/compiler/InputConversion/MHLO/ConvertCollectiveOps.cpp
@@ -0,0 +1,449 @@
+// Copyright 2023 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/compiler/Dialect/Flow/IR/FlowTypes.h>
+
+#include <optional>
+
+#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
+#include "iree/compiler/InputConversion/MHLO/PassDetail.h"
+#include "iree/compiler/InputConversion/MHLO/Passes.h"
+#include "iree/compiler/InputConversion/MHLO/Rewriters.h"
+#include "mhlo/IR/hlo_ops.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace MHLO {
+
+// Work in progress. The implementation is planned as several stages.
+//
+// For the first stage, a few simplifications are made to support simple models.
+//
+// 1. Single stream with deterministic order of execution
+// 2. Single replica group for all collective ops
+// 3. Only replicas without partition_id used
+//
+// These allow us to use a default channel for all communications, and there is
+// 1:1 mapping from the replica IDs to the communication ranks. The attribute,
+// use_global_device_ids, is always set in this case.
+//
+// The next stage is to support multiple replica groups. This needs a channel
+// creation with a subset of processes, which should have another communication
+// among the group. A possible strategy is to have the root process in the group
+// (the first rank of the group) creates a channel and the other processes query
+// the channel info from the root process. A key-value store using gRPC might be
+// a good solution.
+//
+// Supporting partition_id comes next. This includes the support for various
+// mode combinations for cross-replica and cross partition communication. See
+// the stablehlo specification for more details about the different modes.
+
+namespace {
+
+static std::optional<IREE::Flow::CollectiveElementType>
+convertToFlowCollectiveElementType(Type type) {
+ if (type.isF32()) {
+ return IREE::Flow::CollectiveElementType::Float32;
+ }
+
+ if (type.isInteger(32)) {
+ if (type.isSignedInteger()) {
+ return IREE::Flow::CollectiveElementType::Sint32;
+ } else {
+ return IREE::Flow::CollectiveElementType::Uint32;
+ }
+ }
+
+ if (type.isF16()) {
+ return IREE::Flow::CollectiveElementType::Float16;
+ }
+
+ if (type.isInteger(8)) {
+ if (type.isSignedInteger()) {
+ return IREE::Flow::CollectiveElementType::Sint8;
+ } else {
+ return IREE::Flow::CollectiveElementType::Uint8;
+ }
+ }
+
+ if (type.isInteger(16)) {
+ if (type.isSignedInteger()) {
+ return IREE::Flow::CollectiveElementType::Sint16;
+ } else {
+ return IREE::Flow::CollectiveElementType::Uint16;
+ }
+ }
+
+ if (type.isBF16()) {
+ return IREE::Flow::CollectiveElementType::BFloat16;
+ }
+
+ if (type.isF64()) {
+ return IREE::Flow::CollectiveElementType::Float64;
+ }
+
+ if (type.isInteger(64)) {
+ if (type.isSignedInteger()) {
+ return IREE::Flow::CollectiveElementType::Sint64;
+ } else {
+ return IREE::Flow::CollectiveElementType::Uint64;
+ }
+ }
+
+ return std::nullopt;
+}
+
+static std::optional<IREE::Flow::CollectiveReductionOp>
+convertToFlowCollectiveReductionOp(const Operation &op) {
+ if (isa<mhlo::AddOp>(op)) {
+ return IREE::Flow::CollectiveReductionOp::ReductionSum;
+ } else if (isa<mhlo::MulOp>(op)) {
+ return IREE::Flow::CollectiveReductionOp::ReductionProduct;
+ } else if (isa<mhlo::MinOp>(op)) {
+ return IREE::Flow::CollectiveReductionOp::ReductionMinimum;
+ } else if (isa<mhlo::MaxOp>(op)) {
+ return IREE::Flow::CollectiveReductionOp::ReductionMaximum;
+ } else {
+ // TODO: we may be able to detect an average operation and convert it
+ // into IREE::Flow::CollectiveReductionOp::ReductionAverage.
+ return std::nullopt;
+ }
+}
+
+static IREE::Flow::CollectiveElementTypeAttr getCollectiveElementTypeAttr(
+ MLIRContext *context, RankedTensorType type) {
+ std::optional<IREE::Flow::CollectiveElementType> collectiveElemType =
+ convertToFlowCollectiveElementType(type.getElementType());
+ if (!collectiveElemType) {
+ return IREE::Flow::CollectiveElementTypeAttr();
+ }
+ return IREE::Flow::CollectiveElementTypeAttr::get(context,
+ *collectiveElemType);
+}
+
+} // namespace
+
+/// Converts mhlo.replica_id to flow.channel.default + flow.channel.rank.
+/// TODO(okkwon): this assumes that there is no partition so that there is a 1:1
+/// mapping between the replica ID and the process ID.
+struct ReplicaIdOpConversion : public OpConversionPattern<mhlo::ReplicaIdOp> {
+ using OpConversionPattern<mhlo::ReplicaIdOp>::OpConversionPattern;
+
+ LogicalResult matchAndRewrite(
+ mhlo::ReplicaIdOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto channel = rewriter.create<IREE::Flow::ChannelDefaultOp>(loc);
+ auto rank = rewriter.create<IREE::Flow::ChannelRankOp>(loc, channel);
+ auto resultType = op.getType().cast<RankedTensorType>(); // tensor<ui32>
+ auto elemType = resultType.getElementType();
+ // index -> ui32
+ auto rankElem = rewriter.create<arith::IndexCastUIOp>(loc, elemType, rank);
+ // tensor<ui32>
+ auto rankTensor = rewriter.create<tensor::FromElementsOp>(
+ loc, resultType, rankElem.getResult());
+ rewriter.replaceOp(op, rankTensor.getResult());
+ return success();
+ }
+};
+
+struct AllGatherOpConversion : public OpConversionPattern<mhlo::AllGatherOp> {
+ using OpConversionPattern<mhlo::AllGatherOp>::OpConversionPattern;
+
+ LogicalResult matchAndRewrite(
+ mhlo::AllGatherOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+
+ if (!op.getUseGlobalDeviceIds()) {
+ return rewriter.notifyMatchFailure(op, "must use global device IDs");
+ }
+
+ // Check there is only one group in the replica_groups
+ ShapedType replicaGroupType = op.getReplicaGroups().getType();
+ if (replicaGroupType.getRank() != 2 ||
+ replicaGroupType.getDimSize(0) != 1) {
+ return rewriter.notifyMatchFailure(op,
+ "must have a single replica group");
+ }
+
+ // Currently only the default channel is used.
+
+ // Create a default channel.
+ auto channel = rewriter.create<IREE::Flow::ChannelDefaultOp>(loc);
+
+ // Get the collective element type attribute.
+ auto resultType = op.getResult().getType().cast<RankedTensorType>();
+ IREE::Flow::CollectiveElementTypeAttr elementTypeAttr =
+ getCollectiveElementTypeAttr(op.getContext(), resultType);
+ if (!elementTypeAttr) {
+ return rewriter.notifyMatchFailure(
+ op, "unsupported element type for collective op");
+ }
+
+ // When all_gather_dim != 0, we need to transpose between 0 and
+ // all_gather_dim before and after the flow allgather op.
+ uint64_t allGatherDim = op.getAllGatherDim();
+ auto inputType = op.getOperand().getType().cast<RankedTensorType>();
+ SmallVector<int64_t> gatherInputShape(inputType.getShape());
+ Value gatherInput = op.getOperand();
+ DenseIntElementsAttr permutationAttr;
+ SmallVector<int64_t> gatherResultShape(resultType.getShape());
+
+ if (allGatherDim != 0) {
+ SmallVector<int64_t> permutation =
+ llvm::to_vector(llvm::seq<int64_t>(0, gatherResultShape.size()));
+ std::swap(permutation[0], permutation[allGatherDim]);
+ permutationAttr = rewriter.getI64VectorAttr(permutation);
+ std::swap(gatherInputShape[0], gatherInputShape[allGatherDim]);
+ std::swap(gatherResultShape[0], gatherResultShape[allGatherDim]);
+ // Transpose the input.
+ gatherInput = rewriter
+ .create<mhlo::TransposeOp>(
+ loc,
+ RankedTensorType::get(gatherInputShape,
+ resultType.getElementType()),
+ gatherInput, permutationAttr)
+ .getResult();
+ }
+
+ // Create an empty tensor for the result.
+ Value target = rewriter.create<tensor::EmptyOp>(
+ loc, gatherResultShape, resultType.getElementType());
+ Value gatherResult =
+ rewriter
+ .create<IREE::Flow::CollectiveAllGatherOp>(
+ op.getLoc(), elementTypeAttr, target, gatherInput, channel)
+ .getResult();
+
+ if (allGatherDim != 0) {
+ gatherResult = rewriter
+ .create<mhlo::TransposeOp>(
+ loc, resultType, gatherResult, permutationAttr)
+ .getResult();
+ }
+
+ rewriter.replaceOp(op, gatherResult);
+ return success();
+ }
+};
+
+struct AllReduceOpConversion : public OpConversionPattern<mhlo::AllReduceOp> {
+ using OpConversionPattern<mhlo::AllReduceOp>::OpConversionPattern;
+
+ LogicalResult matchAndRewrite(
+ mhlo::AllReduceOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+
+ if (!op.getUseGlobalDeviceIds()) {
+ return rewriter.notifyMatchFailure(op, "must use global device IDs");
+ }
+
+ // Check there is only one group in the replica_groups.
+ ShapedType replicaGroupType = op.getReplicaGroups().getType();
+ if (replicaGroupType.getRank() != 2 ||
+ replicaGroupType.getDimSize(0) != 1) {
+ return rewriter.notifyMatchFailure(op,
+ "must have a single replica group");
+ }
+
+ // Only single elementwise op is supported.
+ Block &block = op.getComputation().front();
+
+ if (block.empty() || llvm::hasSingleElement(block) ||
+ std::next(block.begin(), 2) != block.end()) {
+ return rewriter.notifyMatchFailure(op, "must have two ops in the block");
+ }
+
+ if (block.getNumArguments() != 2) {
+ return rewriter.notifyMatchFailure(op, "must have two block args");
+ }
+
+ Operation &op1 = block.front();
+ Operation &op2 = *(++block.begin());
+
+ if (op1.getNumResults() != 1 ||
+ !op1.hasTrait<::mlir::OpTrait::Elementwise>()) {
+ return rewriter.notifyMatchFailure(op, "must have elementwise trait");
+ }
+
+ // Convert mhlo reduction op into flow reduction op.
+ std::optional<IREE::Flow::CollectiveReductionOp> redOp =
+ convertToFlowCollectiveReductionOp(op1);
+ if (!redOp) {
+ return rewriter.notifyMatchFailure(op, "unsupported operation.");
+ }
+
+ if (!op2.mightHaveTrait<OpTrait::IsTerminator>()) {
+ return rewriter.notifyMatchFailure(op,
+ "the second op must be a terminator");
+ }
+ // Currently only the default channel is used.
+
+ // Create a default channel.
+ auto channel = rewriter.create<IREE::Flow::ChannelDefaultOp>(loc);
+
+ // Convert mhlo reduction op into flow reduction op.
+ auto reductionOpAttr =
+ IREE::Flow::CollectiveReductionOpAttr::get(op.getContext(), *redOp);
+
+ auto inputType = op.getOperand().getType().cast<RankedTensorType>();
+
+ // Get the collective element type attribute.
+ IREE::Flow::CollectiveElementTypeAttr elementTypeAttr =
+ getCollectiveElementTypeAttr(op.getContext(), inputType);
+ if (!elementTypeAttr) {
+ return rewriter.notifyMatchFailure(op, "unsupported input type");
+ }
+
+ // Create an empty tensor for the result.
+ ArrayRef<int64_t> inputShape = inputType.getShape();
+ Value target = rewriter.create<tensor::EmptyOp>(loc, inputShape,
+ inputType.getElementType());
+ auto allReduceOp = rewriter.create<IREE::Flow::CollectiveAllReduceOp>(
+ op.getLoc(), reductionOpAttr, elementTypeAttr, target, op.getOperand(),
+ channel);
+ rewriter.replaceOp(op, allReduceOp.getResult());
+ return success();
+ }
+};
+
+struct ReduceScatterOpConversion
+ : public OpConversionPattern<mhlo::ReduceScatterOp> {
+ using OpConversionPattern<mhlo::ReduceScatterOp>::OpConversionPattern;
+
+ LogicalResult matchAndRewrite(
+ mhlo::ReduceScatterOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+
+ if (!op.getUseGlobalDeviceIds()) {
+ return rewriter.notifyMatchFailure(op, "must use global device IDs");
+ }
+
+ // Check if there is only one group in the replica_groups.
+ ShapedType replicaGroupType = op.getReplicaGroups().getType();
+ if (replicaGroupType.getRank() != 2 ||
+ replicaGroupType.getDimSize(0) != 1) {
+ return rewriter.notifyMatchFailure(op,
+ "must have a single replica group");
+ }
+
+ // Only single elementwise op is supported.
+ Block &block = op.getComputation().front();
+
+ if (block.empty() || llvm::hasSingleElement(block) ||
+ std::next(block.begin(), 2) != block.end()) {
+ return rewriter.notifyMatchFailure(op, "must have two ops in the block");
+ }
+
+ if (block.getNumArguments() != 2) {
+ return rewriter.notifyMatchFailure(op, "must have two block args");
+ }
+
+ Operation &op1 = block.front();
+ Operation &op2 = *(++block.begin());
+
+ if (op1.getNumResults() != 1 ||
+ !op1.hasTrait<::mlir::OpTrait::Elementwise>()) {
+ return rewriter.notifyMatchFailure(op, "must have elementwise trait");
+ }
+
+ // Convert mhlo reduction op into flow reduction op.
+ std::optional<IREE::Flow::CollectiveReductionOp> redOp =
+ convertToFlowCollectiveReductionOp(op1);
+ if (!redOp) {
+ return rewriter.notifyMatchFailure(op, "unsupported operation.");
+ }
+
+ if (!op2.mightHaveTrait<OpTrait::IsTerminator>()) {
+ return rewriter.notifyMatchFailure(op,
+ "the second op must be a terminator");
+ }
+
+ // Convert mhlo reduction op into flow reduction op.
+ auto reductionOpAttr =
+ IREE::Flow::CollectiveReductionOpAttr::get(op.getContext(), *redOp);
+
+ // Currently only the default channel is used.
+
+ // Create a default channel.
+ auto channel = rewriter.create<IREE::Flow::ChannelDefaultOp>(loc);
+
+ // Get the collective element type attribute.
+ auto resultType = op.getResult().getType().cast<RankedTensorType>();
+ IREE::Flow::CollectiveElementTypeAttr elementTypeAttr =
+ getCollectiveElementTypeAttr(op.getContext(), resultType);
+ if (!elementTypeAttr) {
+ return rewriter.notifyMatchFailure(op, "unsupported input type");
+ }
+
+ // When scatter_dimension != 0, we need to transpose between 0 and
+ // scatter_dimension before and after the flow reduce_scatter op.
+ uint64_t scatterDim = op.getScatterDimension();
+ auto inputType = op.getOperand().getType().cast<RankedTensorType>();
+ SmallVector<int64_t> reduceInputShape(inputType.getShape());
+ Value reduceInput = op.getOperand();
+ DenseIntElementsAttr permutationAttr;
+
+ SmallVector<int64_t> scatterResultShape(resultType.getShape());
+ auto elemType = resultType.getElementType();
+
+ if (scatterDim != 0) {
+ SmallVector<int64_t> permutation =
+ llvm::to_vector(llvm::seq<int64_t>(0, scatterResultShape.size()));
+ std::swap(permutation[0], permutation[scatterDim]);
+ permutationAttr = rewriter.getI64VectorAttr(permutation);
+ std::swap(reduceInputShape[0], reduceInputShape[scatterDim]);
+ std::swap(scatterResultShape[0], scatterResultShape[scatterDim]);
+ // Transpose the input.
+ reduceInput =
+ rewriter
+ .create<mhlo::TransposeOp>(
+ loc, RankedTensorType::get(reduceInputShape, elemType),
+ reduceInput, permutationAttr)
+ .getResult();
+ }
+
+ // Create an empty tensor for the result.
+ Value target = rewriter.create<tensor::EmptyOp>(
+ loc, scatterResultShape, resultType.getElementType());
+ Value scatterResult = rewriter
+ .create<IREE::Flow::CollectiveReduceScatterOp>(
+ op.getLoc(), reductionOpAttr, elementTypeAttr,
+ target, reduceInput, channel)
+ .getResult();
+
+ if (scatterDim != 0) {
+ scatterResult = rewriter
+ .create<mhlo::TransposeOp>(
+ loc, resultType, scatterResult, permutationAttr)
+ .getResult();
+ }
+
+ rewriter.replaceOp(op, scatterResult);
+ return success();
+ }
+};
+
+void populateMHLOCollectiveOpsConversionPatterns(MLIRContext *context,
+ TypeConverter &typeConverter,
+ RewritePatternSet &patterns) {
+ patterns.insert<AllGatherOpConversion>(typeConverter, context);
+ patterns.insert<AllReduceOpConversion>(typeConverter, context);
+ patterns.insert<ReduceScatterOpConversion>(typeConverter, context);
+ patterns.insert<ReplicaIdOpConversion>(typeConverter, context);
+}
+
+} // namespace MHLO
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/compiler/src/iree/compiler/InputConversion/MHLO/MHLOToLinalgOnTensors.cpp b/compiler/src/iree/compiler/InputConversion/MHLO/MHLOToLinalgOnTensors.cpp
index 03aa781..ea51f00 100644
--- a/compiler/src/iree/compiler/InputConversion/MHLO/MHLOToLinalgOnTensors.cpp
+++ b/compiler/src/iree/compiler/InputConversion/MHLO/MHLOToLinalgOnTensors.cpp
@@ -356,6 +356,7 @@
// TODO: Collapse/rework all of these patterns once the consolidation
// lands. There is little reason to have these so spread out.
populateMHLOToFlowPatterns(context, patterns);
+
chlo::populateDecomposeChloPatterns(context, &patterns);
populateMHLOBroadcastingToLinalgPatterns(context, *typeConverter, patterns);
mhlo::populateScalarHloToArithmeticConversionPatterns(
@@ -365,6 +366,8 @@
patterns);
populateMHLOComplexToRealPatterns(context, *typeConverter, patterns);
+ populateMHLOCollectiveOpsConversionPatterns(context, *typeConverter,
+ patterns);
// TODO(*): expose patterns that do this much better from
// iree/compiler/Dialect/Util/Transforms/ConvertPrimitiveType.cpp
@@ -386,8 +389,13 @@
context);
patterns.insert<GenericTypeConvert>(
ml_program::GlobalStoreOp::getOperationName(), *typeConverter, context);
-
+ // This is needed when converting mhlo::ReplicaIDOp.
+ patterns.insert<GenericTypeConvert>(
+ tensor::FromElementsOp::getOperationName(), *typeConverter, context);
+ patterns.insert<GenericTypeConvert>(
+ arith::IndexCastUIOp::getOperationName(), *typeConverter, context);
ConversionTarget target(getContext());
+
auto isIllegalType = [&](Type t) { return !typeConverter->isLegal(t); };
auto isLegallyTypedOp = [&](Operation *op) -> bool {
for (Type type : op->getResultTypes()) {
diff --git a/compiler/src/iree/compiler/InputConversion/MHLO/Passes.td b/compiler/src/iree/compiler/InputConversion/MHLO/Passes.td
index 1855004..984b92a 100644
--- a/compiler/src/iree/compiler/InputConversion/MHLO/Passes.td
+++ b/compiler/src/iree/compiler/InputConversion/MHLO/Passes.td
@@ -55,4 +55,5 @@
let constructor = "mlir::iree_compiler::MHLO::createTestMHLOConvertComplexToRealPass()";
}
+
#endif // IREE_COMPILER_INPUTCONVERSION_MHLO_PASSES
diff --git a/compiler/src/iree/compiler/InputConversion/MHLO/Rewriters.h b/compiler/src/iree/compiler/InputConversion/MHLO/Rewriters.h
index ac996f0..e839f38 100644
--- a/compiler/src/iree/compiler/InputConversion/MHLO/Rewriters.h
+++ b/compiler/src/iree/compiler/InputConversion/MHLO/Rewriters.h
@@ -27,6 +27,11 @@
TypeConverter &typeConverter,
RewritePatternSet &patterns);
+/// Populates patterns to convert MHLO collective ops to Stream ops.
+void populateMHLOCollectiveOpsConversionPatterns(MLIRContext *context,
+ TypeConverter &typeConverter,
+ RewritePatternSet &patterns);
+
/// Populates patterns to convert MHLO/CHLO arithmetic on complex tensors to
/// equivalent HLO level real arithmetic.
void populateMHLOComplexToRealPatterns(MLIRContext *context,
diff --git a/compiler/src/iree/compiler/InputConversion/MHLO/test/BUILD b/compiler/src/iree/compiler/InputConversion/MHLO/test/BUILD
index 99e5b4f..5ac112d 100644
--- a/compiler/src/iree/compiler/InputConversion/MHLO/test/BUILD
+++ b/compiler/src/iree/compiler/InputConversion/MHLO/test/BUILD
@@ -20,6 +20,7 @@
[
"broadcasting.mlir",
"convert_mhlo_to_linalg_ext.mlir",
+ "convert_collective_ops.mlir",
"convert_complex_to_real.mlir",
"convert_structural_types.mlir",
"dynamic_shape.mlir",
diff --git a/compiler/src/iree/compiler/InputConversion/MHLO/test/CMakeLists.txt b/compiler/src/iree/compiler/InputConversion/MHLO/test/CMakeLists.txt
index 8404a20..8b9aa57 100644
--- a/compiler/src/iree/compiler/InputConversion/MHLO/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/InputConversion/MHLO/test/CMakeLists.txt
@@ -15,6 +15,7 @@
lit
SRCS
"broadcasting.mlir"
+ "convert_collective_ops.mlir"
"convert_complex_to_real.mlir"
"convert_mhlo_to_linalg_ext.mlir"
"convert_structural_types.mlir"
diff --git a/compiler/src/iree/compiler/InputConversion/MHLO/test/convert_collective_ops.mlir b/compiler/src/iree/compiler/InputConversion/MHLO/test/convert_collective_ops.mlir
new file mode 100644
index 0000000..a93557d
--- /dev/null
+++ b/compiler/src/iree/compiler/InputConversion/MHLO/test/convert_collective_ops.mlir
@@ -0,0 +1,167 @@
+// RUN: iree-opt --split-input-file --iree-mhlo-to-linalg-on-tensors --canonicalize -cse %s | FileCheck %s
+
+// CHECK-LABEL: @replica_id
+func.func @replica_id() -> tensor<ui32> {
+ // CHECK-DAG: [[CHANNEL:%.+]] = flow.channel.default : !flow.channel
+ // CHECK-DAG: [[RANK:%.+]] = flow.channel.rank [[CHANNEL]] : index
+ // CHECK-DAG: [[CAST:%.+]] = arith.index_castui [[RANK]] : index to i32
+ // CHECK-DAG: [[TENSOR:%.+]] = tensor.from_elements [[CAST]] : tensor<i32>
+ // CHECK-DAG: return [[TENSOR]] : tensor<i32>
+ %id = mhlo.replica_id : tensor<ui32>
+ return %id : tensor<ui32>
+}
+
+// -----
+
+// CHECK-LABEL: @all_reduce_sum
+// CHECK-SAME: ([[ARG0:%.+]]: tensor<2304xf32>)
+func.func @all_reduce_sum(%input : tensor<2304xf32>) -> tensor<2304xf32> {
+ // CHECK: [[CHANNEL:%.+]] = flow.channel.default : !flow.channel
+ // CHECK: [[EMPTY:%.+]] = tensor.empty() : tensor<2304xf32>
+ // CHECK: [[ALLREDUCE:%.+]] = flow.collective.all_reduce sum, f32, [[EMPTY]], [[ARG0]], %channel_default : (tensor<2304xf32>, tensor<2304xf32>, !flow.channel) -> [[EMPTY]] as tensor<2304xf32>
+ // CHECK: return [[ALLREDUCE]] : tensor<2304xf32>
+ %out = "mhlo.all_reduce"(%input) ({
+ ^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
+ %sum = mhlo.add %arg0, %arg1 : tensor<f32>
+ mhlo.return %sum : tensor<f32>
+ }) {channel_handle = #mhlo.channel_handle<handle = 1, type = 1>,
+ replica_groups = dense<[[0, 1, 2, 3, 4, 5, 6, 7]]> : tensor<1x8xi64>,
+ use_global_device_ids} : (tensor<2304xf32>) -> tensor<2304xf32>
+ return %out : tensor<2304xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @all_reduce_product
+// CHECK-SAME: ([[ARG0:%.+]]: tensor<2304xf32>)
+func.func @all_reduce_product(%input : tensor<2304xf32>) -> tensor<2304xf32> {
+ // CHECK: [[CHANNEL:%.+]] = flow.channel.default : !flow.channel
+ // CHECK: [[EMPTY:%.+]] = tensor.empty() : tensor<2304xf32>
+ // CHECK: [[OP:%.+]] = flow.collective.all_reduce product, f32, [[EMPTY]], [[ARG0]], %channel_default : (tensor<2304xf32>, tensor<2304xf32>, !flow.channel) -> [[EMPTY]] as tensor<2304xf32>
+ // CHECK: return [[OP]] : tensor<2304xf32>
+ %out = "mhlo.all_reduce"(%input) ({
+ ^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
+ %mul = mhlo.multiply %arg0, %arg1 : tensor<f32>
+ mhlo.return %mul : tensor<f32>
+ }) {channel_handle = #mhlo.channel_handle<handle = 1, type = 1>,
+ replica_groups = dense<[[0, 1, 2, 3, 4, 5, 6, 7]]> : tensor<1x8xi64>,
+ use_global_device_ids} : (tensor<2304xf32>) -> tensor<2304xf32>
+ return %out : tensor<2304xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @all_reduce_minimum
+// CHECK-SAME: ([[ARG0:%.+]]: tensor<2304xf32>)
+func.func @all_reduce_minimum(%input : tensor<2304xf32>) -> tensor<2304xf32> {
+ // CHECK: [[CHANNEL:%.+]] = flow.channel.default : !flow.channel
+ // CHECK: [[EMPTY:%.+]] = tensor.empty() : tensor<2304xf32>
+ // CHECK: [[OP:%.+]] = flow.collective.all_reduce minimum, f32, [[EMPTY]], [[ARG0]], %channel_default : (tensor<2304xf32>, tensor<2304xf32>, !flow.channel) -> [[EMPTY]] as tensor<2304xf32>
+ // CHECK: return [[OP]] : tensor<2304xf32>
+ %out = "mhlo.all_reduce"(%input) ({
+ ^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
+ %mul = mhlo.minimum %arg0, %arg1 : tensor<f32>
+ mhlo.return %mul : tensor<f32>
+ }) {channel_handle = #mhlo.channel_handle<handle = 1, type = 1>,
+ replica_groups = dense<[[0, 1, 2, 3, 4, 5, 6, 7]]> : tensor<1x8xi64>,
+ use_global_device_ids} : (tensor<2304xf32>) -> tensor<2304xf32>
+ return %out : tensor<2304xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @all_reduce_maximum
+// CHECK-SAME: ([[ARG0:%.+]]: tensor<2304xf32>)
+func.func @all_reduce_maximum(%input : tensor<2304xf32>) -> tensor<2304xf32> {
+ // CHECK: [[CHANNEL:%.+]] = flow.channel.default : !flow.channel
+ // CHECK: [[EMPTY:%.+]] = tensor.empty() : tensor<2304xf32>
+ // CHECK: [[OP:%.+]] = flow.collective.all_reduce maximum, f32, [[EMPTY]], [[ARG0]], %channel_default : (tensor<2304xf32>, tensor<2304xf32>, !flow.channel) -> [[EMPTY]] as tensor<2304xf32>
+ // CHECK: return [[OP]] : tensor<2304xf32>
+ %out = "mhlo.all_reduce"(%input) ({
+ ^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
+ %mul = mhlo.maximum %arg0, %arg1 : tensor<f32>
+ mhlo.return %mul : tensor<f32>
+ }) {channel_handle = #mhlo.channel_handle<handle = 1, type = 1>,
+ replica_groups = dense<[[0, 1, 2, 3, 4, 5, 6, 7]]> : tensor<1x8xi64>,
+ use_global_device_ids} : (tensor<2304xf32>) -> tensor<2304xf32>
+ return %out : tensor<2304xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @all_gather_dim_0
+// CHECK-SAME: ([[ARG0:%.+]]: tensor<512xf32>) -> tensor<1024xf32>
+func.func @all_gather_dim_0(%input : tensor<512xf32>) -> tensor<1024xf32> {
+ // CHECK: [[CHANNEL:%.+]] = flow.channel.default : !flow.channel
+ // CHECK: [[EMPTY:%.+]] = tensor.empty() : tensor<1024xf32>
+ // CHECK: [[OP:%.+]] = flow.collective.all_gather f32, [[EMPTY]], [[ARG0]], %channel_default : (tensor<1024xf32>, tensor<512xf32>, !flow.channel) -> [[EMPTY]] as tensor<1024xf32>
+ // CHECK: return [[OP]] : tensor<1024xf32>
+ %out = "mhlo.all_gather"(%input) {all_gather_dim = 0 : i64,
+ channel_handle = #mhlo.channel_handle<handle = 1, type = 1>,
+ replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
+ use_global_device_ids} : (tensor<512xf32>) -> tensor<1024xf32>
+ return %out : tensor<1024xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @all_gather_dim_1
+// CHECK-SAME: ([[ARG0:%.+]]: tensor<2x2xf32>) -> tensor<2x4xf32>
+func.func @all_gather_dim_1(%input : tensor<2x2xf32>) -> tensor<2x4xf32> {
+ // CHECK: [[CHANNEL:%.+]] = flow.channel.default : !flow.channel
+ // CHECK: tensor.empty() : tensor<2x2xf32>
+ // CHECK: [[TRANSPOSE_ARG:%.+]] = linalg.generic
+ // CHECK: [[EMPTY:%.+]] = tensor.empty() : tensor<4x2xf32>
+ // CHECK: [[OP:%.+]] = flow.collective.all_gather f32, [[EMPTY]], [[TRANSPOSE_ARG]], %channel_default : (tensor<4x2xf32>, tensor<2x2xf32>, !flow.channel) -> [[EMPTY]] as tensor<4x2xf32>
+ // CHECK: tensor.empty() : tensor<2x4xf32>
+ // CHECK: [[TRANSPOSE_OUT:%.+]] = linalg.generic
+ // CHECK: return [[TRANSPOSE_OUT]] : tensor<2x4xf32>
+ %out = "mhlo.all_gather"(%input) {all_gather_dim = 1 : i64,
+ channel_handle = #mhlo.channel_handle<handle = 1, type = 1>,
+ replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
+ use_global_device_ids} : (tensor<2x2xf32>) -> tensor<2x4xf32>
+ return %out : tensor<2x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @reduce_scatter_dim_0
+// CHECK-SAME: ([[ARG0:%.+]]: tensor<4x2xf32>) -> tensor<2x2xf32>
+func.func @reduce_scatter_dim_0(%input : tensor<4x2xf32>) -> tensor<2x2xf32> {
+ // CHECK: [[CHANNEL:%.+]] = flow.channel.default : !flow.channel
+ // CHECK: [[EMPTY:%.+]] = tensor.empty() : tensor<2x2xf32>
+ // CHECK: [[OP:%.+]] = flow.collective.reduce_scatter sum, f32, [[EMPTY]], [[ARG0]], %channel_default : (tensor<2x2xf32>, tensor<4x2xf32>, !flow.channel) -> [[EMPTY]] as tensor<2x2xf32>
+ // CHECK: return [[OP]] : tensor<2x2xf32>
+ %out = "mhlo.reduce_scatter"(%input) ({
+ ^bb0(%arg0: tensor<f32> , %arg1: tensor<f32>) :
+ %sum = mhlo.add %arg0, %arg1 : tensor<f32>
+ mhlo.return %sum : tensor<f32>
+ }) {scatter_dimension = 0 : i64,
+ channel_handle = #mhlo.channel_handle<handle = 1, type = 1>,
+ replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
+ use_global_device_ids} : (tensor<4x2xf32>) -> tensor<2x2xf32>
+ return %out : tensor<2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @reduce_scatter_dim_1
+// CHECK-SAME: ([[ARG0:%.+]]: tensor<2x4xf32>) -> tensor<2x2xf32>
+func.func @reduce_scatter_dim_1(%input : tensor<2x4xf32>) -> tensor<2x2xf32> {
+ // CHECK: [[CHANNEL:%.+]] = flow.channel.default : !flow.channel
+ // CHECK: tensor.empty() : tensor<4x2xf32>
+ // CHECK: [[TRANSPOSE_ARG:%.+]] = linalg.generic
+ // CHECK: [[EMPTY:%.+]] = tensor.empty() : tensor<2x2xf32>
+ // CHECK: [[OP:%.+]] = flow.collective.reduce_scatter sum, f32, [[EMPTY]], [[TRANSPOSE_ARG]], %channel_default : (tensor<2x2xf32>, tensor<4x2xf32>, !flow.channel) -> [[EMPTY]] as tensor<2x2xf32>
+ // CHECK: [[TRANSPOSE_OUT:%.+]] = linalg.generic
+ // CHECK: return [[TRANSPOSE_OUT]] : tensor<2x2xf32>
+ %out = "mhlo.reduce_scatter"(%input) ({
+ ^bb0(%arg0: tensor<f32> , %arg1: tensor<f32>) :
+ %sum = mhlo.add %arg0, %arg1 : tensor<f32>
+ mhlo.return %sum : tensor<f32>
+ }) {scatter_dimension = 1 : i64,
+ channel_handle = #mhlo.channel_handle<handle = 1, type = 1>,
+ replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
+ use_global_device_ids} : (tensor<2x4xf32>) -> tensor<2x2xf32>
+ return %out : tensor<2x2xf32>
+}
diff --git a/runtime/src/iree/hal/command_buffer.h b/runtime/src/iree/hal/command_buffer.h
index c88700e..6e5fd65 100644
--- a/runtime/src/iree/hal/command_buffer.h
+++ b/runtime/src/iree/hal/command_buffer.h
@@ -296,8 +296,10 @@
// Specifies the reduction operator of a collective reduction operation.
enum iree_hal_collective_reduction_e {
+ // Specifies that the reduction operation is unspecified.
+ IREE_HAL_COLLECTIVE_REDUCTION_NONE = 0,
// Specifies that the reduction operation computes a sum (addition).
- IREE_HAL_COLLECTIVE_REDUCTION_SUM = 0,
+ IREE_HAL_COLLECTIVE_REDUCTION_SUM = 1,
// Specifies that the reduction operation computes a product (multiplication).
IREE_HAL_COLLECTIVE_REDUCTION_PRODUCT,
// Specifies that the reduction operation computes a minimum (min).
diff --git a/runtime/src/iree/hal/drivers/cuda/cuda_device.c b/runtime/src/iree/hal/drivers/cuda/cuda_device.c
index ddb3c8b..ef0a82c 100644
--- a/runtime/src/iree/hal/drivers/cuda/cuda_device.c
+++ b/runtime/src/iree/hal/drivers/cuda/cuda_device.c
@@ -302,7 +302,8 @@
// We could multiplex channels but it'd be better to surface that to the
// compiler so that it can emit the right rank math.
int requested_count = iree_math_count_ones_u64(queue_affinity);
- if (requested_count != 1) {
+ // TODO(#12206): properly assign affinity in the compiler.
+ if (requested_count != 64 && requested_count != 1) {
return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
"exactly one participant is allowed in a "
"channel but %d were specified",
diff --git a/tests/e2e/models/generated_e2e_model_tests.cmake b/tests/e2e/models/generated_e2e_model_tests.cmake
index ea0a24f..f33f379 100644
--- a/tests/e2e/models/generated_e2e_model_tests.cmake
+++ b/tests/e2e/models/generated_e2e_model_tests.cmake
@@ -32,6 +32,7 @@
RUNNER_ARGS
"--function=main"
"--input=1x224x224x3xf32=0"
+ "--device_allocator=caching"
UNSUPPORTED_PLATFORMS
"riscv32-Linux"
"android-arm64-v8a"
@@ -49,6 +50,7 @@
RUNNER_ARGS
"--function=main"
"--input=1x224x224x3xui8=0"
+ "--device_allocator=caching"
UNSUPPORTED_PLATFORMS
"android-arm64-v8a"
)
@@ -66,6 +68,7 @@
"--function=main"
"--input=1x257x257x3xf32=0"
"--expected_f32_threshold=0.001"
+ "--device_allocator=caching"
UNSUPPORTED_PLATFORMS
"riscv32-Linux"
)
@@ -82,6 +85,7 @@
RUNNER_ARGS
"--function=main"
"--input=1x96x96x1xi8=0"
+ "--device_allocator=caching"
UNSUPPORTED_PLATFORMS
"android-arm64-v8a"
)
diff --git a/tests/transform_dialect/cuda/BUILD b/tests/transform_dialect/cuda/BUILD
index f818876..022af65 100644
--- a/tests/transform_dialect/cuda/BUILD
+++ b/tests/transform_dialect/cuda/BUILD
@@ -26,6 +26,7 @@
iree_lit_test_suite(
name = "lit",
srcs = [
+ "mma.mlir",
"reduction.mlir",
"reduction_eltwise.mlir",
"reduction_v2.mlir",
diff --git a/tests/transform_dialect/cuda/CMakeLists.txt b/tests/transform_dialect/cuda/CMakeLists.txt
index 4bd5475..9c34317 100644
--- a/tests/transform_dialect/cuda/CMakeLists.txt
+++ b/tests/transform_dialect/cuda/CMakeLists.txt
@@ -18,6 +18,7 @@
NAME
lit
SRCS
+ "mma.mlir"
"reduction.mlir"
"reduction_eltwise.mlir"
"reduction_v2.mlir"
diff --git a/tests/transform_dialect/cuda/mma.mlir b/tests/transform_dialect/cuda/mma.mlir
new file mode 100644
index 0000000..d7cdeec
--- /dev/null
+++ b/tests/transform_dialect/cuda/mma.mlir
@@ -0,0 +1,71 @@
+// RUN: iree-opt %s --split-input-file --iree-transform-dialect-interpreter | FileCheck %s
+
+
+#matmat_accesses = [
+ affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (k, n)>,
+ affine_map<(m, n, k) -> (m, n)>
+]
+
+#matmat_trait = {
+ indexing_maps = #matmat_accesses,
+ iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+func.func @wmma(%a: memref<16x16xf32>, %b: memref<16x16xf32>, %c: memref<16x16xf32>) {
+ %c0 = arith.constant 0: index
+ %cst = arith.constant 0.0: f32
+ %va = vector.transfer_read %a[%c0, %c0], %cst: memref<16x16xf32>, vector<16x16xf32>
+ %vb = vector.transfer_read %b[%c0, %c0], %cst: memref<16x16xf32>, vector<16x16xf32>
+ %vc = vector.transfer_read %c[%c0, %c0], %cst: memref<16x16xf32>, vector<16x16xf32>
+
+ // CHECK-NOT: vector.contract
+ // CHECK: gpu.subgroup_mma_compute
+ %vres = vector.contract #matmat_trait %va, %vb, %vc
+ : vector<16x16xf32>, vector<16x16xf32> into vector<16x16xf32>
+ vector.transfer_write %vres, %c[%c0, %c0]: vector<16x16xf32>, memref<16x16xf32>
+ return
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%module: !pdl.operation):
+ %func = transform.structured.match ops{["func.func"]} in %module
+ : (!pdl.operation) -> !pdl.operation
+ %func_2 = transform.iree.apply_patterns %func { unroll_vectors_gpu_wmma }
+ transform.iree.vector.vector_to_mma_conversion %func_2 { use_wmma }
+}
+
+// -----
+
+#matmat_accesses = [
+ affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (k, n)>,
+ affine_map<(m, n, k) -> (m, n)>
+]
+#matmat_trait = {
+ indexing_maps = #matmat_accesses,
+ iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+func.func @mma_sync(%a: memref<16x16xf32>, %b: memref<16x16xf32>, %c: memref<16x16xf32>) {
+ %c0 = arith.constant 0: index
+ %cst = arith.constant 0.0: f32
+ %va = vector.transfer_read %a[%c0, %c0], %cst: memref<16x16xf32>, vector<16x16xf32>
+ %vb = vector.transfer_read %b[%c0, %c0], %cst: memref<16x16xf32>, vector<16x16xf32>
+ %vc = vector.transfer_read %c[%c0, %c0], %cst: memref<16x16xf32>, vector<16x16xf32>
+
+ // CHECK-NOT: vector.contract
+ // CHECK: nvgpu.mma.sync
+ %vres = vector.contract #matmat_trait %va, %vb, %vc
+ : vector<16x16xf32>, vector<16x16xf32> into vector<16x16xf32>
+ vector.transfer_write %vres, %c[%c0, %c0]: vector<16x16xf32>, memref<16x16xf32>
+ return
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%module: !pdl.operation):
+ %func = transform.structured.match ops{["func.func"]} in %module
+ : (!pdl.operation) -> !pdl.operation
+ %func_2 = transform.iree.apply_patterns %func { unroll_vectors_gpu_mma_sync }
+ transform.iree.vector.vector_to_mma_conversion %func_2 { use_mma_sync }
+}
diff --git a/third_party/llvm-project b/third_party/llvm-project
index 66616cb..677ea5e 160000
--- a/third_party/llvm-project
+++ b/third_party/llvm-project
@@ -1 +1 @@
-Subproject commit 66616cb9136ed360b0b8b3545fbd64050b32c8f6
+Subproject commit 677ea5eb2e5b25b137495221a21877504ee5f2ce
