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