diff --git a/.github/workflows/build_package.yml b/.github/workflows/build_package.yml
index 44103a9..79b7e9d 100644
--- a/.github/workflows/build_package.yml
+++ b/.github/workflows/build_package.yml
@@ -263,4 +263,4 @@
           workflow: Validate and Publish Release
           token: ${{ secrets.WRITE_ACCESS_TOKEN }}
           ref: "${{ env.tag_name }}"
-          inputs: '{"release_id": "${{ github.event.inputs.release_id }}"}'
+          inputs: '{"release_id": "${{ github.event.inputs.release_id }}", "package_version": "${{ github.event.inputs.package_version }}", "build_run_id": "${{ github.run_id }}"}'
diff --git a/.github/workflows/validate_and_publish_release.yml b/.github/workflows/validate_and_publish_release.yml
index 2124c8d..a06b660 100644
--- a/.github/workflows/validate_and_publish_release.yml
+++ b/.github/workflows/validate_and_publish_release.yml
@@ -1,5 +1,5 @@
-# Publishes a draft snapshot, after performing validation to make sure all
-# packages are stable (coming soon).
+# Publishes a draft snapshot, after performing validation to make sure all linux
+# packages are stable.
 name: Validate and Publish Release
 
 on:
@@ -8,13 +8,71 @@
       release_id:
         description: 'Release id to publish'
         required: true
-        default: ''
-
-# TODO(jennik): Add more validation. https://github.com/google/iree/issues/6316
+      package_version:
+        description: 'Version of the package'
+        required: true
+      build_run_id:
+        description: 'Run ID for the build_package.yml workflow that triggered this workflow'
+        required: true
 
 jobs:
+  validate_packages:
+    name: "Validate packages"
+    # TODO(jennik): Look into testing windows and macos builds.
+    runs-on: ubuntu-18.04
+    steps:
+      - name: Download packages
+        id: download_packages
+        uses: dawidd6/action-download-artifact@v2
+        with:
+          github_token: ${{secrets.WRITE_ACCESS_TOKEN}}
+          workflow: build_package.yml
+          run_id: ${{ github.event.inputs.build_run_id }}
+      - name: Extract and display downloaded files
+        run: |
+          tar -xf artifact/iree-dist-snapshot-${{ github.event.inputs.package_version }}-linux-x86_64.tar.xz
+          pwd
+          ls -R
+      - name: Set up python
+        id: set_up_python
+        uses: actions/setup-python@v2
+        with:
+          python-version: '3.x'
+      - name: Install python packages
+        id: install_python_packages
+        run: |
+          python -m pip install -f file://$PWD/artifact/ iree-compiler-snapshot iree-jax-snapshot iree-runtime-snapshot iree-tools-tflite-snapshot iree-tools-tf-snapshot iree-tools-xla-snapshot
+      - name: Run iree-benchmark-module
+        id: run_iree_benchmark_module
+        run: ./bin/iree-benchmark-module --help
+      - name: Run iree-benchmark-trace
+        id: run_iree_benchmark_trace
+        run: ./bin/iree-benchmark-trace --help
+      - name: Run iree-flatcc-cli
+        id: run_iree_flatcc_cli
+        run: ./bin/iree-flatcc-cli --help
+      - name: Run iree-opt
+        id: run_iree_opt
+        run: ./bin/iree-opt --help
+      - name: Run iree-run-mlir
+        id: run_iree_run_mlir
+        run: ./bin/iree-run-mlir --help
+      - name: Run iree-run-module
+        id: run_iree_run_module
+        run: ./bin/iree-run-module --help
+      - name: Run iree-run-trace
+        id: run_iree_run_trace
+        run: ./bin/iree-run-trace --help
+      - name: Run iree-tblgen
+        id: run_iree_tblgen
+        run: ./bin/iree-tblgen --help
+      - name: Run iree-translate
+        id: run_iree_translate
+        run: ./bin/iree-translate --help
+
   publish_release:
     name: "Publish release"
+    needs: validate_packages
     runs-on: ubuntu-18.04
     steps:
       - name: Publish Release
diff --git a/bindings/python/iree/compiler/core.py b/bindings/python/iree/compiler/core.py
index a74466a..d9baedc 100644
--- a/bindings/python/iree/compiler/core.py
+++ b/bindings/python/iree/compiler/core.py
@@ -172,9 +172,6 @@
   cl = [
       iree_translate,
       input_file,
-      # People are only looking at stderr here if the tool fails and there's
-      # already lots of output, so adding this doesn't cost much.
-      "--print-main-address",
       f"--iree-input-type={options.input_type.value}",
       f"--iree-vm-bytecode-module-output-format={options.output_format.value}",
   ]
diff --git a/docs/developers/debugging/releases.md b/docs/developers/debugging/releases.md
index efd7ba9..6f82ac3 100644
--- a/docs/developers/debugging/releases.md
+++ b/docs/developers/debugging/releases.md
@@ -84,3 +84,50 @@
   Release mode `main-dist` package above once, which will drop binaries in the
   `iree-install` directory. Then build the `py-runtime-pkg` or equiv and
   iterate further in the build directory. Ditto for TF/XLA/etc.
+
+## Testing releases on your fork
+
+To avoid interrupting the regular releases published on the IREE github, you
+can test any changes to the release process on your own fork.  Some setup is
+required before these github actions will work on your fork and development
+branch.
+
+To run
+[`schedule_snapshot_release.yml`](https://github.com/google/iree/blob/main/.github/workflows/schedule_snapshot_release.yml),
+comment out
+[this line](https://github.com/google/iree/blob/392449e986493bf710e3da637ebf807715da9ffe/.github/workflows/schedule_snapshot_release.yml#L14):
+```yaml
+# Don't run this in everyone's forks.
+if: github.repository == 'google/iree'
+```
+
+And change the branch from 'main' to the branch you are developing on
+[here](https://github.com/google/iree/blob/392449e986493bf710e3da637ebf807715da9ffe/.github/workflows/schedule_snapshot_release.yml#L37):
+```yaml
+- name: Pushing changes
+  uses: ad-m/github-push-action@v0.6.0
+  with:
+    github_token: ${{ secrets.WRITE_ACCESS_TOKEN }}
+    branch: main
+    tags: true
+```
+
+To speed up
+[`build_package.yml`](https://github.com/google/iree/blob/main/.github/workflows/build_package.yml),
+you may want to comment out some of the builds
+[here](https://github.com/google/iree/blob/392449e986493bf710e3da637ebf807715da9ffe/.github/workflows/build_package.yml#L34-L87).
+The
+[`py-pure-pkgs`](https://github.com/google/iree/blob/392449e986493bf710e3da637ebf807715da9ffe/.github/workflows/build_package.yml#L52)
+build takes only ~2 minutes and the
+[`py-runtime-pkg`](https://github.com/google/iree/blob/392449e986493bf710e3da637ebf807715da9ffe/.github/workflows/build_package.yml#L39)
+build takes ~5, while the others can take several hours.
+
+From your development branch, you can manually run the
+[Schedule Snapshot Release](https://github.com/google/iree/actions/workflows/schedule_snapshot_release.yml)
+action, which invokes the
+[Build Native Release Packages](https://github.com/google/iree/actions/workflows/build_package.yml)
+action, which finally invokes the
+[Validate and Publish Release](https://github.com/google/iree/actions/workflows/validate_and_publish_release.yml)
+action.  If you already have a draft release and know the release id, package
+version, and run ID from a previous Build Native Release Packages run, you can
+also manually run just the Validate and Publish Release action.
diff --git a/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp b/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp
index c12fac9..27deff8 100644
--- a/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp
+++ b/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp
@@ -176,6 +176,15 @@
       IREE::HAL::InterfaceBindingOp binding = input.queryBindingOp();
       llvmInputTypes[argMapping[binding]] = llvmType;
     });
+    // As a convention with HAL, push constants are appended as kernel arguments
+    // after all the binding inputs.
+    uint64_t numConstants = 0;
+    funcOp.walk([&](IREE::HAL::InterfaceLoadConstantOp constant) {
+      numConstants =
+          std::max(constant.offset().getZExtValue() + 1, numConstants);
+    });
+    llvmInputTypes.resize(argMapping.size() + numConstants,
+                          rewriter.getI32Type());
     signatureConverter.addInputs(llvmInputTypes);
 
     // Construct newFunc with all attributes except return type & symbol name.
@@ -256,14 +265,72 @@
           rewriter, loc, *getTypeConverter(), memrefType, llvmBufferBasePtr);
       rewriter.replaceOp(op, {desc});
     } else {
-      // TODO: pull those paramters from HAL constants.
-      assert(0 && "TODO: implement dynamic shape");
+      ValueRange dynamicDims = adaptor.dynamic_dims();
+      assert(memrefType.getNumDynamicDims() == dynamicDims.size());
+      int64_t rank = memrefType.getRank();
+
+      // Build MemRef descriptor for this interface binding.
+      auto desc = MemRefDescriptor::undef(
+          rewriter, loc, typeConverter->convertType(memrefType));
+      desc.setAllocatedPtr(rewriter, loc, llvmBufferBasePtr);
+      desc.setAlignedPtr(rewriter, loc, llvmBufferBasePtr);
+      desc.setConstantOffset(rewriter, loc, 0);
+
+      // Update memref descriptor shape. Dynamic dimensions can be mixed with
+      // static dimensions, like [128, ?, 128].
+      int dynamicDimIndex = 0;
+      for (int i = 0; i < rank; ++i) {
+        if (memrefType.isDynamicDim(i)) {
+          desc.setSize(rewriter, loc, i, dynamicDims[dynamicDimIndex++]);
+        } else {
+          desc.setConstantSize(rewriter, loc, i, memrefType.getDimSize(i));
+        }
+      }
+
+      // Compute and update strides. Assume that MemRefs are row-major, that is,
+      // following index linearization:
+      //   x[i, j, k] = i * x.dim[1] * x.dim[2] + j * x.dim[2] + k
+      desc.setConstantStride(rewriter, loc, rank - 1, 1);
+      for (int i = rank - 2; i >= 0; --i) {
+        auto stride = desc.stride(rewriter, loc, i + 1);
+        auto dim = desc.size(rewriter, loc, i + 1);
+        Value strideVal = rewriter.create<LLVM::MulOp>(loc, stride, dim);
+        desc.setStride(rewriter, loc, i, strideVal);
+      }
+      rewriter.replaceOp(op, {desc});
     }
 
     return success();
   }
 };
 
+class ConvertIREEConstantOp : public ConvertToLLVMPattern {
+ public:
+  explicit ConvertIREEConstantOp(MLIRContext *context,
+                                 LLVMTypeConverter &converter)
+      : ConvertToLLVMPattern(
+            IREE::HAL::InterfaceLoadConstantOp::getOperationName(), context,
+            converter) {}
+  LogicalResult matchAndRewrite(
+      Operation *op, ArrayRef<Value> operands,
+      ConversionPatternRewriter &rewriter) const override {
+    // Bail until nested under an LLVMFuncOp.
+    auto llvmFuncOp = op->getParentOfType<LLVM::LLVMFuncOp>();
+    if (!llvmFuncOp) return failure();
+    assert(llvmFuncOp.getNumArguments() > 0);
+
+    llvm::SmallDenseMap<Operation *, size_t> argMapping =
+        getKernelArgMapping(llvmFuncOp);
+    auto ireeConstantOp = cast<IREE::HAL::InterfaceLoadConstantOp>(op);
+    mlir::BlockArgument llvmBufferArg = llvmFuncOp.getArgument(
+        argMapping.size() + ireeConstantOp.offset().getZExtValue());
+    assert(llvmBufferArg.getType().isInteger(32));
+    Type dstType = getTypeConverter()->convertType(ireeConstantOp.getType());
+    rewriter.replaceOpWithNewOp<LLVM::ZExtOp>(op, dstType, llvmBufferArg);
+    return success();
+  }
+};
+
 /// A pattern to convert hal.interface.workgroup.id/count/size into
 /// corresponding NVVM/ROCDL ops.
 template <typename InterfaceOpTy, typename XOp, typename YOp, typename ZOp>
@@ -305,7 +372,8 @@
                                     OwningRewritePatternList &patterns,
                                     LLVMTypeConverter &converter,
                                     bool useROCM) {
-  patterns.insert<ConvertFunc, ConvertIREEBindingOp>(context, converter);
+  patterns.insert<ConvertFunc, ConvertIREEBindingOp, ConvertIREEConstantOp>(
+      context, converter);
   if (useROCM) {
     patterns.insert<HALInterfaceWorkgroupOpsConverter<
                         IREE::HAL::InterfaceWorkgroupIDOp, ROCDL::BlockIdXOp,
diff --git a/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp b/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
index a069bed..bcb7bf5 100644
--- a/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
+++ b/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
@@ -43,8 +43,8 @@
   // Infer the MxN size of the matmul based on operands and indexing maps.
   auto lhsShape = getUntiledShape(op.getInputOperand(0)->get());
   auto rhsShape = getUntiledShape(op.getInputOperand(1)->get());
-  int64_t sizeM = -1;
-  int64_t sizeN = -1;
+  int64_t sizeM = ShapedType::kDynamicSize;
+  int64_t sizeN = ShapedType::kDynamicSize;
   auto outputMap = op.getTiedIndexingMap(op.getOutputOperand(0));
   for (unsigned i = 0; i < lhsShape.size(); i++) {
     if (op.getTiedIndexingMap(op.getInputOperand(0)).getDimPosition(i) ==
@@ -65,25 +65,29 @@
   int64_t tileY = 256;
   int64_t tileK = 4;
   SmallVector<int64_t, 3> workgroupSize = {2 * cudaWarpSize, 1, 1};
-  // Special case for very small matrices.
-  if (sizeM * sizeN <= cudaWarpSize) {
-    tileX = sizeN;
-    tileY = sizeM;
-    workgroupSize = {sizeM, sizeN, 1};
-  }
-  SmallVector<TileWorkgroupSizePair> tileSizeConfig;
-  // Query the best configuration.
-  getMatmulConfig(tileSizeConfig);
-  // Pick the best configuration where the original shape is aligned on the tile
-  // size.
-  for (TileWorkgroupSizePair &config : tileSizeConfig) {
-    if (sizeN % config.tileSize[1] == 0 && sizeM % config.tileSize[0] == 0) {
-      tileX = config.tileSize[0];
-      tileY = config.tileSize[1];
-      tileK = config.tileSize[2];
-      workgroupSize.assign(config.workgroupSize.begin(),
-                           config.workgroupSize.end());
-      break;
+  bool isStaticSize =
+      sizeM != ShapedType::kDynamicSize && sizeN != ShapedType::kDynamicSize;
+  if (isStaticSize) {
+    // Special case for very small matrices.
+    if (sizeM * sizeN <= cudaWarpSize) {
+      tileX = sizeN;
+      tileY = sizeM;
+      workgroupSize = {sizeM, sizeN, 1};
+    }
+    SmallVector<TileWorkgroupSizePair> tileSizeConfig;
+    // Query the best configuration.
+    getMatmulConfig(tileSizeConfig);
+    // Pick the best configuration where the original shape is aligned on the
+    // tile size.
+    for (TileWorkgroupSizePair &config : tileSizeConfig) {
+      if (sizeN % config.tileSize[1] == 0 && sizeM % config.tileSize[0] == 0) {
+        tileX = config.tileSize[0];
+        tileY = config.tileSize[1];
+        tileK = config.tileSize[2];
+        workgroupSize.assign(config.workgroupSize.begin(),
+                             config.workgroupSize.end());
+        break;
+      }
     }
   }
   // Currently just a basic tile size to enable tiling and vectorization.
diff --git a/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir b/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir
index 00e5c31..7ee0fcc 100644
--- a/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir
+++ b/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt -iree-convert-to-nvvm %s | IreeFileCheck %s
+// RUN: iree-opt -iree-convert-to-nvvm -split-input-file %s | IreeFileCheck %s
 
 // Test that that standard and GPU ops are converted to LLVM and NVVM.
 func @abs_ex_dispatch_0() {
@@ -37,3 +37,44 @@
 //       CHECK:   llvm.insertvalue %[[PTR]], %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<1 x i64>, array<1 x i64>)>
 //      CHECK:    nvvm.read.ptx.sreg.tid.x
 //      CHECK:    llvm.fadd
+
+// -----
+
+func @abs_dynamic() {
+  %c0 = constant 0 : index
+  %c128 = constant 128 : index
+  %s = hal.interface.load.constant offset = 1 : index
+  %0 = hal.interface.binding.subspan @io::@arg0[%c128] : memref<?xf32>{%s}
+  %1 = hal.interface.binding.subspan @io::@arg1[%c0] : memref<16xi32>
+  %2 = hal.interface.binding.subspan @io::@ret0[%c0] : memref<16xf32>
+  %3 = "gpu.block_id"() {dimension = "x"} : () -> index
+  %4 = "gpu.block_dim"() {dimension = "x"} : () -> index
+  %5 = "gpu.thread_id"() {dimension = "x"} : () -> index
+  %6 = muli %3, %4 : index
+  %7 = addi %6, %5 : index
+  %9 = memref.load %0[%7] : memref<?xf32>
+  %10 = memref.load %1[%7] : memref<16xi32>
+  %11 = sitofp %10 : i32 to f32
+  %12 = addf %9, %11 : f32
+  memref.store %12, %2[%7] : memref<16xf32>
+  return
+}
+hal.interface @io attributes {sym_visibility = "private"} {
+  hal.interface.binding @arg0, set=0, binding=4, type="StorageBuffer", access="Read"
+  hal.interface.binding @arg1, set=0, binding=0, type="StorageBuffer", access="Read"
+  hal.interface.binding @ret0, set=1, binding=2, type="StorageBuffer", access="Write|Discard"
+}
+
+// CHECK-LABEL: llvm.func @abs_dynamic
+//  CHECK-SAME: (%[[ARG0:.+]]: !llvm.ptr<i32> {llvm.align = 16 : i32},
+//  CHECK-SAME:  %[[ARG1:.+]]: !llvm.ptr<f32> {llvm.align = 16 : i32}, %[[ARG2:.+]]: !llvm.ptr<f32> {llvm.align = 16 : i32},
+//  CHECK-SAME:  %[[ARG3:.+]]: i32, %[[ARG4:.+]]: i32)
+//       CHECK:   %[[C128:.+]] = llvm.mlir.constant(128 : index) : i64
+//       CHECK:   %{{.*}} = llvm.zext %[[ARG4]] : i32 to i64
+//       CHECK:   %[[PTRI8:.+]] = llvm.bitcast %[[ARG1]] : !llvm.ptr<f32> to !llvm.ptr<i8>
+//       CHECK:   %[[OFF:.+]] = llvm.getelementptr %[[PTRI8]][%[[C128]]] : (!llvm.ptr<i8>, i64) -> !llvm.ptr<i8>
+//       CHECK:   %[[PTR:.+]] = llvm.bitcast %[[OFF]] : !llvm.ptr<i8> to !llvm.ptr<f32>
+//       CHECK:   llvm.insertvalue %[[PTR]], %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<1 x i64>, array<1 x i64>)>
+//       CHECK:   llvm.insertvalue %[[PTR]], %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<1 x i64>, array<1 x i64>)>
+//      CHECK:    nvvm.read.ptx.sreg.tid.x
+//      CHECK:    llvm.fadd
diff --git a/iree/compiler/Codegen/SPIRV/AdrenoConfig.cpp b/iree/compiler/Codegen/SPIRV/AdrenoConfig.cpp
new file mode 100644
index 0000000..c8b441f
--- /dev/null
+++ b/iree/compiler/Codegen/SPIRV/AdrenoConfig.cpp
@@ -0,0 +1,125 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+//===- AdrenoConfig.h - Adreno CodeGen Configurations ---------------------===//
+//
+// This file contains CodeGen configurations for Adreno GPUs.
+//
+//===----------------------------------------------------------------------===//
+
+#include <array>
+
+#include "iree/compiler/Codegen/SPIRV/KernelConfig.h"
+#include "iree/compiler/Codegen/Transforms/Transforms.h"
+#include "iree/compiler/Codegen/Utils/Utils.h"
+#include "llvm/Support/MathExtras.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace detail {
+
+//===----------------------------------------------------------------------===//
+// Matmul
+//===----------------------------------------------------------------------===//
+
+static LogicalResult setOpConfig(linalg::MatmulOp op) {
+  ArrayRef<int64_t> lhsShape = getUntiledShape(op.inputs()[0]);
+  ArrayRef<int64_t> rhsShape = getUntiledShape(op.inputs()[1]);
+  if (llvm::any_of(lhsShape, ShapedType::isDynamic)) return success();
+  if (llvm::any_of(rhsShape, ShapedType::isDynamic)) return success();
+
+  int64_t dimM = lhsShape[0], dimK = lhsShape[1], dimN = rhsShape[1];
+
+  // The core idea is to distribute the matmul M/N dimension to the workgroup
+  // Y/X dimension, with each thread in a workgroup handling multiple vector
+  // elements. We start from the best (X, Y) and the tiling sizes for (M, N, K)
+  // and try different configurations by scaling them down until we find a
+  // configuration that can perfectly tile the input matmul.
+
+  const int64_t bestX = 32, bestY = 2;
+  const int64_t bestThreadM = 8, bestThreadN = 8, bestThreadK = 4;
+
+  int64_t residualThreads = bestX * bestY;
+  int64_t residualTilingFactor = (bestThreadM + bestThreadK) * bestThreadN;
+
+  SmallVector<int64_t, 3> workgroupSize(3, 1);        // (X, Y, Z)
+  SmallVector<int64_t, 4> workgroupTileSizes(3, 0);   // (M, N, K)
+  SmallVector<int64_t, 4> invocationTileSizes(3, 0);  // (M, N, K)
+
+  // Deduce the configuration for the N dimension. Start with the best workgroup
+  // X size, and reduce by a factor of two each time.
+  for (int64_t x = bestX; x >= 2; x >>= 1) {
+    // Handle 4 elements per thread. We need this for vectorized load.
+    int64_t chosenTileSize = 4;
+    if (dimN % (x * chosenTileSize) == 0) {
+      workgroupSize[0] = x;
+      workgroupTileSizes[1] = x * chosenTileSize;
+      invocationTileSizes[1] = chosenTileSize;
+      residualThreads /= x;
+      assert(residualTilingFactor % chosenTileSize == 0);
+      residualTilingFactor /= chosenTileSize;
+      break;
+    }
+  }
+  if (workgroupTileSizes[1] == 0) return success();
+
+  // Deduce the configuration for the M dimension. Start with the best workgroup
+  // Y size, and reduce by a factor of two each time.
+  for (int64_t y = residualThreads; y >= 1; y >>= 1) {
+    int64_t chosenTileSize = 0;
+    // Reduce the thread tiling size by one each time. We read one row each
+    // time; so it's fine to not be some power of two here.
+    for (int64_t t = bestThreadM; t >= 1; --t) {
+      if (dimM % (y * t) == 0) {
+        chosenTileSize = t;
+        break;
+      }
+    }
+    if (chosenTileSize) {
+      workgroupSize[1] = y;
+      workgroupTileSizes[0] = y * chosenTileSize;
+      invocationTileSizes[0] = chosenTileSize;
+      assert(residualTilingFactor > chosenTileSize);
+      residualTilingFactor -= chosenTileSize;
+      break;
+    }
+  }
+  if (workgroupTileSizes[0] == 0) return success();
+
+  // Deduce the configuration for the K dimension. We need some power of two
+  // here so that we can do vector load.
+  for (int64_t t = llvm::PowerOf2Floor(residualTilingFactor); t >= 1; t >>= 1) {
+    if (dimK % t == 0) {
+      workgroupTileSizes[2] = invocationTileSizes[2] = t;
+      break;
+    }
+  }
+
+  auto pipeline = IREE::HAL::DispatchLoweringPassPipeline::SPIRVVectorize;
+  TileSizesListType tileSizes;
+  tileSizes.push_back(workgroupTileSizes);
+  tileSizes.emplace_back();
+  tileSizes.push_back(invocationTileSizes);
+  return setOpConfigAndEntryPointFnTranslation(op->getParentOfType<FuncOp>(),
+                                               op, tileSizes, {}, pipeline,
+                                               workgroupSize);
+}
+
+//===----------------------------------------------------------------------===//
+// Entry Point
+//===----------------------------------------------------------------------===//
+
+LogicalResult setAdrenoCodeGenConfig(const spirv::TargetEnv &,
+                                     Operation *rootOp) {
+  return TypeSwitch<Operation *, LogicalResult>(rootOp)
+      .Case<linalg::MatmulOp>([](auto op) { return setOpConfig(op); })
+      .Default([](Operation *) { return success(); });
+}
+
+}  // namespace detail
+}  // namespace iree_compiler
+}  // namespace mlir
diff --git a/iree/compiler/Codegen/SPIRV/BUILD b/iree/compiler/Codegen/SPIRV/BUILD
index 4d68549..4915dce 100644
--- a/iree/compiler/Codegen/SPIRV/BUILD
+++ b/iree/compiler/Codegen/SPIRV/BUILD
@@ -13,6 +13,7 @@
 cc_library(
     name = "SPIRV",
     srcs = [
+        "AdrenoConfig.cpp",
         "ConvertToSPIRVPass.cpp",
         "KernelConfig.cpp",
         "MaliConfig.cpp",
diff --git a/iree/compiler/Codegen/SPIRV/CMakeLists.txt b/iree/compiler/Codegen/SPIRV/CMakeLists.txt
index 313e988..3cadce3 100644
--- a/iree/compiler/Codegen/SPIRV/CMakeLists.txt
+++ b/iree/compiler/Codegen/SPIRV/CMakeLists.txt
@@ -18,6 +18,7 @@
     "MemorySpace.h"
     "Utils.h"
   SRCS
+    "AdrenoConfig.cpp"
     "ConvertToSPIRVPass.cpp"
     "KernelConfig.cpp"
     "MaliConfig.cpp"
diff --git a/iree/compiler/Codegen/SPIRV/KernelConfig.cpp b/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
index a2c536a..d7e72f3 100644
--- a/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
+++ b/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
@@ -252,6 +252,9 @@
     case spirv::Vendor::NVIDIA:
       result = detail::setNVIDIACodeGenConfig(targetEnv, rootOp);
       break;
+    case spirv::Vendor::Qualcomm:
+      result = detail::setAdrenoCodeGenConfig(targetEnv, rootOp);
+      break;
     default:
       break;
   }
diff --git a/iree/compiler/Codegen/SPIRV/KernelConfig.h b/iree/compiler/Codegen/SPIRV/KernelConfig.h
index 68660c0..61f2a13 100644
--- a/iree/compiler/Codegen/SPIRV/KernelConfig.h
+++ b/iree/compiler/Codegen/SPIRV/KernelConfig.h
@@ -26,21 +26,20 @@
 
 namespace detail {
 
-/// If the given `rootOp` has known good CodeGen configuration for Mali GPUs,
-/// attaches a `translation.info` attribute to the entry point containing
-/// `rootOp` and a `lowering.config` attribute to `rootOp`.
+/// Sets CodeGen configuration for GPUs from a specific vendor.
+///
+/// If the given `rootOp` has known good CodeGen configuration, attaches a
+/// `translation.info` attribute to the entry point containing `rootOp` and a
+/// `lowering.config` attribute to `rootOp`.
+///
 /// Returns success when either no configuration is found or a configuration is
 /// successfullly attached as attribute. Returns failure only when there is an
 /// issue attaching the attribute.
+
+LogicalResult setAdrenoCodeGenConfig(const spirv::TargetEnv &targetEnv,
+                                     Operation *rootOp);
 LogicalResult setMaliCodeGenConfig(const spirv::TargetEnv &targetEnv,
                                    Operation *rootOp);
-
-/// If the given `rootOp` has known good CodeGen configuration for NVIDIA GPUs,
-/// attaches a `translation.info` attribute to the entry point containing
-/// `rootOp` and a `lowering.config` attribute to `rootOp`.
-/// Returns success when either no configuration is found or a configuration is
-/// successfullly attached as attribute. Returns failure only when there is an
-/// issue attaching the attribute.
 LogicalResult setNVIDIACodeGenConfig(const spirv::TargetEnv &targetEnv,
                                      Operation *rootOp);
 
diff --git a/iree/compiler/Codegen/SPIRV/test/BUILD b/iree/compiler/Codegen/SPIRV/test/BUILD
index fc4402b..330e6a3 100644
--- a/iree/compiler/Codegen/SPIRV/test/BUILD
+++ b/iree/compiler/Codegen/SPIRV/test/BUILD
@@ -19,13 +19,14 @@
     name = "lit",
     srcs = enforce_glob(
         [
+            "config_adreno_matmul.mlir",
+            "config_linalg_ext_ops.mlir",
             "convert_to_spirv.mlir",
             "distribute_to_global_id.mlir",
             "fold_gpu_procid_uses.mlir",
             "pipeline_matmul_cooperative_matrix.mlir",
             "pipeline_matmul_vectorization.mlir",
             "remove_one_trip_tiled_loop.mlir",
-            "set_lowering_config.mlir",
             "tile_and_distribute_scatter.mlir",
             "tile_and_distribute_sort.mlir",
             "tile_and_vectorize.mlir",
diff --git a/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt b/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt
index 31ccf04..ef888fb 100644
--- a/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt
+++ b/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt
@@ -14,13 +14,14 @@
   NAME
     lit
   SRCS
+    "config_adreno_matmul.mlir"
+    "config_linalg_ext_ops.mlir"
     "convert_to_spirv.mlir"
     "distribute_to_global_id.mlir"
     "fold_gpu_procid_uses.mlir"
     "pipeline_matmul_cooperative_matrix.mlir"
     "pipeline_matmul_vectorization.mlir"
     "remove_one_trip_tiled_loop.mlir"
-    "set_lowering_config.mlir"
     "tile_and_distribute_scatter.mlir"
     "tile_and_distribute_sort.mlir"
     "tile_and_vectorize.mlir"
diff --git a/iree/compiler/Codegen/SPIRV/test/config_adreno_matmul.mlir b/iree/compiler/Codegen/SPIRV/test/config_adreno_matmul.mlir
new file mode 100644
index 0000000..aee34bd
--- /dev/null
+++ b/iree/compiler/Codegen/SPIRV/test/config_adreno_matmul.mlir
@@ -0,0 +1,379 @@
+// RUN: iree-opt -split-input-file -mlir-print-local-scope -pass-pipeline='hal.executable(hal.executable.variant(iree-spirv-lower-executable-target-pass{test-lowering-configuration=true}))' %s | IreeFileCheck %s
+
+// Large matmul that can match the best tiling scheme.
+
+hal.executable @matmul_1024x2048x512 {
+  hal.interface @io {
+    hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard"
+  }
+  hal.executable.variant @vulkan_spirv_fb, target = #hal.executable.target<"vulkan", "vulkan-spirv-fb", {
+      spv.target_env = #spv.target_env<#spv.vce<v1.4, [Shader], []>, Qualcomm:IntegratedGPU, {
+        max_compute_shared_memory_size = 32768 : i32,
+        max_compute_workgroup_invocations = 1024 : i32,
+        max_compute_workgroup_size = dense<[1024, 1024, 64]> : vector<3xi32>,
+        subgroup_size = 64 : i32}>
+    }> {
+    hal.executable.entry_point @matmul_1024x2048x512 attributes {interface = @io, ordinal = 0 : index}
+    builtin.module  {
+      func @matmul_1024x2048x512() {
+        %c0 = constant 0 : index
+        %c2048 = constant 2048 : index
+        %c1024 = constant 1024 : index
+        %cst = constant 0.000000e+00 : f32
+        %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:1024x512xf32>
+        %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:512x2048xf32>
+        %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:1024x2048xf32>
+        %workgroup_size_x = hal.interface.workgroup.size[0] : index
+        %workgroup_size_y = hal.interface.workgroup.size[1] : index
+        %workgroup_id_x = hal.interface.workgroup.id[0] : index
+        %workgroup_count_x = hal.interface.workgroup.count[0] : index
+        %workgroup_id_y = hal.interface.workgroup.id[1] : index
+        %workgroup_count_y = hal.interface.workgroup.count[1] : index
+        %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y]
+        %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y]
+        scf.for %arg0 = %3 to %c1024 step %4 {
+          %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x]
+          %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x]
+          scf.for %arg1 = %5 to %c2048 step %6 {
+            %7 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 1024)>(%arg0)[%workgroup_size_y]
+            %8 = flow.dispatch.tensor.load %0, offsets = [%arg0, 0], sizes = [%7, 512], strides = [1, 1] : !flow.dispatch.tensor<readonly:1024x512xf32> -> tensor<?x512xf32>
+            %9 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 2048)>(%arg1)[%workgroup_size_x]
+            %10 = flow.dispatch.tensor.load %1, offsets = [0, %arg1], sizes = [512, %9], strides = [1, 1] : !flow.dispatch.tensor<readonly:512x2048xf32> -> tensor<512x?xf32>
+            %11 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 1024)>(%arg0)[%workgroup_size_y]
+            %12 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 2048)>(%arg1)[%workgroup_size_x]
+            %13 = affine.min affine_map<(d0)[s0] -> (-d0 + 1024, s0)>(%arg0)[%workgroup_size_y]
+            %14 = affine.min affine_map<(d0)[s0] -> (-d0 + 2048, s0)>(%arg1)[%workgroup_size_x]
+            %15 = linalg.init_tensor [%13, %14] : tensor<?x?xf32>
+            %16 = linalg.fill(%cst, %15) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
+            %17 = linalg.matmul {__internal_linalg_transform__ = "workgroup"} ins(%8, %10 : tensor<?x512xf32>, tensor<512x?xf32>) outs(%16 : tensor<?x?xf32>) -> tensor<?x?xf32>
+            flow.dispatch.tensor.store %17, %2, offsets = [%arg0, %arg1], sizes = [%11, %12], strides = [1, 1] : tensor<?x?xf32> -> !flow.dispatch.tensor<writeonly:1024x2048xf32>
+          }
+        }
+        return
+      }
+      hal.interface @io attributes {sym_visibility = "private"} {
+        hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read"
+        hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read"
+        hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard"
+      }
+    }
+  }
+}
+
+//          CHECK-LABEL: hal.executable.entry_point public @matmul_1024x2048x512
+//           CHECK-SAME:   translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [128, 16]}
+//           CHECK-SAME:   workgroup_size = [32 : index, 2 : index, 1 : index]
+//           CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index):
+//           CHECK-NEXT:   %[[ONE:.+]] = constant 1 : index
+//           CHECK-NEXT:   %[[X_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 128)>()[%[[X]]]
+//           CHECK-NEXT:   %[[Y_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 16)>()[%[[Y]]]
+//           CHECK-NEXT:   hal.return %[[X_COUNT]], %[[Y_COUNT]], %[[ONE]]
+
+//                CHECK: func @matmul_1024x2048x512()
+//                CHECK:   linalg.matmul
+//  CHECK-SAME{LITERAL}:     lowering.config = {tileSizes = [[16, 128, 16], [], [8, 4, 16]]}
+
+// -----
+
+// Small matmul N that can still tile to all threads in a workgroup.
+
+hal.executable @matmul_3136x24x96 {
+  hal.interface @io {
+    hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard"
+  }
+  hal.executable.variant @vulkan_spirv_fb, target = #hal.executable.target<"vulkan", "vulkan-spirv-fb", {
+      spv.target_env = #spv.target_env<#spv.vce<v1.4, [Shader], []>, Qualcomm:IntegratedGPU, {
+        max_compute_shared_memory_size = 32768 : i32,
+        max_compute_workgroup_invocations = 1024 : i32,
+        max_compute_workgroup_size = dense<[1024, 1024, 64]> : vector<3xi32>,
+        subgroup_size = 64 : i32}>
+    }> {
+    hal.executable.entry_point @matmul_3136x24x96 attributes {interface = @io, ordinal = 0 : index}
+    builtin.module  {
+      func @matmul_3136x24x96() {
+        %c0 = constant 0 : index
+        %c24 = constant 24 : index
+        %c3136 = constant 3136 : index
+        %cst = constant 0.000000e+00 : f32
+        %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:3136x96xf32>
+        %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:96x24xf32>
+        %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:3136x24xf32>
+        %workgroup_size_x = hal.interface.workgroup.size[0] : index
+        %workgroup_size_y = hal.interface.workgroup.size[1] : index
+        %workgroup_id_x = hal.interface.workgroup.id[0] : index
+        %workgroup_count_x = hal.interface.workgroup.count[0] : index
+        %workgroup_id_y = hal.interface.workgroup.id[1] : index
+        %workgroup_count_y = hal.interface.workgroup.count[1] : index
+        %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y]
+        %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y]
+        scf.for %arg0 = %3 to %c3136 step %4 {
+          %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x]
+          %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x]
+          scf.for %arg1 = %5 to %c24 step %6 {
+            %7 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 3136)>(%arg0)[%workgroup_size_y]
+            %8 = flow.dispatch.tensor.load %0, offsets = [%arg0, 0], sizes = [%7, 96], strides = [1, 1] : !flow.dispatch.tensor<readonly:3136x96xf32> -> tensor<?x96xf32>
+            %9 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 24)>(%arg1)[%workgroup_size_x]
+            %10 = flow.dispatch.tensor.load %1, offsets = [0, %arg1], sizes = [96, %9], strides = [1, 1] : !flow.dispatch.tensor<readonly:96x24xf32> -> tensor<96x?xf32>
+            %11 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 3136)>(%arg0)[%workgroup_size_y]
+            %12 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 24)>(%arg1)[%workgroup_size_x]
+            %13 = affine.min affine_map<(d0)[s0] -> (-d0 + 3136, s0)>(%arg0)[%workgroup_size_y]
+            %14 = affine.min affine_map<(d0)[s0] -> (-d0 + 24, s0)>(%arg1)[%workgroup_size_x]
+            %15 = linalg.init_tensor [%13, %14] : tensor<?x?xf32>
+            %16 = linalg.fill(%cst, %15) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
+            %17 = linalg.matmul {__internal_linalg_transform__ = "workgroup"} ins(%8, %10 : tensor<?x96xf32>, tensor<96x?xf32>) outs(%16 : tensor<?x?xf32>) -> tensor<?x?xf32>
+            flow.dispatch.tensor.store %17, %2, offsets = [%arg0, %arg1], sizes = [%11, %12], strides = [1, 1] : tensor<?x?xf32> -> !flow.dispatch.tensor<writeonly:3136x24xf32>
+          }
+        }
+        return
+      }
+      hal.interface @io attributes {sym_visibility = "private"} {
+        hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read"
+        hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read"
+        hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard"
+      }
+    }
+  }
+}
+
+//          CHECK-LABEL: hal.executable.entry_point public @matmul_3136x24x96
+//           CHECK-SAME:   translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [8, 224]}
+//           CHECK-SAME:   workgroup_size = [2 : index, 32 : index, 1 : index]
+//           CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index):
+//           CHECK-NEXT:   %[[ONE:.+]] = constant 1 : index
+//           CHECK-NEXT:   %[[X_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%[[X]]]
+//           CHECK-NEXT:   %[[Y_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 224)>()[%[[Y]]]
+//           CHECK-NEXT:   hal.return %[[X_COUNT]], %[[Y_COUNT]], %[[ONE]]
+
+//                CHECK: func @matmul_3136x24x96()
+//                CHECK:   linalg.matmul
+//  CHECK-SAME{LITERAL}:     lowering.config = {tileSizes = [[224, 8, 16], [], [7, 4, 16]]}
+
+// -----
+
+// Small matmul M that can still tile to all threads in a workgroup.
+
+hal.executable @matmul_196x64x192 {
+  hal.interface @io {
+    hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard"
+  }
+  hal.executable.variant @vulkan_spirv_fb, target = #hal.executable.target<"vulkan", "vulkan-spirv-fb", {
+      spv.target_env = #spv.target_env<#spv.vce<v1.4, [Shader], []>, Qualcomm:IntegratedGPU, {
+        max_compute_shared_memory_size = 32768 : i32,
+        max_compute_workgroup_invocations = 1024 : i32,
+        max_compute_workgroup_size = dense<[1024, 1024, 64]> : vector<3xi32>,
+        subgroup_size = 64 : i32}>
+    }> {
+    hal.executable.entry_point @matmul_196x64x192 attributes {interface = @io, ordinal = 0 : index}
+    builtin.module  {
+      func @matmul_196x64x192() {
+        %c0 = constant 0 : index
+        %c64 = constant 64 : index
+        %c196 = constant 196 : index
+        %cst = constant 0.000000e+00 : f32
+        %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:196x192xf32>
+        %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:192x64xf32>
+        %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:196x64xf32>
+        %workgroup_size_x = hal.interface.workgroup.size[0] : index
+        %workgroup_size_y = hal.interface.workgroup.size[1] : index
+        %workgroup_id_x = hal.interface.workgroup.id[0] : index
+        %workgroup_count_x = hal.interface.workgroup.count[0] : index
+        %workgroup_id_y = hal.interface.workgroup.id[1] : index
+        %workgroup_count_y = hal.interface.workgroup.count[1] : index
+        %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y]
+        %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y]
+        scf.for %arg0 = %3 to %c196 step %4 {
+          %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x]
+          %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x]
+          scf.for %arg1 = %5 to %c64 step %6 {
+            %7 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 196)>(%arg0)[%workgroup_size_y]
+            %8 = flow.dispatch.tensor.load %0, offsets = [%arg0, 0], sizes = [%7, 192], strides = [1, 1] : !flow.dispatch.tensor<readonly:196x192xf32> -> tensor<?x192xf32>
+            %9 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 64)>(%arg1)[%workgroup_size_x]
+            %10 = flow.dispatch.tensor.load %1, offsets = [0, %arg1], sizes = [192, %9], strides = [1, 1] : !flow.dispatch.tensor<readonly:192x64xf32> -> tensor<192x?xf32>
+            %11 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 196)>(%arg0)[%workgroup_size_y]
+            %12 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 64)>(%arg1)[%workgroup_size_x]
+            %13 = affine.min affine_map<(d0)[s0] -> (-d0 + 196, s0)>(%arg0)[%workgroup_size_y]
+            %14 = affine.min affine_map<(d0)[s0] -> (-d0 + 64, s0)>(%arg1)[%workgroup_size_x]
+            %15 = linalg.init_tensor [%13, %14] : tensor<?x?xf32>
+            %16 = linalg.fill(%cst, %15) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
+            %17 = linalg.matmul {__internal_linalg_transform__ = "workgroup"} ins(%8, %10 : tensor<?x192xf32>, tensor<192x?xf32>) outs(%16 : tensor<?x?xf32>) -> tensor<?x?xf32>
+            flow.dispatch.tensor.store %17, %2, offsets = [%arg0, %arg1], sizes = [%11, %12], strides = [1, 1] : tensor<?x?xf32> -> !flow.dispatch.tensor<writeonly:196x64xf32>
+          }
+        }
+        return
+      }
+      hal.interface @io attributes {sym_visibility = "private"} {
+        hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read"
+        hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read"
+        hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard"
+      }
+    }
+  }
+}
+
+//          CHECK-LABEL: hal.executable.entry_point public @matmul_196x64x192
+//           CHECK-SAME:   translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [64, 28]}
+//           CHECK-SAME:   workgroup_size = [16 : index, 4 : index, 1 : index]
+//           CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index):
+//           CHECK-NEXT:   %[[ONE:.+]] = constant 1 : index
+//           CHECK-NEXT:   %[[X_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 64)>()[%[[X]]]
+//           CHECK-NEXT:   %[[Y_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 28)>()[%[[Y]]]
+//           CHECK-NEXT:   hal.return %[[X_COUNT]], %[[Y_COUNT]], %[[ONE]]
+
+//                CHECK: func @matmul_196x64x192()
+//                CHECK:   linalg.matmul
+//  CHECK-SAME{LITERAL}:      lowering.config = {tileSizes = [[28, 64, 16], [], [7, 4, 16]]}
+
+// -----
+
+// Small matmul K that can still tile to all threads in a workgroup.
+
+hal.executable @matmul_12544x96x16 {
+  hal.interface @io {
+    hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard"
+  }
+  hal.executable.variant @vulkan_spirv_fb, target = #hal.executable.target<"vulkan", "vulkan-spirv-fb", {
+      spv.target_env = #spv.target_env<#spv.vce<v1.4, [Shader], []>, Qualcomm:IntegratedGPU, {
+        max_compute_shared_memory_size = 32768 : i32,
+        max_compute_workgroup_invocations = 1024 : i32,
+        max_compute_workgroup_size = dense<[1024, 1024, 64]> : vector<3xi32>,
+        subgroup_size = 64 : i32}>
+    }> {
+    hal.executable.entry_point @matmul_12544x96x16 attributes {interface = @io, ordinal = 0 : index}
+    builtin.module  {
+      func @matmul_12544x96x16() {
+        %c0 = constant 0 : index
+        %c96 = constant 96 : index
+        %c12544 = constant 12544 : index
+        %cst = constant 0.000000e+00 : f32
+        %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : memref<12544x16xf32>
+        %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : memref<16x96xf32>
+        %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : memref<12544x96xf32>
+        %workgroup_size_x = hal.interface.workgroup.size[0] : index
+        %workgroup_size_y = hal.interface.workgroup.size[1] : index
+        %workgroup_id_x = hal.interface.workgroup.id[0] : index
+        %workgroup_count_x = hal.interface.workgroup.count[0] : index
+        %workgroup_id_y = hal.interface.workgroup.id[1] : index
+        %workgroup_count_y = hal.interface.workgroup.count[1] : index
+        %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y]
+        %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y]
+        scf.for %arg0 = %3 to %c12544 step %4 {
+          %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x]
+          %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x]
+          scf.for %arg1 = %5 to %c96 step %6 {
+            %7 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 12544)>(%arg0)[%workgroup_size_y]
+            %8 = memref.subview %0[%arg0, 0] [%7, 16] [1, 1] : memref<12544x16xf32> to memref<?x16xf32, affine_map<(d0, d1)[s0] -> (d0 * 16 + s0 + d1)>>
+            %9 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 96)>(%arg1)[%workgroup_size_x]
+            %10 = memref.subview %1[0, %arg1] [16, %9] [1, 1] : memref<16x96xf32> to memref<16x?xf32, affine_map<(d0, d1)[s0] -> (d0 * 96 + s0 + d1)>>
+            %11 = memref.subview %2[%arg0, %arg1] [%7, %9] [1, 1] : memref<12544x96xf32> to memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * 96 + s0 + d1)>>
+            linalg.fill(%cst, %11) : f32, memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * 96 + s0 + d1)>>
+            linalg.matmul {__internal_linalg_transform__ = "workgroup"} ins(%8, %10 : memref<?x16xf32, affine_map<(d0, d1)[s0] -> (d0 * 16 + s0 + d1)>>, memref<16x?xf32, affine_map<(d0, d1)[s0] -> (d0 * 96 + s0 + d1)>>) outs(%11 : memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * 96 + s0 + d1)>>)
+          }
+        }
+        return
+      }
+      hal.interface @io attributes {sym_visibility = "private"} {
+        hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read"
+        hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read"
+        hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard"
+      }
+    }
+  }
+}
+
+//          CHECK-LABEL: hal.executable.entry_point public @matmul_12544x96x16
+//           CHECK-SAME:   translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [32, 64]}
+//           CHECK-SAME:   workgroup_size = [8 : index, 8 : index, 1 : index]
+//           CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index):
+//           CHECK-NEXT:   %[[ONE:.+]] = constant 1 : index
+//           CHECK-NEXT:   %[[X_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 32)>()[%[[X]]]
+//           CHECK-NEXT:   %[[Y_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 64)>()[%[[Y]]]
+//           CHECK-NEXT:   hal.return %[[X_COUNT]], %[[Y_COUNT]], %[[ONE]]
+
+//                CHECK: func @matmul_12544x96x16()
+//                CHECK:   linalg.matmul
+//  CHECK-SAME{LITERAL}:     lowering.config =  {tileSizes = [[64, 32, 16], [], [8, 4, 16]]}
+
+// -----
+
+// Odd matmul M and small N that cannot utilize all threads in a workgroup.
+
+hal.executable @matmul_49x160x576 {
+  hal.interface @io {
+    hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard"
+  }
+  hal.executable.variant @vulkan_spirv_fb, target = #hal.executable.target<"vulkan", "vulkan-spirv-fb", {
+      spv.target_env = #spv.target_env<#spv.vce<v1.4, [Shader], []>, Qualcomm:IntegratedGPU, {
+        max_compute_shared_memory_size = 32768 : i32,
+        max_compute_workgroup_invocations = 1024 : i32,
+        max_compute_workgroup_size = dense<[1024, 1024, 64]> : vector<3xi32>,
+        subgroup_size = 64 : i32}>
+    }> {
+    hal.executable.entry_point @matmul_49x160x576 attributes {interface = @io, ordinal = 0 : index}
+    builtin.module  {
+      func @matmul_49x160x576() {
+        %c0 = constant 0 : index
+        %c160 = constant 160 : index
+        %c49 = constant 49 : index
+        %cst = constant 0.000000e+00 : f32
+        %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:49x576xf32>
+        %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:576x160xf32>
+        %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:49x160xf32>
+        %workgroup_size_x = hal.interface.workgroup.size[0] : index
+        %workgroup_size_y = hal.interface.workgroup.size[1] : index
+        %workgroup_id_x = hal.interface.workgroup.id[0] : index
+        %workgroup_count_x = hal.interface.workgroup.count[0] : index
+        %workgroup_id_y = hal.interface.workgroup.id[1] : index
+        %workgroup_count_y = hal.interface.workgroup.count[1] : index
+        %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y]
+        %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y]
+        scf.for %arg0 = %3 to %c49 step %4 {
+          %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x]
+          %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x]
+          scf.for %arg1 = %5 to %c160 step %6 {
+            %7 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 49)>(%arg0)[%workgroup_size_y]
+            %8 = flow.dispatch.tensor.load %0, offsets = [%arg0, 0], sizes = [%7, 576], strides = [1, 1] : !flow.dispatch.tensor<readonly:49x576xf32> -> tensor<?x576xf32>
+            %9 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 160)>(%arg1)[%workgroup_size_x]
+            %10 = flow.dispatch.tensor.load %1, offsets = [0, %arg1], sizes = [576, %9], strides = [1, 1] : !flow.dispatch.tensor<readonly:576x160xf32> -> tensor<576x?xf32>
+            %11 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 49)>(%arg0)[%workgroup_size_y]
+            %12 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 160)>(%arg1)[%workgroup_size_x]
+            %13 = affine.min affine_map<(d0)[s0] -> (-d0 + 49, s0)>(%arg0)[%workgroup_size_y]
+            %14 = affine.min affine_map<(d0)[s0] -> (-d0 + 160, s0)>(%arg1)[%workgroup_size_x]
+            %15 = linalg.init_tensor [%13, %14] : tensor<?x?xf32>
+            %16 = linalg.fill(%cst, %15) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
+            %17 = linalg.matmul {__internal_linalg_transform__ = "workgroup"} ins(%8, %10 : tensor<?x576xf32>, tensor<576x?xf32>) outs(%16 : tensor<?x?xf32>) -> tensor<?x?xf32>
+            flow.dispatch.tensor.store %17, %2, offsets = [%arg0, %arg1], sizes = [%11, %12], strides = [1, 1] : tensor<?x?xf32> -> !flow.dispatch.tensor<writeonly:49x160xf32>
+          }
+        }
+        return
+      }
+      hal.interface @io attributes {sym_visibility = "private"} {
+        hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read"
+        hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read"
+        hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard"
+      }
+    }
+  }
+}
+
+//          CHECK-LABEL: hal.executable.entry_point public @matmul_49x160x576
+//           CHECK-SAME:   translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [32, 7]}
+//           CHECK-SAME:   workgroup_size = [8 : index, 1 : index, 1 : index]
+//           CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index):
+//           CHECK-NEXT:   %[[ONE:.+]] = constant 1 : index
+//           CHECK-NEXT:   %[[X_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 32)>()[%[[X]]]
+//           CHECK-NEXT:   %[[Y_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 7)>()[%[[Y]]]
+//           CHECK-NEXT:   hal.return %[[X_COUNT]], %[[Y_COUNT]], %[[ONE]]
+
+//                CHECK: func @matmul_49x160x576()
+//                CHECK:   linalg.matmul
+//  CHECK-SAME{LITERAL}:     lowering.config = {tileSizes = [[7, 32, 16], [], [7, 4, 16]]}
diff --git a/iree/compiler/Codegen/SPIRV/test/set_lowering_config.mlir b/iree/compiler/Codegen/SPIRV/test/config_linalg_ext_ops.mlir
similarity index 100%
rename from iree/compiler/Codegen/SPIRV/test/set_lowering_config.mlir
rename to iree/compiler/Codegen/SPIRV/test/config_linalg_ext_ops.mlir
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
index 87affe9..975a41b 100644
--- a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
+++ b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
@@ -104,8 +104,7 @@
   // The generated ABI wrappers assume such an expansion and will generate code
   // to produce it from the original reflection metadata captured in the
   // previous pass.
-  passManager.addNestedPass<mlir::FuncOp>(
-      Shape::createExpandFunctionDynamicDimsPass());
+  passManager.addPass(Shape::createExpandFunctionDynamicDimsPass());
 
   // Special case peephole optimizations.
   {
diff --git a/iree/compiler/Dialect/HAL/Transforms/Passes.cpp b/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
index 3b3df2c..fc00b5d 100644
--- a/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
+++ b/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
@@ -96,8 +96,7 @@
   // on explicit shape types (such as ranked_shape). After this pass, these
   // composite types will be expanded to primitives (i.e. one 'index' for each
   // dynamic dim in the case of ranked_shape).
-  passManager.addNestedPass<FuncOp>(
-      Shape::createExpandFunctionRankedShapeDimsPass());
+  passManager.addPass(Shape::createExpandFunctionRankedShapeDimsPass());
 
   passManager.addNestedPass<FuncOp>(createCanonicalizerPass());
   passManager.addNestedPass<FuncOp>(createCSEPass());
diff --git a/iree/compiler/Dialect/Shape/Transforms/FunctionSignatureExpansionPass.cpp b/iree/compiler/Dialect/Shape/Transforms/FunctionSignatureExpansionPass.cpp
index 316cd72..edd0814 100644
--- a/iree/compiler/Dialect/Shape/Transforms/FunctionSignatureExpansionPass.cpp
+++ b/iree/compiler/Dialect/Shape/Transforms/FunctionSignatureExpansionPass.cpp
@@ -21,10 +21,95 @@
 namespace mlir {
 namespace iree_compiler {
 namespace Shape {
+
+static LogicalResult expandCallGraphTypesForCall(
+    CallOp callOp, const TypeExpander &typeExpander) {
+  Location loc = callOp.getLoc();
+
+  OpBuilder builder(callOp);
+  // Create the operands list of the new `CallOp`.
+  SmallVector<Value> newOperands;
+  if (failed(typeExpander.expandSourceValuesToTarget(
+          loc, llvm::to_vector<6>(callOp.getOperands()), newOperands,
+          builder))) {
+    return failure();
+  }
+
+  // Create the new result types for the new `CallOp` and track the indices
+  // in the new call op's results that correspond to the old call op's
+  // results.
+  //
+  // expandedResultIndices[i] = "list of new result indices that old result
+  // i expanded to".
+  SmallVector<Type> newResultTypes;
+  SmallVector<SmallVector<unsigned>> expandedResultIndices;
+  for (Type resultType : callOp.getResultTypes()) {
+    unsigned oldSize = newResultTypes.size();
+    if (failed(typeExpander.convertType(resultType, newResultTypes))) {
+      return failure();
+    }
+    auto &resultMapping = expandedResultIndices.emplace_back();
+    for (unsigned i = oldSize, e = newResultTypes.size(); i < e; i++)
+      resultMapping.push_back(i);
+  }
+
+  CallOp newCallOp = builder.create<CallOp>(loc, callOp.getCallee(),
+                                            newResultTypes, newOperands);
+
+  // Build a replacement value for each result to replace its uses. If a
+  // result has multiple mapping values, it needs to be materialized as a
+  // single value.
+  SmallVector<Value, 2> replacedValues;
+  replacedValues.reserve(callOp.getNumResults());
+  for (unsigned i = 0, e = callOp.getNumResults(); i < e; ++i) {
+    auto decomposedValues = llvm::to_vector<6>(
+        llvm::map_range(expandedResultIndices[i],
+                        [&](unsigned i) { return newCallOp.getResult(i); }));
+    if (decomposedValues.empty()) {
+      // No replacement is required.
+      replacedValues.push_back(nullptr);
+    } else if (decomposedValues.size() == 1) {
+      replacedValues.push_back(decomposedValues.front());
+    } else {
+      Value materialized = typeExpander.castToSource(loc, callOp.getType(i),
+                                                     decomposedValues, builder);
+      replacedValues.push_back(materialized);
+    }
+  }
+  callOp.replaceAllUsesWith(replacedValues);
+  return success();
+}
+
+static LogicalResult expandCallGraphTypes(ModuleOp module,
+                                          const TypeExpander &typeExpander) {
+  for (auto funcOp : module.getOps<FuncOp>()) {
+    OpBuilder builder(funcOp);
+    if (failed(typeExpander.expandFunctionSignature(funcOp, builder)) ||
+        failed(typeExpander.expandAllReturnLikeTerminators<mlir::ReturnOp>(
+            funcOp, builder))) {
+      return funcOp.emitError()
+             << "couldn't expand function signature or return";
+    }
+    bool hadError = false;
+    SmallVector<Operation *> toErase;
+    funcOp.walk([&](CallOp callOp) {
+      if (failed(expandCallGraphTypesForCall(callOp, typeExpander))) {
+        hadError = true;
+        return;
+      }
+      toErase.push_back(callOp);
+    });
+    if (hadError) return failure();
+    for (Operation *op : toErase) op->erase();
+  }
+  return success();
+}
+
 namespace {
 
 class ExpandFunctionDynamicDimsPass
-    : public PassWrapper<ExpandFunctionDynamicDimsPass, FunctionPass> {
+    : public PassWrapper<ExpandFunctionDynamicDimsPass,
+                         OperationPass<ModuleOp>> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<ShapeDialect>();
   }
@@ -37,20 +122,16 @@
     return "Expands dynamic dimensions in function signatures.";
   }
 
-  void runOnFunction() override {
-    auto funcOp = getFunction();
+  void runOnOperation() override {
     auto &typeExpander = getDynamicShapeTypeExpander();
-    OpBuilder builder(funcOp);
-    if (failed(typeExpander.expandFunctionSignature(funcOp, builder)) ||
-        failed(typeExpander.expandAllReturnLikeTerminators<mlir::ReturnOp>(
-            funcOp, builder))) {
-      return signalPassFailure();
-    }
+    if (failed(expandCallGraphTypes(getOperation(), typeExpander)))
+      signalPassFailure();
   }
 };
 
 class ExpandFunctionRankedShapeDimsPass
-    : public PassWrapper<ExpandFunctionRankedShapeDimsPass, FunctionPass> {
+    : public PassWrapper<ExpandFunctionRankedShapeDimsPass,
+                         OperationPass<ModuleOp>> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<ShapeDialect>();
   }
@@ -63,15 +144,10 @@
     return "Expands ranked_shape types at function boundaries to loose dims.";
   }
 
-  void runOnFunction() override {
-    auto funcOp = getFunction();
+  void runOnOperation() override {
     auto &typeExpander = getShapeToPrimitiveTypeExpander();
-    OpBuilder builder(funcOp);
-    if (failed(typeExpander.expandFunctionSignature(funcOp, builder)) ||
-        failed(typeExpander.expandAllReturnLikeTerminators<mlir::ReturnOp>(
-            funcOp, builder))) {
-      return signalPassFailure();
-    }
+    if (failed(expandCallGraphTypes(getOperation(), typeExpander)))
+      signalPassFailure();
   }
 };
 
@@ -79,14 +155,14 @@
 
 // For any function which contains dynamic dims in its inputs or results,
 // rewrites it so that the dynamic dims are passed in/out.
-std::unique_ptr<OperationPass<FuncOp>> createExpandFunctionDynamicDimsPass() {
+std::unique_ptr<OperationPass<ModuleOp>> createExpandFunctionDynamicDimsPass() {
   return std::make_unique<Shape::ExpandFunctionDynamicDimsPass>();
 }
 
 // For any function which contains ranked_shape argument/result types,
 // expands them to individual dynamic dimensions, inserting appropriate casts
 // within the function.
-std::unique_ptr<OperationPass<FuncOp>>
+std::unique_ptr<OperationPass<ModuleOp>>
 createExpandFunctionRankedShapeDimsPass() {
   return std::make_unique<Shape::ExpandFunctionRankedShapeDimsPass>();
 }
diff --git a/iree/compiler/Dialect/Shape/Transforms/Passes.h b/iree/compiler/Dialect/Shape/Transforms/Passes.h
index ab1544a..3cc2889 100644
--- a/iree/compiler/Dialect/Shape/Transforms/Passes.h
+++ b/iree/compiler/Dialect/Shape/Transforms/Passes.h
@@ -18,12 +18,12 @@
 
 // For any function which contains dynamic dims in its inputs or results,
 // rewrites it so that the dynamic dims are passed in/out.
-std::unique_ptr<OperationPass<FuncOp>> createExpandFunctionDynamicDimsPass();
+std::unique_ptr<OperationPass<ModuleOp>> createExpandFunctionDynamicDimsPass();
 
 // For any function which contains ranked_shape argument/result types,
 // expands them to individual dynamic dimensions, inserting appropriate casts
 // within the function.
-std::unique_ptr<OperationPass<FuncOp>>
+std::unique_ptr<OperationPass<ModuleOp>>
 createExpandFunctionRankedShapeDimsPass();
 
 // Folds tensor.dim/memref.dim ops taking shape carrying ops as operands.
diff --git a/iree/compiler/Dialect/Shape/Transforms/test/expand_function_dynamic_dims.mlir b/iree/compiler/Dialect/Shape/Transforms/test/expand_function_dynamic_dims.mlir
index d0dfe66..be98e44 100644
--- a/iree/compiler/Dialect/Shape/Transforms/test/expand_function_dynamic_dims.mlir
+++ b/iree/compiler/Dialect/Shape/Transforms/test/expand_function_dynamic_dims.mlir
@@ -36,3 +36,18 @@
   // CHECK: return %[[RESULT]], %[[SHAPE]] : tensor<?x3xf32>, !shapex.ranked_shape<[?,3]>
   return %1 : tensor<?x3xf32>
 }
+
+// -----
+// CHECK-LABEL:   func @calls(
+// CHECK-SAME:                %[[ARG:.*]]: tensor<1x?x2x?xf32>,
+// CHECK-SAME:                %[[ARG_SHAPE:.*]]: !shapex.ranked_shape<[1,?,2,?]>) -> (tensor<1x?x2x?xf32>, !shapex.ranked_shape<[1,?,2,?]>) {
+func @calls(%arg0 : tensor<1x?x2x?xf32>) -> tensor<1x?x2x?xf32> {
+  // CHECK:           %[[ARG_TIED:.*]] = shapex.tie_shape %[[ARG]], %[[ARG_SHAPE]] : tensor<1x?x2x?xf32>, !shapex.ranked_shape<[1,?,2,?]>
+  // CHECK:           %[[ARG_TIED_SHAPE:.*]] = shapex.get_ranked_shape %[[ARG_TIED]] : tensor<1x?x2x?xf32> -> !shapex.ranked_shape<[1,?,2,?]>
+  // CHECK:           %[[CALL_RESULT:.*]]:2 = call @calls(%[[ARG_TIED]], %[[ARG_TIED_SHAPE]]) : (tensor<1x?x2x?xf32>, !shapex.ranked_shape<[1,?,2,?]>) -> (tensor<1x?x2x?xf32>, !shapex.ranked_shape<[1,?,2,?]>)
+  %0 = std.call @calls(%arg0) : (tensor<1x?x2x?xf32>) -> tensor<1x?x2x?xf32>
+  // CHECK:           %[[CALL_RESULT_TIED:.*]] = shapex.tie_shape %[[CALL_RESULT]]#0, %[[CALL_RESULT]]#1 : tensor<1x?x2x?xf32>, !shapex.ranked_shape<[1,?,2,?]>
+  // CHECK:           %[[CALL_RESULT_SHAPE:.*]] = shapex.get_ranked_shape %[[CALL_RESULT_TIED]] : tensor<1x?x2x?xf32> -> !shapex.ranked_shape<[1,?,2,?]>
+  // CHECK:           return %[[CALL_RESULT_TIED]], %[[CALL_RESULT_SHAPE]] : tensor<1x?x2x?xf32>, !shapex.ranked_shape<[1,?,2,?]>
+  return %0 : tensor<1x?x2x?xf32>
+}
diff --git a/iree/compiler/Dialect/Shape/Transforms/test/expand_function_ranked_shape_dims.mlir b/iree/compiler/Dialect/Shape/Transforms/test/expand_function_ranked_shape_dims.mlir
index 60e0cc4..9ad16d5 100644
--- a/iree/compiler/Dialect/Shape/Transforms/test/expand_function_ranked_shape_dims.mlir
+++ b/iree/compiler/Dialect/Shape/Transforms/test/expand_function_ranked_shape_dims.mlir
@@ -34,3 +34,12 @@
   // CHECK: return %[[CAST_IDX0]], %[[CAST_IDX1]]
   return %rs : !shapex.ranked_shape<[?,?]>
 }
+
+// -----
+// CHECK-LABEL: @calls
+// CHECK-SAME: (%arg0: index, %arg1: index) -> (index, index)
+func @calls(%arg0 :!shapex.ranked_shape<[?,?]>) -> !shapex.ranked_shape<[?,?]> {
+  // CHECK: call @calls(%{{.*}}, %{{.*}}) : (index, index) -> (index, index)
+  %0 = std.call @calls(%arg0) : (!shapex.ranked_shape<[?,?]>) -> !shapex.ranked_shape<[?,?]>
+  return %0 : !shapex.ranked_shape<[?,?]>
+}
diff --git a/iree/hal/cuda/executable_layout.c b/iree/hal/cuda/executable_layout.c
index a89c7d9..21858ba 100644
--- a/iree/hal/cuda/executable_layout.c
+++ b/iree/hal/cuda/executable_layout.c
@@ -15,6 +15,7 @@
 typedef struct iree_hal_cuda_executable_layout_t {
   iree_hal_resource_t resource;
   iree_hal_cuda_context_wrapper_t* context;
+  iree_host_size_t push_constant_base_index;
   iree_host_size_t set_layout_count;
   iree_hal_descriptor_set_layout_t* set_layouts[];
 } iree_hal_cuda_executable_layout_t;
@@ -28,6 +29,21 @@
   return (iree_hal_cuda_executable_layout_t*)base_value;
 }
 
+static void iree_hal_cuda_executable_layout_destroy(
+    iree_hal_executable_layout_t* base_executable_layout) {
+  iree_hal_cuda_executable_layout_t* executable_layout =
+      iree_hal_cuda_executable_layout_cast(base_executable_layout);
+  iree_allocator_t host_allocator = executable_layout->context->host_allocator;
+  IREE_TRACE_ZONE_BEGIN(z0);
+
+  for (iree_host_size_t i = 0; i < executable_layout->set_layout_count; ++i) {
+    iree_hal_descriptor_set_layout_release(executable_layout->set_layouts[i]);
+  }
+  iree_allocator_free(host_allocator, executable_layout);
+
+  IREE_TRACE_ZONE_END(z0);
+}
+
 iree_status_t iree_hal_cuda_executable_layout_create(
     iree_hal_cuda_context_wrapper_t* context, iree_host_size_t set_layout_count,
     iree_hal_descriptor_set_layout_t** set_layouts,
@@ -52,31 +68,20 @@
                                  &executable_layout->resource);
     executable_layout->context = context;
     executable_layout->set_layout_count = set_layout_count;
+    iree_host_size_t binding_number = 0;
     for (iree_host_size_t i = 0; i < set_layout_count; ++i) {
       executable_layout->set_layouts[i] = set_layouts[i];
       iree_hal_descriptor_set_layout_retain(set_layouts[i]);
+      binding_number +=
+          iree_hal_cuda_descriptor_set_layout_binding_count(set_layouts[i]);
     }
+    executable_layout->push_constant_base_index = binding_number;
     *out_executable_layout = (iree_hal_executable_layout_t*)executable_layout;
   }
   IREE_TRACE_ZONE_END(z0);
   return status;
 }
 
-static void iree_hal_cuda_executable_layout_destroy(
-    iree_hal_executable_layout_t* base_executable_layout) {
-  iree_hal_cuda_executable_layout_t* executable_layout =
-      iree_hal_cuda_executable_layout_cast(base_executable_layout);
-  iree_allocator_t host_allocator = executable_layout->context->host_allocator;
-  IREE_TRACE_ZONE_BEGIN(z0);
-
-  for (iree_host_size_t i = 0; i < executable_layout->set_layout_count; ++i) {
-    iree_hal_descriptor_set_layout_release(executable_layout->set_layouts[i]);
-  }
-  iree_allocator_free(host_allocator, executable_layout);
-
-  IREE_TRACE_ZONE_END(z0);
-}
-
 iree_host_size_t iree_hal_cuda_base_binding_index(
     iree_hal_executable_layout_t* base_executable_layout, uint32_t set) {
   iree_hal_cuda_executable_layout_t* executable_layout =
@@ -91,6 +96,13 @@
   return base_binding;
 }
 
+iree_host_size_t iree_hal_cuda_push_constant_index(
+    iree_hal_executable_layout_t* base_executable_layout) {
+  iree_hal_cuda_executable_layout_t* executable_layout =
+      iree_hal_cuda_executable_layout_cast(base_executable_layout);
+  return executable_layout->push_constant_base_index;
+}
+
 const iree_hal_executable_layout_vtable_t
     iree_hal_cuda_executable_layout_vtable = {
         .destroy = iree_hal_cuda_executable_layout_destroy,
diff --git a/iree/hal/cuda/executable_layout.h b/iree/hal/cuda/executable_layout.h
index 9b5a26e..adf29ea 100644
--- a/iree/hal/cuda/executable_layout.h
+++ b/iree/hal/cuda/executable_layout.h
@@ -26,6 +26,10 @@
 iree_host_size_t iree_hal_cuda_base_binding_index(
     iree_hal_executable_layout_t* executable_layout, uint32_t set);
 
+// Return the base index for push constant data.
+iree_host_size_t iree_hal_cuda_push_constant_index(
+    iree_hal_executable_layout_t* base_executable_layout);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif  // __cplusplus
diff --git a/iree/hal/cuda/graph_command_buffer.c b/iree/hal/cuda/graph_command_buffer.c
index f5667fa..19ddb12 100644
--- a/iree/hal/cuda/graph_command_buffer.c
+++ b/iree/hal/cuda/graph_command_buffer.c
@@ -37,6 +37,8 @@
 } iree_hal_cuda_graph_command_buffer_t;
 
 #define IREE_HAL_CUDA_MAX_BINDING_COUNT 64
+// Kernel arguments contains binding and push constants.
+#define IREE_HAL_CUDA_MAX_KERNEL_ARG 128
 
 extern const iree_hal_command_buffer_vtable_t
     iree_hal_cuda_graph_command_buffer_vtable;
@@ -62,8 +64,8 @@
                        "cuGraphCreate");
   iree_hal_cuda_graph_command_buffer_t* command_buffer = NULL;
   size_t total_size = sizeof(*command_buffer) +
-                      IREE_HAL_CUDA_MAX_BINDING_COUNT * sizeof(void*) +
-                      IREE_HAL_CUDA_MAX_BINDING_COUNT * sizeof(CUdeviceptr);
+                      IREE_HAL_CUDA_MAX_KERNEL_ARG * sizeof(void*) +
+                      IREE_HAL_CUDA_MAX_KERNEL_ARG * sizeof(CUdeviceptr);
   iree_status_t status = iree_allocator_malloc(
       context->host_allocator, total_size, (void**)&command_buffer);
   if (iree_status_is_ok(status)) {
@@ -79,8 +81,8 @@
 
     CUdeviceptr* device_ptrs =
         (CUdeviceptr*)(command_buffer->current_descriptor +
-                       IREE_HAL_CUDA_MAX_BINDING_COUNT);
-    for (size_t i = 0; i < IREE_HAL_CUDA_MAX_BINDING_COUNT; i++) {
+                       IREE_HAL_CUDA_MAX_KERNEL_ARG);
+    for (size_t i = 0; i < IREE_HAL_CUDA_MAX_KERNEL_ARG; i++) {
       command_buffer->current_descriptor[i] = &device_ptrs[i];
     }
 
@@ -330,8 +332,16 @@
     iree_hal_command_buffer_t* base_command_buffer,
     iree_hal_executable_layout_t* executable_layout, iree_host_size_t offset,
     const void* values, iree_host_size_t values_length) {
-  return iree_make_status(IREE_STATUS_UNIMPLEMENTED,
-                          "need cuda implementation");
+  iree_hal_cuda_graph_command_buffer_t* command_buffer =
+      iree_hal_cuda_graph_command_buffer_cast(base_command_buffer);
+  iree_host_size_t constant_base_index =
+      iree_hal_cuda_push_constant_index(executable_layout) +
+      offset / sizeof(int32_t);
+  for (iree_host_size_t i = 0; i < values_length / sizeof(int32_t); i++) {
+    *((uint32_t*)command_buffer->current_descriptor[i + constant_base_index]) =
+        ((uint32_t*)values)[i];
+  }
+  return iree_ok_status();
 }
 
 // Tie together the binding index and its index in |bindings| array.
diff --git a/iree/test/e2e/regression/BUILD b/iree/test/e2e/regression/BUILD
index 9aeca97..a034ab7 100644
--- a/iree/test/e2e/regression/BUILD
+++ b/iree/test/e2e/regression/BUILD
@@ -19,6 +19,14 @@
 )
 
 BACKEND_TESTS = [
+    "dynamic_abs.mlir",
+    "dynamic_add.mlir",
+    "dynamic_dot.mlir",
+    "dynamic_reduce_min.mlir",
+    "dynamic_torch_index_select_high_rank.mlir",
+    "dynamic_torch_index_select_negative.mlir",
+    "dynamic_torch_index_select_scalar.mlir",
+    "dynamic_torch_index_select_vector.mlir",
     "linalg_ops.mlir",
 ]
 
@@ -26,14 +34,6 @@
     name = "lit",
     srcs = enforce_glob(
         [
-            "dynamic_abs.mlir",
-            "dynamic_add.mlir",
-            "dynamic_dot.mlir",
-            "dynamic_reduce_min.mlir",
-            "dynamic_torch_index_select_high_rank.mlir",
-            "dynamic_torch_index_select_negative.mlir",
-            "dynamic_torch_index_select_scalar.mlir",
-            "dynamic_torch_index_select_vector.mlir",
             "globals.mlir",
             "scalar.mlir",
             "tensor_cast.mlir",
@@ -64,7 +64,7 @@
 )
 
 iree_check_single_backend_test_suite(
-    name = "check_regression_linalg_ops_dylib-llvm-aot",
+    name = "check_regression_dylib-llvm-aot",
     srcs = [
         "linalg_ext_ops.mlir",
         "lowering_config.mlir",
@@ -75,7 +75,7 @@
 )
 
 iree_check_single_backend_test_suite(
-    name = "check_regression_linalg_ops_vmvx",
+    name = "check_regression_vmvx",
     srcs = BACKEND_TESTS,
     compiler_flags = ["-iree-input-type=mhlo"],
     driver = "vmvx",
@@ -83,9 +83,25 @@
 )
 
 iree_check_single_backend_test_suite(
-    name = "check_regression_linalg_ops_vulkan-spirv",
+    name = "check_regression_vulkan-spirv",
     srcs = BACKEND_TESTS,
     compiler_flags = ["-iree-input-type=mhlo"],
     driver = "vulkan",
     target_backend = "vulkan-spirv",
 )
+
+iree_check_single_backend_test_suite(
+    name = "check_regression_cuda",
+    srcs = BACKEND_TESTS,
+    compiler_flags = ["-iree-input-type=mhlo"],
+    driver = "cuda",
+    tags = [
+        # CUDA cuInit fails with sanitizer on.
+        "noasan",
+        "nomsan",
+        "notsan",
+        "noubsan",
+        "requires-gpu-nvidia",
+    ],
+    target_backend = "cuda",
+)
diff --git a/iree/test/e2e/regression/CMakeLists.txt b/iree/test/e2e/regression/CMakeLists.txt
index 38d6648..f812bb3 100644
--- a/iree/test/e2e/regression/CMakeLists.txt
+++ b/iree/test/e2e/regression/CMakeLists.txt
@@ -14,14 +14,6 @@
   NAME
     lit
   SRCS
-    "dynamic_abs.mlir"
-    "dynamic_add.mlir"
-    "dynamic_dot.mlir"
-    "dynamic_reduce_min.mlir"
-    "dynamic_torch_index_select_high_rank.mlir"
-    "dynamic_torch_index_select_negative.mlir"
-    "dynamic_torch_index_select_scalar.mlir"
-    "dynamic_torch_index_select_vector.mlir"
     "globals.mlir"
     "scalar.mlir"
     "tensor_cast.mlir"
@@ -38,8 +30,16 @@
 
 iree_check_single_backend_test_suite(
   NAME
-    check_regression_linalg_ops_dylib-llvm-aot
+    check_regression_dylib-llvm-aot
   SRCS
+    "dynamic_abs.mlir"
+    "dynamic_add.mlir"
+    "dynamic_dot.mlir"
+    "dynamic_reduce_min.mlir"
+    "dynamic_torch_index_select_high_rank.mlir"
+    "dynamic_torch_index_select_negative.mlir"
+    "dynamic_torch_index_select_scalar.mlir"
+    "dynamic_torch_index_select_vector.mlir"
     "linalg_ext_ops.mlir"
     "linalg_ops.mlir"
     "lowering_config.mlir"
@@ -53,8 +53,16 @@
 
 iree_check_single_backend_test_suite(
   NAME
-    check_regression_linalg_ops_vmvx
+    check_regression_vmvx
   SRCS
+    "dynamic_abs.mlir"
+    "dynamic_add.mlir"
+    "dynamic_dot.mlir"
+    "dynamic_reduce_min.mlir"
+    "dynamic_torch_index_select_high_rank.mlir"
+    "dynamic_torch_index_select_negative.mlir"
+    "dynamic_torch_index_select_scalar.mlir"
+    "dynamic_torch_index_select_vector.mlir"
     "linalg_ops.mlir"
   TARGET_BACKEND
     "vmvx"
@@ -66,8 +74,16 @@
 
 iree_check_single_backend_test_suite(
   NAME
-    check_regression_linalg_ops_vulkan-spirv
+    check_regression_vulkan-spirv
   SRCS
+    "dynamic_abs.mlir"
+    "dynamic_add.mlir"
+    "dynamic_dot.mlir"
+    "dynamic_reduce_min.mlir"
+    "dynamic_torch_index_select_high_rank.mlir"
+    "dynamic_torch_index_select_negative.mlir"
+    "dynamic_torch_index_select_scalar.mlir"
+    "dynamic_torch_index_select_vector.mlir"
     "linalg_ops.mlir"
   TARGET_BACKEND
     "vulkan-spirv"
@@ -77,4 +93,31 @@
     "-iree-input-type=mhlo"
 )
 
+iree_check_single_backend_test_suite(
+  NAME
+    check_regression_cuda
+  SRCS
+    "dynamic_abs.mlir"
+    "dynamic_add.mlir"
+    "dynamic_dot.mlir"
+    "dynamic_reduce_min.mlir"
+    "dynamic_torch_index_select_high_rank.mlir"
+    "dynamic_torch_index_select_negative.mlir"
+    "dynamic_torch_index_select_scalar.mlir"
+    "dynamic_torch_index_select_vector.mlir"
+    "linalg_ops.mlir"
+  TARGET_BACKEND
+    "cuda"
+  DRIVER
+    "cuda"
+  COMPILER_FLAGS
+    "-iree-input-type=mhlo"
+  LABELS
+    "noasan"
+    "nomsan"
+    "notsan"
+    "noubsan"
+    "requires-gpu-nvidia"
+)
+
 ### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/test/e2e/regression/dynamic_abs.mlir b/iree/test/e2e/regression/dynamic_abs.mlir
index e1c6343..fa63d0d 100644
--- a/iree/test/e2e/regression/dynamic_abs.mlir
+++ b/iree/test/e2e/regression/dynamic_abs.mlir
@@ -1,12 +1,8 @@
-// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot %s | IreeFileCheck %s)
-// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vmvx %s | IreeFileCheck %s)
-// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv %s | IreeFileCheck %s)
-
-// CHECK-LABEL: EXEC @dynamic_tensor
-func @dynamic_tensor() -> tensor<?x?xf32> {
+func @dynamic_tensor() {
   %input = util.dynamic_shape_constant dense<[[-1.0, 2.0, -3.0], [4.0, -5.0, 6.0]]> : tensor<2x3xf32> -> tensor<?x?xf32>
   %res = "mhlo.abs"(%input) : (tensor<?x?xf32>) -> tensor<?x?xf32>
-  return %res : tensor<?x?xf32>
+  %dshape = util.do_not_optimize(%res) : tensor<?x?xf32>
+  %result = tensor.cast %dshape : tensor<?x?xf32> to tensor<2x3xf32>
+  check.expect_almost_eq_const(%result, dense<[[1.0, 2.0, 3.0],[4.0, 5.0, 6.0]]> : tensor<2x3xf32>) : tensor<2x3xf32>
+  return
 }
-
-// CHECK: 2x3xf32=[1 2 3][4 5 6]
diff --git a/iree/test/e2e/regression/dynamic_add.mlir b/iree/test/e2e/regression/dynamic_add.mlir
index b9bb49d..ab5ee2c 100644
--- a/iree/test/e2e/regression/dynamic_add.mlir
+++ b/iree/test/e2e/regression/dynamic_add.mlir
@@ -1,11 +1,9 @@
-// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot -function-input="2x4xf32=[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]" -function-input="2x4xf32=[[5.0,6.0,7.0,8.0],[-5.0,-6.0,-7.0,-8.0]]" %s | IreeFileCheck %s)
-// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vmvx -function-input="2x4xf32=[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]" -function-input="2x4xf32=[[5.0,6.0,7.0,8.0],[-5.0,-6.0,-7.0,-8.0]]" %s | IreeFileCheck %s)
-// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv -function-input="2x4xf32=[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]" -function-input="2x4xf32=[[5.0,6.0,7.0,8.0],[-5.0,-6.0,-7.0,-8.0]]" %s | IreeFileCheck %s)
-
-// CHECK: EXEC @main
-// CHECK: 2x4xf32=[6 8 10 12][-6 -8 -10 -12]
-
-func @main(%arg0: tensor<?x4xf32>, %arg1: tensor<?x4xf32>) -> tensor<?x4xf32> {
-  %0 = "mhlo.add"(%arg0, %arg1) : (tensor<?x4xf32>, tensor<?x4xf32>) -> tensor<?x4xf32>
-  return %0: tensor<?x4xf32>
+func @main() {
+  %lhs = util.dynamic_shape_constant  dense<[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]> : tensor<2x4xf32> -> tensor<?x4xf32>
+  %rhs = util.dynamic_shape_constant  dense<[[5.0,6.0,7.0,8.0],[-5.0,-6.0,-7.0,-8.0]]> : tensor<2x4xf32> -> tensor<?x4xf32>
+  %2 = "mhlo.add"(%lhs, %rhs) : (tensor<?x4xf32>, tensor<?x4xf32>) -> tensor<?x4xf32>
+  %3 = util.do_not_optimize(%2) : tensor<?x4xf32>
+  %result = tensor.cast %3 : tensor<?x4xf32> to tensor<2x4xf32>
+  check.expect_almost_eq_const(%result, dense<[[6.0, 8.0, 10.0, 12.0],[-6.0, -8.0, -10.0, -12.0]]> : tensor<2x4xf32>) : tensor<2x4xf32>
+  return
 }
diff --git a/iree/test/e2e/regression/dynamic_dot.mlir b/iree/test/e2e/regression/dynamic_dot.mlir
index e9d82d7..c60d2b2 100644
--- a/iree/test/e2e/regression/dynamic_dot.mlir
+++ b/iree/test/e2e/regression/dynamic_dot.mlir
@@ -1,9 +1,4 @@
-// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot %s | IreeFileCheck %s)
-// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vmvx %s | IreeFileCheck %s)
-// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv %s | IreeFileCheck %s)
-
-// CHECK-LABEL: EXEC @dynamic_dot
-func @dynamic_dot() -> tensor<?x?xf32> {
+func @dynamic_dot() {
   %lhs = util.dynamic_shape_constant dense<[
     [15.0, 14.0, 13.0],
     [12.0, 11.0, 10.0],
@@ -15,7 +10,13 @@
     [10.0, 09.0, 08.0, 07.0, 06.0],
     [05.0, 04.0, 03.0, 02.0, 01.0]]> : tensor<3x5xf32> -> tensor<?x?xf32>
   %res = "mhlo.dot"(%lhs, %rhs) : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
-  return %res : tensor<?x?xf32>
+  %dshape = util.do_not_optimize(%res) : tensor<?x?xf32>
+  %result = tensor.cast %dshape : tensor<?x?xf32> to tensor<5x5xf32>
+  check.expect_almost_eq_const(%result,
+    dense<[[430.0, 388.0, 346.0, 304.0, 262.0],
+           [340.0, 307.0, 274.0, 241.0, 208.0],
+           [250.0, 226.0, 202.0, 178.0, 154.0],
+           [160.0, 145.0, 130.0, 115.0, 100.0],
+           [70.0, 64.0, 58.0, 52.0, 46.0]]> : tensor<5x5xf32>) : tensor<5x5xf32>
+  return
 }
-
-// CHECK: 5x5xf32=[430 388 346 304 262][340 307 274 241 208][250 226 202 178 154][160 145 130 115 100][70 64 58 52 46]
diff --git a/iree/test/e2e/regression/dynamic_reduce_min.mlir b/iree/test/e2e/regression/dynamic_reduce_min.mlir
index cb5b0ae..b912c84 100644
--- a/iree/test/e2e/regression/dynamic_reduce_min.mlir
+++ b/iree/test/e2e/regression/dynamic_reduce_min.mlir
@@ -1,16 +1,12 @@
-// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot -function-input="2x4xf32=[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]" %s | IreeFileCheck %s)
-// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vmvx -function-input="2x4xf32=[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]" %s | IreeFileCheck %s)
-// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv -function-input="2x4xf32=[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]" %s | IreeFileCheck %s)
-
-func @reduce_min(%arg0: tensor<?x?xf32>) -> tensor<f32> {
+func @reduce_min() {
+  %input = util.dynamic_shape_constant
+    dense<[[1.0, 2.0, 3.0, 4.0],[-1.0 ,-2.0 ,-3.0 ,-4.0]]> : tensor<2x4xf32> -> tensor<?x?xf32>
   %0 = mhlo.constant dense<0x7F800000> : tensor<f32>
-  %1 = "mhlo.reduce"(%arg0, %0) ( {
+  %1 = "mhlo.reduce"(%input, %0) ( {
   ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
     %2 = mhlo.minimum %arg1, %arg2 : tensor<f32>
     "mhlo.return"(%2) : (tensor<f32>) -> ()
   }) {dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<?x?xf32>, tensor<f32>) -> tensor<f32>
-  return %1 : tensor<f32>
+  check.expect_almost_eq_const(%1, dense<-4.0> : tensor<f32>) : tensor<f32>
+  return
 }
-
-// CHECK: EXEC @reduce_min
-// CHECK: f32=-4
diff --git a/iree/test/e2e/regression/dynamic_torch_index_select_high_rank.mlir b/iree/test/e2e/regression/dynamic_torch_index_select_high_rank.mlir
index 570f8a5..b15c6a3 100644
--- a/iree/test/e2e/regression/dynamic_torch_index_select_high_rank.mlir
+++ b/iree/test/e2e/regression/dynamic_torch_index_select_high_rank.mlir
@@ -1,59 +1,31 @@
-// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot -function-input="2x2xi32=[6,7][8,9]" -function-input="2x2x2x2xi32=[[[0,1][1,0]][[0,0][1,1]]][[[1,1][0,0]][[0,1][1,0]]]" | IreeFileCheck %s)
-// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vmvx -function-input="2x2xi32=[6,7][8,9]" -function-input="2x2x2x2xi32=[[[0,1][1,0]][[0,0][1,1]]][[[1,1][0,0]][[0,1][1,0]]]" | IreeFileCheck %s)
-// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv -function-input="2x2xi32=[6,7][8,9]" -function-input="2x2x2x2xi32=[[[0,1][1,0]][[0,0][1,1]]][[[1,1][0,0]][[0,1][1,0]]]" | IreeFileCheck %s)
-
-// CHECK-LABEL: EXEC @torch_index_select1
-func @torch_index_select1(%arg0: tensor<?x?xi32>, %arg1: tensor<?x?x?x?xi32>) -> tensor<?x?x?x?xi32> {
-  %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = 1 : i64, dim = 1 : i64} : (tensor<?x?xi32>, tensor<?x?x?x?xi32>) -> tensor<?x?x?x?xi32>
-  return %0 : tensor<?x?x?x?xi32>
+func @torch_index_select1() {
+  %lhs = util.dynamic_shape_constant  dense<[[6,7],[8,9]]> : tensor<2x2xi32> -> tensor<?x?xi32>
+  %rhs = util.dynamic_shape_constant  dense<[[[[0,1],[1,0]],[[0,0],[1,1]]],[[[1,1],[0,0]],[[0,1],[1,0]]]]> : tensor<2x2x2x2xi32> -> tensor<?x?x?x?xi32>
+  %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = 1 : i64, dim = 1 : i64} : (tensor<?x?xi32>, tensor<?x?x?x?xi32>) -> tensor<?x?x?x?xi32>
+  %dshape = util.do_not_optimize(%0) : tensor<?x?x?x?xi32>
+  %result = tensor.cast %dshape : tensor<?x?x?x?xi32> to tensor<2x2x2x2xi32>
+  check.expect_eq_const(%result,
+    dense<[[[[6, 7],[7, 6]],
+            [[6, 6],[7, 7]]],
+           [[[9, 9],[8, 8]],
+            [[8, 9],[9, 8]]]]> : tensor<2x2x2x2xi32>) : tensor<2x2x2x2xi32>
+  return
 }
 
-// CHECK: 2x2x2x2xi32=[
-// CHECK-SAME:   [
-// CHECK-SAME:     [6 7][7 6]
-// CHECK-SAME:   ][
-// CHECK-SAME:     [6 6][7 7]
-// CHECK-SAME:   ]
-// CHECK-SAME: ][
-// CHECK-SAME:   [
-// CHECK-SAME:     [9 9][8 8]
-// CHECK-SAME:   ][
-// CHECK-SAME:     [8 9][9 8]
-// CHECK-SAME:   ]
-// CHECK-SAME: ]
-
-// CHECK-LABEL: EXEC @torch_index_select2
-func @torch_index_select2(%arg0: tensor<?x?xi32>, %arg1: tensor<?x?x?x?xi32>) -> tensor<?x?x?x?x?xi32> {
-  %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = 0 : i64, dim = 0 : i64} : (tensor<?x?xi32>, tensor<?x?x?x?xi32>) -> tensor<?x?x?x?x?xi32>
-  return %0 : tensor<?x?x?x?x?xi32>
+func @torch_index_select2() {
+  %lhs = util.dynamic_shape_constant  dense<[[6,7],[8,9]]> : tensor<2x2xi32> -> tensor<?x?xi32>
+  %rhs = util.dynamic_shape_constant  dense<[[[[0,1],[1,0]],[[0,0],[1,1]]],[[[1,1],[0,0]],[[0,1],[1,0]]]]> : tensor<2x2x2x2xi32> -> tensor<?x?x?x?xi32>
+  %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = 0 : i64, dim = 0 : i64} : (tensor<?x?xi32>, tensor<?x?x?x?xi32>) -> tensor<?x?x?x?x?xi32>
+  %dshape = util.do_not_optimize(%0) : tensor<?x?x?x?x?xi32>
+  %result = tensor.cast %dshape : tensor<?x?x?x?x?xi32> to tensor<2x2x2x2x2xi32>
+  check.expect_eq_const(%result,
+    dense<[[[[[6, 7],[8, 9]],
+             [[8, 9],[6, 7]]],
+            [[[6, 7],[6, 7]],
+             [[8, 9],[8, 9]]]],
+           [[[[8, 9],[8, 9]],
+             [[6, 7],[6, 7]]],
+            [[[6, 7],[8, 9]],
+             [[8, 9],[6, 7]]]]]> : tensor<2x2x2x2x2xi32>) : tensor<2x2x2x2x2xi32>
+  return
 }
-
-// CHECK: 2x2x2x2x2xi32=[
-// CHECK-SAME:   [
-// CHECK-SAME:     [
-// CHECK-SAME:       [6 7][8 9]
-// CHECK-SAME:     ][
-// CHECK-SAME:       [8 9][6 7]
-// CHECK-SAME:     ]
-// CHECK-SAME:   ][
-// CHECK-SAME:     [
-// CHECK-SAME:       [6 7][6 7]
-// CHECK-SAME:     ][
-// CHECK-SAME:       [8 9][8 9]
-// CHECK-SAME:     ]
-// CHECK-SAME:   ]
-// CHECK-SAME: ][
-// CHECK-SAME:   [
-// CHECK-SAME:     [
-// CHECK-SAME:       [8 9][8 9]
-// CHECK-SAME:     ][
-// CHECK-SAME:       [6 7][6 7]
-// CHECK-SAME:     ]
-// CHECK-SAME:   ][
-// CHECK-SAME:     [
-// CHECK-SAME:       [6 7][8 9]
-// CHECK-SAME:     ][
-// CHECK-SAME:       [8 9][6 7]
-// CHECK-SAME:     ]
-// CHECK-SAME:   ]
-// CHECK-SAME: ]
diff --git a/iree/test/e2e/regression/dynamic_torch_index_select_negative.mlir b/iree/test/e2e/regression/dynamic_torch_index_select_negative.mlir
index f23dae5..19f8367 100644
--- a/iree/test/e2e/regression/dynamic_torch_index_select_negative.mlir
+++ b/iree/test/e2e/regression/dynamic_torch_index_select_negative.mlir
@@ -1,15 +1,11 @@
-// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot -function-input="2x2x2xi32=[[100, 101] [110, 111]] [[200, 201] [210, 211]]" -function-input="2x2x2xi32=[[0, 1] [1, 0]] [[0, 0] [1, 1]]" | IreeFileCheck %s)
-// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vmvx -function-input="2x2x2xi32=[[100, 101] [110, 111]] [[200, 201] [210, 211]]" -function-input="2x2x2xi32=[[0, 1] [1, 0]] [[0, 0] [1, 1]]" | IreeFileCheck %s)
-// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv -function-input="2x2x2xi32=[[100, 101] [110, 111]] [[200, 201] [210, 211]]" -function-input="2x2x2xi32=[[0, 1] [1, 0]] [[0, 0] [1, 1]]" | IreeFileCheck %s)
-
-// CHECK-LABEL: EXEC @torch_index_select1
-func @torch_index_select1(%arg0: tensor<?x?x?xi32>, %arg1: tensor<?x?x?xi32>) -> tensor<?x?x?xi32> {
-  %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = -1 : i64, dim = -1 : i64} : (tensor<?x?x?xi32>, tensor<?x?x?xi32>) -> tensor<?x?x?xi32>
-  return %0 : tensor<?x?x?xi32>
+func @torch_index_select1() {
+  %lhs = util.dynamic_shape_constant  dense<[[[100, 101],[110, 111]],[[200, 201],[210, 211]]]> : tensor<2x2x2xi32> -> tensor<?x?x?xi32>
+  %rhs = util.dynamic_shape_constant  dense<[[[0, 1],[1, 0]],[[0, 0],[1, 1]]]> : tensor<2x2x2xi32> -> tensor<?x?x?xi32>
+  %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = -1 : i64, dim = -1 : i64} : (tensor<?x?x?xi32>, tensor<?x?x?xi32>) -> tensor<?x?x?xi32>
+  %dshape = util.do_not_optimize(%0) : tensor<?x?x?xi32>
+  %result = tensor.cast %dshape : tensor<?x?x?xi32> to tensor<2x2x2xi32>
+  check.expect_eq_const(%result,
+    dense<[[[100, 101],[111, 110]],
+           [[200, 200],[211, 211]]]> : tensor<2x2x2xi32>) : tensor<2x2x2xi32>
+  return
 }
-
-// CHECK: 2x2x2xi32=[
-// CHECK-SAME:   [100 101][111 110]
-// CHECK-SAME: ][
-// CHECK-SAME:   [200 200][211 211]
-// CHECK-SAME: ]
diff --git a/iree/test/e2e/regression/dynamic_torch_index_select_scalar.mlir b/iree/test/e2e/regression/dynamic_torch_index_select_scalar.mlir
index e97cf8f..f68962d 100644
--- a/iree/test/e2e/regression/dynamic_torch_index_select_scalar.mlir
+++ b/iree/test/e2e/regression/dynamic_torch_index_select_scalar.mlir
@@ -1,19 +1,35 @@
-// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot -function-input="5x1x5xi32=[[1,2,3,4,5]] [[6,7,8,9,10]] [[11,12,13,14,15]] [[16,17,18,19,20]] [[21,22,23,24,25]]" -function-input="i32=0" | IreeFileCheck %s)
-// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vmvx -function-input="5x1x5xi32=[[1,2,3,4,5]] [[6,7,8,9,10]] [[11,12,13,14,15]] [[16,17,18,19,20]] [[21,22,23,24,25]]" -function-input="i32=0" | IreeFileCheck %s)
-// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv -function-input="5x1x5xi32=[[1,2,3,4,5]] [[6,7,8,9,10]] [[11,12,13,14,15]] [[16,17,18,19,20]] [[21,22,23,24,25]]" -function-input="i32=0" | IreeFileCheck %s)
-
-// CHECK-LABEL: EXEC @torch_index_select1
-func @torch_index_select1(%arg0: tensor<?x?x?xi32>, %arg1: tensor<i32>) -> tensor<?x?xi32> {
-  %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = 0 : i64, dim = 0 : i64} : (tensor<?x?x?xi32>, tensor<i32>) -> tensor<?x?xi32>
-  return %0 : tensor<?x?xi32>
+func @torch_index_select1() {
+  %lhs = util.dynamic_shape_constant
+    dense<[[[1,2,3,4,5]],
+           [[6,7,8,9,10]],
+           [[11,12,13,14,15]],
+           [[16,17,18,19,20]],
+           [[21,22,23,24,25]]]> : tensor<5x1x5xi32> -> tensor<?x?x?xi32>
+  %rhs = util.unfoldable_constant dense<0> : tensor<i32>
+  %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = 0 : i64, dim = 0 : i64} : (tensor<?x?x?xi32>, tensor<i32>) -> tensor<?x?xi32>
+  %dshape = util.do_not_optimize(%0) : tensor<?x?xi32>
+  %result = tensor.cast %dshape : tensor<?x?xi32> to tensor<1x5xi32>
+  check.expect_eq_const(%result,
+    dense<[[1, 2, 3, 4, 5]]> : tensor<1x5xi32>) : tensor<1x5xi32>
+  return
 }
 
-// CHECK: 1x5xi32=[1 2 3 4 5]
-
-// CHECK-LABEL: EXEC @torch_index_select2
-func @torch_index_select2(%arg0: tensor<?x?x?xi32>, %arg1: tensor<i32>) -> tensor<?x?xi32> {
-  %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = 0 : i64, dim = 1 : i64} : (tensor<?x?x?xi32>, tensor<i32>) -> tensor<?x?xi32>
-  return %0 : tensor<?x?xi32>
+func @torch_index_select2() {
+   %lhs = util.dynamic_shape_constant
+    dense<[[[1,2,3,4,5]],
+           [[6,7,8,9,10]],
+           [[11,12,13,14,15]],
+           [[16,17,18,19,20]],
+           [[21,22,23,24,25]]]> : tensor<5x1x5xi32> -> tensor<?x?x?xi32>
+  %rhs = util.unfoldable_constant dense<0> : tensor<i32>
+  %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = 0 : i64, dim = 1 : i64} : (tensor<?x?x?xi32>, tensor<i32>) -> tensor<?x?xi32>
+  %dshape = util.do_not_optimize(%0) : tensor<?x?xi32>
+  %result = tensor.cast %dshape : tensor<?x?xi32> to tensor<5x5xi32>
+  check.expect_eq_const(%result,
+    dense<[[1, 2, 3, 4, 5],
+           [6, 7, 8, 9, 10],
+           [11, 12, 13, 14, 15],
+           [16, 17, 18, 19, 20],
+           [21, 22, 23, 24, 25]]> : tensor<5x5xi32>) : tensor<5x5xi32>
+  return
 }
-
-// CHECK: 5x5xi32=[1 2 3 4 5][6 7 8 9 10][11 12 13 14 15][16 17 18 19 20][21 22 23 24 25]
diff --git a/iree/test/e2e/regression/dynamic_torch_index_select_vector.mlir b/iree/test/e2e/regression/dynamic_torch_index_select_vector.mlir
index 64e13dd..1ef9007 100644
--- a/iree/test/e2e/regression/dynamic_torch_index_select_vector.mlir
+++ b/iree/test/e2e/regression/dynamic_torch_index_select_vector.mlir
@@ -1,29 +1,26 @@
-// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot -function-input="3x2x2xi32=[[1, 2] [3, 4]] [[5, 6] [7, 8]] [[9, 10] [11, 12]]" -function-input="2xi32=[0, 1]" | IreeFileCheck %s)
-// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vmvx -function-input="3x2x2xi32=[[1, 2] [3, 4]] [[5, 6] [7, 8]] [[9, 10] [11, 12]]" -function-input="2xi32=[0, 1]" | IreeFileCheck %s)
-// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv -function-input="3x2x2xi32=[[1, 2] [3, 4]] [[5, 6] [7, 8]] [[9, 10] [11, 12]]" -function-input="2xi32=[0, 1]" | IreeFileCheck %s)
-
-// CHECK-LABEL: EXEC @torch_index_select1
-func @torch_index_select1(%arg0: tensor<?x?x?xi32>, %arg1: tensor<?xi32>) -> tensor<?x?x?xi32> {
-  %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = 0 : i64, dim = 1 : i64} : (tensor<?x?x?xi32>, tensor<?xi32>) -> tensor<?x?x?xi32>
-  return %0 : tensor<?x?x?xi32>
+func @torch_index_select1() {
+  %lhs = util.dynamic_shape_constant
+    dense<[[[1, 2],[3, 4]],[[5, 6],[7, 8]],[[9, 10],[11, 12]]]> : tensor<3x2x2xi32> -> tensor<?x?x?xi32>
+  %rhs = util.dynamic_shape_constant dense<[0, 1]> : tensor<2xi32> -> tensor<?xi32>
+  %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = 0 : i64, dim = 1 : i64} : (tensor<?x?x?xi32>, tensor<?xi32>) -> tensor<?x?x?xi32>
+  %dshape = util.do_not_optimize(%0) : tensor<?x?x?xi32>
+  %result = tensor.cast %dshape : tensor<?x?x?xi32> to tensor<3x2x2xi32>
+  check.expect_eq_const(%result,
+    dense<[[[1, 2],[3, 4]],
+           [[5, 6],[7, 8]],
+           [[9, 10],[11, 12]]]> : tensor<3x2x2xi32>) : tensor<3x2x2xi32>
+  return
 }
 
-// CHECK: 3x2x2xi32=[
-// CHECK-SAME:   [1 2][3 4]
-// CHECK-SAME: ][
-// CHECK-SAME:   [5 6][7 8]
-// CHECK-SAME: ][
-// CHECK-SAME:   [9 10][11 12]
-// CHECK-SAME: ]
-
-// CHECK-LABEL: EXEC @torch_index_select2
-func @torch_index_select2(%arg0: tensor<?x?x?xi32>, %arg1: tensor<?xi32>) -> tensor<?x?x?xi32> {
-  %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = 0 : i64, dim = 0 : i64} : (tensor<?x?x?xi32>, tensor<?xi32>) -> tensor<?x?x?xi32>
-  return %0 : tensor<?x?x?xi32>
+func @torch_index_select2() {
+  %lhs = util.dynamic_shape_constant
+    dense<[[[1, 2],[3, 4]],[[5, 6],[7, 8]],[[9, 10],[11, 12]]]> : tensor<3x2x2xi32> -> tensor<?x?x?xi32>
+  %rhs = util.dynamic_shape_constant dense<[0, 1]> : tensor<2xi32> -> tensor<?xi32>
+  %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = 0 : i64, dim = 0 : i64} : (tensor<?x?x?xi32>, tensor<?xi32>) -> tensor<?x?x?xi32>
+  %dshape = util.do_not_optimize(%0) : tensor<?x?x?xi32>
+  %result = tensor.cast %dshape : tensor<?x?x?xi32> to tensor<2x2x2xi32>
+  check.expect_eq_const(%result,
+    dense<[[[1, 2],[3, 4]],
+           [[5, 6],[7, 8]]]> : tensor<2x2x2xi32>) : tensor<2x2x2xi32>
+  return
 }
-
-// CHECK: 2x2x2xi32=[
-// CHECK-SAME:   [1 2][3 4]
-// CHECK-SAME: ][
-// CHECK-SAME:   [5 6][7 8]
-// CHECK-SAME: ]
diff --git a/iree/test/e2e/xla_ops/BUILD b/iree/test/e2e/xla_ops/BUILD
index b9175a1..9a4e62b 100644
--- a/iree/test/e2e/xla_ops/BUILD
+++ b/iree/test/e2e/xla_ops/BUILD
@@ -62,8 +62,10 @@
             "remainder.mlir",
             "reshape.mlir",
             "reverse.mlir",
+            "rng_uniform.mlir",
             "rsqrt.mlir",
             "scatter.mlir",
+            "scatter_dynamic.mlir",
             "select.mlir",
             "sine.mlir",
             "slice.mlir",
@@ -78,7 +80,6 @@
         include = ["*.mlir"],
         exclude = [
             "round.mlir",
-            "scatter_dynamic.mlir",
         ],
     ),
     compiler_flags = ["-iree-input-type=mhlo"],
@@ -207,6 +208,7 @@
             "remainder.mlir",
             "reshape.mlir",
             "reverse.mlir",
+            "rng_uniform.mlir",
             "rsqrt.mlir",
             "scatter.mlir",
             "scatter_dynamic.mlir",
@@ -291,6 +293,7 @@
             "scatter.mlir",  # TODO(GH-6601): Enable the test.
             "scatter_dynamic.mlir",  # TODO(GH-6601): Enable the test.
             "sort.mlir",
+            "rng_uniform.mlir",  # TODO(GH-6995): Enable the test.
         ],
     ),
     compiler_flags = ["-iree-input-type=mhlo"],
@@ -340,6 +343,7 @@
             "remainder.mlir",
             "reshape.mlir",
             "reverse.mlir",
+            "rng_uniform.mlir",
             "rsqrt.mlir",
             "scatter.mlir",
             "scatter_dynamic.mlir",
diff --git a/iree/test/e2e/xla_ops/CMakeLists.txt b/iree/test/e2e/xla_ops/CMakeLists.txt
index 3a7007a..392dfa4 100644
--- a/iree/test/e2e/xla_ops/CMakeLists.txt
+++ b/iree/test/e2e/xla_ops/CMakeLists.txt
@@ -52,8 +52,10 @@
     "remainder.mlir"
     "reshape.mlir"
     "reverse.mlir"
+    "rng_uniform.mlir"
     "rsqrt.mlir"
     "scatter.mlir"
+    "scatter_dynamic.mlir"
     "select.mlir"
     "sine.mlir"
     "slice.mlir"
@@ -120,6 +122,7 @@
     "remainder.mlir"
     "reshape.mlir"
     "reverse.mlir"
+    "rng_uniform.mlir"
     "rsqrt.mlir"
     "scatter.mlir"
     "scatter_dynamic.mlir"
@@ -241,6 +244,7 @@
     "remainder.mlir"
     "reshape.mlir"
     "reverse.mlir"
+    "rng_uniform.mlir"
     "rsqrt.mlir"
     "scatter.mlir"
     "scatter_dynamic.mlir"
diff --git a/iree/test/e2e/xla_ops/rng_uniform.mlir b/iree/test/e2e/xla_ops/rng_uniform.mlir
new file mode 100644
index 0000000..7e496db
--- /dev/null
+++ b/iree/test/e2e/xla_ops/rng_uniform.mlir
@@ -0,0 +1,34 @@
+// Note that they are stateless random generators, so they have fixed results.
+func @rng_uniform_1d() {
+    %min = util.unfoldable_constant dense<-10.0> : tensor<f32>
+    %max = util.unfoldable_constant dense<10.0> : tensor<f32>
+    %shape = util.unfoldable_constant dense<[10]>  : tensor<1xi32>
+    %res = "mhlo.rng_uniform"(%min, %max, %shape) : (tensor<f32>, tensor<f32>, tensor<1xi32>) -> tensor<10xf32>
+    check.expect_almost_eq_const(%res, dense<[
+        -9.99994, -4.8613, 0.277344, 5.41599, -9.44537, -4.30673, 0.831918, 5.97056, -8.8908, -3.75215
+        ]> : tensor<10xf32>) : tensor<10xf32>
+    return
+}
+
+func @rng_uniform_2d() {
+    %min = util.unfoldable_constant dense<-10.0> : tensor<f32>
+    %max = util.unfoldable_constant dense<10.0> : tensor<f32>
+    %shape = util.unfoldable_constant dense<[3, 3]>  : tensor<2xi32>
+    %res = "mhlo.rng_uniform"(%min, %max, %shape) : (tensor<f32>, tensor<f32>, tensor<2xi32>) -> tensor<3x3xf32>
+    check.expect_almost_eq_const(%res, dense<[
+        [6.55154, -8.30982, -3.17117],
+        [1.75741, 6.89606, -7.9653],
+        [-3.03671, 2.10193, 7.24057]]> : tensor<3x3xf32>) : tensor<3x3xf32>
+    return
+}
+
+func @rng_uniform_3d() {
+    %min = util.unfoldable_constant dense<-10.0> : tensor<f32>
+    %max = util.unfoldable_constant dense<10.0> : tensor<f32>
+    %shape = util.unfoldable_constant dense<[2, 2, 2]>  : tensor<3xi32>
+    %res = "mhlo.rng_uniform"(%min, %max, %shape) : (tensor<f32>, tensor<f32>, tensor<3xi32>) -> tensor<2x2x2xf32>
+    check.expect_almost_eq_const(%res, dense<[
+        [[3.04814, 8.18679], [-1.74598, 3.39266]],
+        [[-6.91349, -1.77484], [8.29239, -6.56897]]]> : tensor<2x2x2xf32>) : tensor<2x2x2xf32>
+    return
+}