diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureTensorLayouts.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureTensorLayouts.cpp
index 98c0f5f..83768db 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureTensorLayouts.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureTensorLayouts.cpp
@@ -29,7 +29,6 @@
 #define GEN_PASS_DEF_LLVMGPUCONFIGURETENSORLAYOUTSPASS
 #include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"
 
-using IREE::GPU::MMASingleSubgroupLayout;
 using IREE::VectorExt::NestedLayoutAttr;
 using IREE::VectorExt::ToLayoutOp;
 using IREE::VectorExt::VectorLayoutInterface;
@@ -80,123 +79,81 @@
   return getSubgroupNCount(config).value();
 }
 
-struct ContractionLayout {
-  VectorLayoutInterface lhs;
-  VectorLayoutInterface rhs;
-  VectorLayoutInterface acc;
-};
-
-// Get the layouts to use for the contraction given the intrinsic to use and
-// number of subgroups on the M and N dimension.
-//
-// The contraction is expected to have 3 operands: lhs, rhs and acc of the
-// contraction and a single accumulator.
-static FailureOr<ContractionLayout>
-getContractionLayout(IREE::Codegen::InnerTileDescAttrInterface intrinsic,
-                     int64_t subgroupMCount, int64_t subgroupNCount,
-                     ArrayRef<int64_t> bounds,
-                     ArrayRef<AffineMap> contractIndexingMaps) {
-  auto mmaIntrinsic = dyn_cast<IREE::GPU::MmaInterfaceAttr>(intrinsic);
-  if (!mmaIntrinsic) {
-    return failure();
-  }
-
-  int64_t rank = bounds.size();
-  FailureOr<VectorContractOpInfo> maybeOpInfo =
-      VectorContractOpInfo::inferFromIndexingMaps(contractIndexingMaps);
-  if (failed(maybeOpInfo)) {
-    return failure();
-  }
-  VectorContractOpInfo opInfo = maybeOpInfo.value();
-  // Get the inner dimensions.
-  int64_t innerMDim = opInfo.getMDims().back();
-  int64_t innerNDim = opInfo.getNDims().back();
-  int64_t innerKDim = opInfo.getKDims().back();
-  // Find the number Of subgroups being used from subgroupMCount/subgroupNCount.
-  // Just assign them to the inner-most dimension for now.
-  // TODO: Use subgroup_basis to get this instead and allow distributing
-  // subgroups on multiple dimensions.
-  SmallVector<int64_t> subgroupCounts(rank, 1);
-  SmallVector<int64_t> subgroupStrides(rank, 0);
-  subgroupCounts[innerMDim] = subgroupMCount;
-  subgroupCounts[innerNDim] = subgroupNCount;
-  // Distribute on M and then N.
-  // TODO: Use subgroup_basis to get the strides.
-  subgroupStrides[innerMDim] = 1;
-  subgroupStrides[innerNDim] = subgroupMCount;
-
-  // Since these MMA intrinsics have a given tile size for each subgroup, we can
-  // calculate the batch dimensions without looking at the subgroup layout.
-  SmallVector<int64_t> subgroupSize(rank, 1);
-  auto [mSize, nSize, kSize] = mmaIntrinsic.getMNKShape();
-  subgroupSize[innerMDim] = mSize;
-  subgroupSize[innerNDim] = nSize;
-  subgroupSize[innerKDim] = kSize;
-
-  SmallVector<int64_t> batchCounts(rank);
-  for (auto [batchCount, subgroupCount, subgroupSize, bound] :
-       llvm::zip_equal(batchCounts, subgroupCounts, subgroupSize, bounds)) {
-    int64_t workgroupDimSize = subgroupCount * subgroupSize;
-    batchCount = llvm::divideCeil(bound, workgroupDimSize);
-  }
-
-  // MMA intrinsics can be weird and usually don't have a single subgroup
-  // iteration space, so we need to find their value subgroup iteration space
-  // indvidually.
-  auto getFragmentLayout = [&](IREE::GPU::MMAFragment fragment,
-                               int64_t outerDim, int64_t innerDim,
-                               AffineMap map) -> VectorLayoutInterface {
-    // Note that the struct MMASingleSubgroupLayout contains the partial layout
-    // for the canonical (M, K) x (K, N) -> (M, N) matmul form. We treat the
-    // concrete nested layout as the layout for the innermost M, N, K
-    // dimensions.
-    SmallVector<int64_t> outerCounts(rank, 1);
-    SmallVector<int64_t> elementCounts(rank, 1);
-    SmallVector<int64_t> threadCounts(rank, 1);
-    SmallVector<int64_t> threadStrides(rank, 0);
-
-    MMASingleSubgroupLayout subgroupLayout =
-        IREE::GPU::getSingleSubgroupLayout(intrinsic, fragment);
-    outerCounts[outerDim] = subgroupLayout.outer[0];
-    outerCounts[innerDim] = subgroupLayout.outer[1];
-    threadCounts[outerDim] = subgroupLayout.thread[0];
-    threadCounts[innerDim] = subgroupLayout.thread[1];
-    threadStrides[outerDim] = subgroupLayout.tstrides[0];
-    threadStrides[innerDim] = subgroupLayout.tstrides[1];
-    elementCounts[outerDim] = subgroupLayout.element[0];
-    elementCounts[innerDim] = subgroupLayout.element[1];
-    // Get the fragment layout for the entire iteration space and then project
-    // it. This is significantly easier than trying to create a layout for each
-    // fragment itself.
-    auto fragmentSpaceLayout = NestedLayoutAttr::get(
-        map.getContext(), subgroupCounts, batchCounts, outerCounts,
-        threadCounts, elementCounts, subgroupStrides, threadStrides);
-    return fragmentSpaceLayout.apply(map);
-  };
-
-  VectorLayoutInterface lhs =
-      getFragmentLayout(IREE::GPU::MMAFragment::Lhs, innerMDim, innerKDim,
-                        contractIndexingMaps[0]);
-  VectorLayoutInterface rhs =
-      getFragmentLayout(IREE::GPU::MMAFragment::Rhs, innerKDim, innerNDim,
-                        contractIndexingMaps[1]);
-  VectorLayoutInterface acc =
-      getFragmentLayout(IREE::GPU::MMAFragment::Acc, innerMDim, innerNDim,
-                        contractIndexingMaps[2]);
-
-  return ContractionLayout{lhs, rhs, acc};
+/// Gets a unit vector of the given rank, but fills in the given dimensions
+/// from the 2 element array |counts|. |dim0| is the position in the returned
+/// vector to put the first element of |counts|, and |dim1| is the position to
+/// put the second element. For example,
+///
+/// rank = 3, counts = [5, 7], dim0 = 2, dim1 = 1
+/// returns [1, 5, 7]
+static SmallVector<int64_t> getUnitOfRankWithDims(int64_t rank,
+                                                  ArrayRef<int64_t> counts,
+                                                  int64_t dim0, int64_t dim1) {
+  assert(counts.size() == 2 &&
+         "Unexpected non-rank 2 single subgroup dimension counts");
+  SmallVector<int64_t> res(rank, 1);
+  res[dim0] = counts[0];
+  res[dim1] = counts[1];
+  return res;
 }
 
-SmallVector<int64_t> getIterationSpaceBounds(linalg::LinalgOp linalgOp) {
-  SmallVector<int64_t> bounds = linalgOp.getStaticLoopRanges();
-  std::optional<VectorizationTileSizes> sizes =
-      inferSizesFromIR(linalgOp, std::nullopt);
-  // Even though the opShape could be dynamic, we could potentially
-  // infer the vector shape
-  if (sizes.has_value()) {
-    bounds = sizes.value().vectorSizes;
-  }
-  return bounds;
+/// Constructs the nested layout given the layout for a single subgroup and the
+/// subgroup/batch counts and orders, as well as the dimensions along which to
+/// distribute the intrinsic's layout.
+///
+/// |outerDim| and |innerDim| refer to which dimensions are the outermost and
+/// innermost for a canonical MK_KN_MN matrix multiply, for a particular
+/// fragment. For example, for the B matrix of an MK_NK_MN matrix multiply,
+/// we would have:
+///   outerDim = 1 for the K dim
+///   innerDim = 0 for the N dim
+///
+/// For something like MK_NKN_MN with multiple N dims, it would typically be:
+///   outerDim = 1 for K
+///   innerDim = 2 for the second N dim
+///
+/// Importantly these two dimensions always refer to the actual dimension
+/// positions in the undistributed vector. For each fragment, this means:
+///   A: [outerDim, innerDim] = [innerMostMDim, innerMostKDim]
+///   B: [outerDim, innerDim] = [innerMostKDim, innerMostNDim]
+///   C: [outerDim, innerDim] = [innerMostMDim, innerMostNDim]
+///
+/// And here inner most is referential to the iteration order, not the order
+/// they appear per fragment (because there is no relationship between the
+/// dimension order of M in A and in C, for example).
+static NestedLayoutAttr createNestedLayout(
+    MLIRContext *context, int64_t rank, int64_t outerDim, int64_t innerDim,
+    ArrayRef<int64_t> subgroupSizes, ArrayRef<int64_t> subgroupStrides,
+    ArrayRef<int64_t> batchCount, IREE::GPU::MMASingleSubgroupLayout counts) {
+
+  LLVM_DEBUG({
+    llvm::dbgs() << "Creating Nested Layout for::" << "\n    outerDim = "
+                 << outerDim << "\n    innerDim = " << innerDim
+                 << "\n    subgroupSizes: " << llvm::interleaved(subgroupSizes)
+                 << "\n    subgroupStrides: "
+                 << llvm::interleaved(subgroupStrides)
+                 << "\n    batchCount: " << llvm::interleaved(batchCount)
+                 << "\n    counts.outer: " << llvm::interleaved(counts.outer)
+                 << "\n    counts.thread: " << llvm::interleaved(counts.thread)
+                 << "\n    counts.element: "
+                 << llvm::interleaved(counts.element)
+                 << "\n    counts.tstrides: "
+                 << llvm::interleaved(counts.tstrides) << "\n";
+  });
+
+  SmallVector<int64_t> outerCount =
+      getUnitOfRankWithDims(rank, counts.outer, outerDim, innerDim);
+  SmallVector<int64_t> threadCount =
+      getUnitOfRankWithDims(rank, counts.thread, outerDim, innerDim);
+  SmallVector<int64_t> threadStrides =
+      getUnitOfRankWithDims(rank, counts.tstrides, outerDim, innerDim);
+  SmallVector<int64_t> elementCount =
+      getUnitOfRankWithDims(rank, counts.element, outerDim, innerDim);
+
+  auto layoutAttr = NestedLayoutAttr::get(context, subgroupSizes, batchCount,
+                                          outerCount, threadCount, elementCount,
+                                          subgroupStrides, threadStrides);
+  return layoutAttr;
 }
 
 static LogicalResult setContractionAnchor(IREE::GPU::MMAScheduleAttr schedule,
@@ -212,10 +169,12 @@
   assert(linalg::isaContractionOpInterface(contract) &&
          "cannot set contraction anchor on non contraction op");
 
-  SmallVector<int64_t> bounds = getIterationSpaceBounds(contract);
-  auto layouts = getContractionLayout(
-      schedule.getIntrinsic(), schedule.getSubgroupMCount(),
-      schedule.getSubgroupNCount(), bounds, contract.getIndexingMapsArray());
+  FailureOr<VectorContractOpInfo> opInfo =
+      VectorContractOpInfo::inferFromIndexingMaps(
+          contract.getIndexingMapsArray());
+  assert(succeeded(opInfo) && "contraction should have been inferred");
+
+  auto layouts = getContractionLayout(schedule, opInfo.value(), contract);
   if (failed(layouts)) {
     return contract->emitError("cannot get concrete layout for contraction");
   }
@@ -297,10 +256,15 @@
     map = projectDims(map, filterDims, /*compressDimsFlag=*/false);
   }
 
-  SmallVector<int64_t> bounds = getIterationSpaceBounds(conv);
-  auto layouts = getContractionLayout(
-      schedule.getIntrinsic(), schedule.getSubgroupMCount(),
-      schedule.getSubgroupNCount(), bounds, maps);
+  FailureOr<VectorContractOpInfo> opInfo =
+      VectorContractOpInfo::inferFromIndexingMaps(maps);
+  assert(succeeded(opInfo) &&
+         "unit filter dim convolution should have been infered");
+
+  auto layouts = getContractionLayout(schedule, opInfo.value(), conv);
+  if (failed(layouts)) {
+    return conv->emitError("cannot get concrete layout for convolution");
+  }
 
   auto [aLayout, bLayout, cLayout] = *layouts;
   Location loc = conv.getLoc();
@@ -492,6 +456,26 @@
                               pvMatmul);
 }
 
+// Apply the permuted projection map to the layout.
+static IREE::VectorExt::VectorLayoutInterface
+getLayoutForMap(VectorLayoutInterface layout, AffineMap map) {
+  // Project out unusued dims in layout.
+  SmallVector<bool> projectedDims(layout.getRank(), false);
+  llvm::SmallBitVector unusedBits = getUnusedDimsBitVector(map);
+  for (int dim : unusedBits.set_bits()) {
+    projectedDims[dim] = true;
+  }
+  IREE::VectorExt::VectorLayoutInterface projectedLayout =
+      layout.project(projectedDims);
+
+  // Transpose dims in layout.
+  AffineMap permMap = compressUnusedDims(map);
+  SmallVector<int64_t> identity =
+      llvm::to_vector(llvm::seq<int64_t>(permMap.getNumDims()));
+  SmallVector<int64_t> perm = applyPermutationMap<int64_t>(permMap, identity);
+  return projectedLayout.permute(perm);
+}
+
 static LogicalResult setDerivedThreadConfigLayout(
     IREE::GPU::DerivedThreadConfigAttr config, linalg::LinalgOp linalgOp,
     ArrayRef<int64_t> workgroupSize, RewriterBase &rewriter) {
@@ -568,7 +552,7 @@
   rewriter.setInsertionPointAfter(linalgOp);
   for (OpResult result : linalgOp->getResults()) {
     VectorLayoutInterface resultLayout =
-        layout.apply(linalgOp.getIndexingMapMatchingResult(result));
+        getLayoutForMap(layout, linalgOp.getIndexingMapMatchingResult(result));
     auto toLayout = rewriter.create<ToLayoutOp>(loc, result, resultLayout);
     rewriter.replaceAllUsesExcept(result, toLayout, toLayout);
   }
@@ -691,7 +675,14 @@
   MLIRContext *context = config.getContext();
   Location loc = candidate.getLoc();
 
-  SmallVector<int64_t> bounds = getIterationSpaceBounds(candidate);
+  SmallVector<int64_t> bounds = candidate.getStaticLoopRanges();
+  std::optional<VectorizationTileSizes> sizes =
+      inferSizesFromIR(candidate, std::nullopt);
+  // Even though the opShape could be dynamic, we could potentially
+  // infer the vector shape
+  if (sizes.has_value()) {
+    bounds = sizes.value().vectorSizes;
+  }
 
   // Subgroup distribution layouts.
   SmallVector<int64_t> subgroupSizes, subgroupStrides;
@@ -732,7 +723,7 @@
   rewriter.setInsertionPoint(candidate);
   for (OpOperand &operand : candidate->getOpOperands()) {
     VectorLayoutInterface operandLayout =
-        layout.apply(candidate.getMatchingIndexingMap(&operand));
+        getLayoutForMap(layout, candidate.getMatchingIndexingMap(&operand));
     auto toLayout =
         rewriter.create<ToLayoutOp>(loc, operand.get(), operandLayout);
     // Set shared memory promotion if requested.
@@ -744,7 +735,7 @@
   rewriter.setInsertionPointAfter(candidate);
   for (OpResult result : candidate->getResults()) {
     VectorLayoutInterface resultLayout =
-        layout.apply(candidate.getIndexingMapMatchingResult(result));
+        getLayoutForMap(layout, candidate.getIndexingMapMatchingResult(result));
     auto toLayout = rewriter.create<ToLayoutOp>(loc, result, resultLayout);
     rewriter.replaceAllUsesExcept(result, toLayout, toLayout);
   }
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUUtils.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUUtils.cpp
index 6cc7285..d8b9fc4 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUUtils.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUUtils.cpp
@@ -415,4 +415,335 @@
   packAllocs(builder, funcOp, aliasGroups);
 }
 
+/// Gets a unit vector of the given rank, but fills in the given dimensions
+/// from the 2 element array |counts|. |dim0| is the position in the returned
+/// vector to put the first element of |counts|, and |dim1| is the position to
+/// put the second element. For example,
+///
+/// rank = 3, counts = [5, 7], dim0 = 2, dim1 = 1
+/// returns [1, 5, 7]
+static SmallVector<int64_t> getUnitOfRankWithDims(int64_t rank,
+                                                  ArrayRef<int64_t> counts,
+                                                  int64_t dim0, int64_t dim1) {
+  assert(counts.size() == 2 &&
+         "Unexpected non-rank 2 single subgroup dimension counts");
+  SmallVector<int64_t> res(rank, 1);
+  res[dim0] = counts[0];
+  res[dim1] = counts[1];
+  return res;
+}
+
+/// Constructs the nested layout given the layout for a single subgroup and the
+/// subgroup/batch counts and orders, as well as the dimensions along which to
+/// distribute the intrinsic's layout.
+///
+/// |outerDim| and |innerDim| refer to which dimensions are the outermost and
+/// innermost for a canonical MK_KN_MN matrix multiply, for a particular
+/// fragment. For example, for the B matrix of an MK_NK_MN matrix multiply,
+/// we would have:
+///   outerDim = 1 for the K dim
+///   innerDim = 0 for the N dim
+///
+/// For something like MK_NKN_MN with multiple N dims, it would typically be:
+///   outerDim = 1 for K
+///   innerDim = 2 for the second N dim
+///
+/// Importantly these two dimensions always refer to the actual dimension
+/// positions in the undistributed vector. For each fragment, this means:
+///   A: [outerDim, innerDim] = [innerMostMDim, innerMostKDim]
+///   B: [outerDim, innerDim] = [innerMostKDim, innerMostNDim]
+///   C: [outerDim, innerDim] = [innerMostMDim, innerMostNDim]
+///
+/// And here inner most is referential to the iteration order, not the order
+/// they appear per fragment (because there is no relationship between the
+/// dimension order of M in A and in C, for example).
+static NestedLayoutAttr createNestedLayout(
+    MLIRContext *context, int64_t rank, int64_t outerDim, int64_t innerDim,
+    ArrayRef<int64_t> subgroupSizes, ArrayRef<int64_t> subgroupStrides,
+    ArrayRef<int64_t> batchCount, IREE::GPU::MMASingleSubgroupLayout counts) {
+
+  LLVM_DEBUG({
+    llvm::dbgs() << "Creating Nested Layout for::";
+    llvm::dbgs() << "\n    outerDim = " << outerDim;
+    llvm::dbgs() << "\n    innerDim = " << innerDim;
+    llvm::dbgs() << "\n    subgroupSizes: ";
+    llvm::interleaveComma(subgroupSizes, llvm::dbgs());
+    llvm::dbgs() << "\n    subgroupStrides: ";
+    llvm::interleaveComma(subgroupStrides, llvm::dbgs());
+    llvm::dbgs() << "\n    batchCount: ";
+    llvm::interleaveComma(batchCount, llvm::dbgs());
+    llvm::dbgs() << "\n    counts.outer: ";
+    llvm::interleaveComma(counts.outer, llvm::dbgs());
+    llvm::dbgs() << "\n    counts.thread: ";
+    llvm::interleaveComma(counts.thread, llvm::dbgs());
+    llvm::dbgs() << "\n    counts.element: ";
+    llvm::interleaveComma(counts.element, llvm::dbgs());
+    llvm::dbgs() << "\n    counts.tstrides: ";
+    llvm::interleaveComma(counts.tstrides, llvm::dbgs());
+    llvm::dbgs() << "\n";
+  });
+
+  SmallVector<int64_t> outerCount =
+      getUnitOfRankWithDims(rank, counts.outer, outerDim, innerDim);
+  SmallVector<int64_t> threadCount =
+      getUnitOfRankWithDims(rank, counts.thread, outerDim, innerDim);
+  SmallVector<int64_t> threadStrides =
+      getUnitOfRankWithDims(rank, counts.tstrides, outerDim, innerDim);
+  SmallVector<int64_t> elementCount =
+      getUnitOfRankWithDims(rank, counts.element, outerDim, innerDim);
+
+  auto layoutAttr = NestedLayoutAttr::get(context, subgroupSizes, batchCount,
+                                          outerCount, threadCount, elementCount,
+                                          subgroupStrides, threadStrides);
+  return layoutAttr;
+}
+
+template <typename ContractOpTy>
+static FailureOr<std::tuple<IREE::VectorExt::VectorLayoutInterface,
+                            IREE::VectorExt::VectorLayoutInterface,
+                            IREE::VectorExt::VectorLayoutInterface>>
+getContractionLayoutImpl(IREE::GPU::MMAScheduleAttr schedule,
+                         VectorContractOpInfo &opInfo,
+                         ContractOpTy contractOp) {
+  LLVM_DEBUG({
+    llvm::dbgs() << "Getting mma layouts for:\n" << contractOp << "\n";
+    llvm::dbgs() << "For schedule: " << schedule << "\n";
+  });
+
+  int64_t rank = contractOp.getIteratorTypesArray().size();
+  auto mmaAttr =
+      llvm::dyn_cast<IREE::GPU::MmaInterfaceAttr>(schedule.getIntrinsic());
+  if (!mmaAttr) {
+    return failure();
+  }
+
+  MLIRContext *context = schedule.getContext();
+
+  SmallVector<int64_t> bounds = contractOp.getStaticLoopRanges();
+  if (llvm::any_of(bounds,
+                   [](int64_t x) { return x == ShapedType::kDynamic; })) {
+    return failure();
+  }
+
+  if (!llvm::all_of(opInfo.getBatchDims(),
+                    [&bounds](int64_t dim) { return bounds[dim] == 1; })) {
+    LLVM_DEBUG({ llvm::dbgs() << "non-unit batch dimension\n"; });
+    return failure();
+  }
+
+  // Get the concrete nested layout for each matrix. Note that the struct
+  // MMASingleSubgroupLayout contains the partial layout for the
+  // canonical (M, K) x (K, N) -> (M, N) matmul form; while the specific
+  // contract op we are looking at right now may not be exactly in that form.
+  // So here we need to permute/transpose the canonical layout to match with
+  // the concrete contract op.
+
+  // Note that no matter how we permute/transpose the input contraction
+  // problem, the way we view the hardware warps remain the same--that is,
+  // from the hardware's perspective, a single warp has the same warp ID no
+  // matter what part of the contraction it works on. Similarly here, we are
+  // delinearizing the linearized GPU hardware lane ID into a n-D concatenated
+  // logical warp+thread using the subgroup/thread basis, so the subgroup
+  // basis should remain the same for all A/B/C matrix.
+
+  auto [intrinsicM, intrinsicN, intrinsicK] = mmaAttr.getMNKShape();
+
+  SmallVector<int64_t, 2> subgroupMBasis;
+  SmallVector<int64_t, 2> batchMSizes;
+  int64_t currMCount = schedule.getSubgroupMCount();
+
+  auto divideGreedily = [](int64_t availableSubgroups, int64_t dimSize,
+                           int64_t minDimSize) -> std::pair<int64_t, int64_t> {
+    int64_t dividableDim = dimSize / minDimSize;
+    int64_t subgroupsUsed = std::gcd(availableSubgroups, dividableDim);
+    dividableDim /= subgroupsUsed;
+    int64_t batchesUsed = dividableDim;
+    return {subgroupsUsed, batchesUsed};
+  };
+
+  // Greedily break up the M subgroup and batch counts along the "M" iteration
+  // bounds. We distribute as many residual subgroups as possible per M dim,
+  // and then divide the remaining along batch dims. The inner most M dim is
+  // always the one used for the intrinsic, meaning for a valid schedule, the
+  // computed batch counts and subgroup basis will satisfy totalMSize /
+  // intrinsicM = product(batchMSizes) * product(subgroupMBasis)
+  for (auto dim : opInfo.getMDims()) {
+    // Get the number of subgroups and batches used for this dimension based
+    // on the intrinsic size and the bound size.
+    int64_t subgroupsUsed, batchesUsed;
+    if (dim == opInfo.getMDims().back()) {
+      std::tie(subgroupsUsed, batchesUsed) =
+          divideGreedily(currMCount, bounds[dim], intrinsicM);
+    } else {
+      std::tie(subgroupsUsed, batchesUsed) =
+          divideGreedily(currMCount, bounds[dim], 1);
+    }
+    subgroupMBasis.push_back(subgroupsUsed);
+    batchMSizes.push_back(batchesUsed);
+    // Update available subgroup count.
+    currMCount /= subgroupsUsed;
+  }
+
+  SmallVector<int64_t, 2> subgroupNBasis;
+  SmallVector<int64_t, 2> batchNSizes;
+  int64_t currNCount = schedule.getSubgroupNCount();
+
+  // Do the same for N dims.
+  for (auto dim : opInfo.getNDims()) {
+    // Get the number of subgroups and batches used for this dimension based
+    // on the intrinsic size and the bound size.
+    int64_t subgroupsUsed, batchesUsed;
+    if (dim == opInfo.getNDims().back()) {
+      std::tie(subgroupsUsed, batchesUsed) =
+          divideGreedily(currNCount, bounds[dim], intrinsicN);
+    } else {
+      std::tie(subgroupsUsed, batchesUsed) =
+          divideGreedily(currNCount, bounds[dim], 1);
+    }
+    subgroupNBasis.push_back(subgroupsUsed);
+    batchNSizes.push_back(batchesUsed);
+    // Update available subgroup count.
+    currNCount /= subgroupsUsed;
+  }
+
+  SmallVector<int64_t> subgroupMStrides(subgroupMBasis.size());
+  SmallVector<int64_t> subgroupNStrides(subgroupNBasis.size());
+
+  auto mDimVec = opInfo.getMDims();
+  llvm::SmallDenseSet<int64_t> mDims(mDimVec.begin(), mDimVec.end());
+  auto nDimVec = opInfo.getNDims();
+  llvm::SmallDenseSet<int64_t> nDims(nDimVec.begin(), nDimVec.end());
+  // Because we currently require all batch dimensions to be unit, the
+  // subgroup basis can be constructed from the M and N bases. To keep things
+  // simple, the current heuristic is to distribute the loop dimensions from
+  // outer to inner.
+  int64_t currStride = 1;
+  int64_t currM = subgroupMStrides.size() - 1;
+  int64_t currN = subgroupNStrides.size() - 1;
+  for (int64_t dim : llvm::reverse(llvm::seq<int64_t>(rank))) {
+    if (mDims.contains(dim)) {
+      subgroupMStrides[currM] = currStride;
+      currStride *= subgroupMBasis[currM];
+      currM--;
+      continue;
+    }
+
+    if (nDims.contains(dim)) {
+      subgroupNStrides[currN] = currStride;
+      currStride *= subgroupNBasis[currN];
+      currN--;
+      continue;
+    }
+  }
+
+  // C matrix layout
+  auto [m, n] = opInfo.getResultMNIndex();
+  int64_t cRank = opInfo.getCRank();
+
+  // Get the M and N dims w.r.t. the dimensions of the C matrix. cMDims and
+  // cNDims are the M and N dimensions of the C matrix in the order they are
+  // iterated over in the contraction.
+  SmallVector<int64_t> cMDims = opInfo.outMDims;
+  SmallVector<int64_t> cNDims = opInfo.outNDims;
+  SmallVector<int64_t> cBatchSizes(cRank, 1);
+  SmallVector<int64_t> cSubgroupSizes(cRank, 1);
+  SmallVector<int64_t> cSubgroupStrides(cRank, 0);
+  for (auto [i, dim] : llvm::enumerate(cMDims)) {
+    cBatchSizes[dim] = batchMSizes[i];
+    cSubgroupSizes[dim] = subgroupMBasis[i];
+    cSubgroupStrides[dim] = subgroupMStrides[i];
+  }
+  for (auto [i, dim] : llvm::enumerate(cNDims)) {
+    cBatchSizes[dim] = batchNSizes[i];
+    cSubgroupSizes[dim] = subgroupNBasis[i];
+    cSubgroupStrides[dim] = subgroupNStrides[i];
+  }
+
+  IREE::VectorExt::NestedLayoutAttr cLayout = createNestedLayout(
+      context, cRank, m, n,
+      /*subgroupCount=*/cSubgroupSizes,
+      /*subgroupStrides=*/cSubgroupStrides,
+      /*batchCount=*/cBatchSizes,
+      getSingleSubgroupLayout(mmaAttr, IREE::GPU::MMAFragment::Acc));
+  LLVM_DEBUG({ llvm::dbgs() << "C layout: " << cLayout << "\n"; });
+
+  // A matrix layout
+  auto [afm, bfn] = opInfo.getOperandMNIndex();
+  auto [afk, bfk] = opInfo.getOperandKIndex();
+
+  int64_t aRank = opInfo.getARank();
+
+  SmallVector<int64_t> aMDims = opInfo.lhsMDims;
+  SmallVector<int64_t> aBatchSizes(aRank, 1);
+  SmallVector<int64_t> aSubgroupSizes(aRank, 1);
+  SmallVector<int64_t> aSubgroupStrides(aRank, 0);
+  for (auto [i, dim] : llvm::enumerate(aMDims)) {
+    aBatchSizes[dim] = batchMSizes[i];
+    aSubgroupSizes[dim] = subgroupMBasis[i];
+    aSubgroupStrides[dim] = subgroupMStrides[i];
+  }
+  for (auto [kDim, lhsKDim] :
+       llvm::zip_equal(opInfo.getKDims(), opInfo.lhsKDim)) {
+    aBatchSizes[lhsKDim] = bounds[kDim];
+  }
+  aBatchSizes[afk] = bounds[opInfo.getKDims().back()] / intrinsicK;
+
+  IREE::VectorExt::NestedLayoutAttr aLayout = createNestedLayout(
+      context, aRank, afm, afk,
+      /*subgroupCount=*/aSubgroupSizes,
+      /*subgroupStrides=*/aSubgroupStrides,
+      /*batchCount=*/aBatchSizes,
+      getSingleSubgroupLayout(mmaAttr, IREE::GPU::MMAFragment::Lhs));
+  LLVM_DEBUG({ llvm::dbgs() << "A layout: " << aLayout << "\n"; });
+
+  int64_t bRank = opInfo.getBRank();
+
+  SmallVector<int64_t> bNDims = opInfo.rhsNDims;
+  SmallVector<int64_t> bBatchSizes(bRank, 1);
+  SmallVector<int64_t> bSubgroupSizes(bRank, 1);
+  SmallVector<int64_t> bSubgroupStrides(bRank, 0);
+  for (auto [i, dim] : llvm::enumerate(bNDims)) {
+    bBatchSizes[dim] = batchNSizes[i];
+    bSubgroupSizes[dim] = subgroupNBasis[i];
+    bSubgroupStrides[dim] = subgroupNStrides[i];
+  }
+  for (auto [kDim, rhsKDim] :
+       llvm::zip_equal(opInfo.getKDims(), opInfo.rhsKDim)) {
+    bBatchSizes[rhsKDim] = bounds[kDim];
+  }
+  bBatchSizes[bfk] = bounds[opInfo.getKDims().back()] / intrinsicK;
+
+  IREE::VectorExt::NestedLayoutAttr bLayout = createNestedLayout(
+      context, bRank, bfk, bfn,
+      /*subgroupCount=*/bSubgroupSizes,
+      /*subgroupStrides=*/bSubgroupStrides,
+      /*batchCount=*/bBatchSizes,
+      getSingleSubgroupLayout(mmaAttr, IREE::GPU::MMAFragment::Rhs));
+  LLVM_DEBUG({ llvm::dbgs() << "B layout: " << bLayout << "\n"; });
+
+  std::tuple<VectorLayoutInterface, VectorLayoutInterface,
+             VectorLayoutInterface>
+      result = {aLayout, bLayout, cLayout};
+  return result;
+}
+
+/// Template specializations
+::mlir::FailureOr<::std::tuple<IREE::VectorExt::VectorLayoutInterface,
+                               IREE::VectorExt::VectorLayoutInterface,
+                               IREE::VectorExt::VectorLayoutInterface>>
+IREE::GPU::getContractionLayout(IREE::GPU::MMAScheduleAttr scheduleAttr,
+                                VectorContractOpInfo &opInfo,
+                                linalg::LinalgOp contractOp) {
+  return getContractionLayoutImpl(scheduleAttr, opInfo, contractOp);
+}
+
+::mlir::FailureOr<::std::tuple<IREE::VectorExt::VectorLayoutInterface,
+                               IREE::VectorExt::VectorLayoutInterface,
+                               IREE::VectorExt::VectorLayoutInterface>>
+IREE::GPU::getContractionLayout(IREE::GPU::MMAScheduleAttr scheduleAttr,
+                                VectorContractOpInfo &opInfo,
+                                vector::ContractionOp contractOp) {
+  return getContractionLayoutImpl(scheduleAttr, opInfo, contractOp);
+}
+
 } // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUUtils.h b/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUUtils.h
index 617b774..061e139 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUUtils.h
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUUtils.h
@@ -62,6 +62,24 @@
 void addBarrier(mlir::FunctionOpInterface funcOp, Operation *alloc,
                 ArrayRef<Operation *> aliasGroup, bool hasAsyncCopies = true);
 
+namespace IREE {
+namespace GPU {
+class MMAScheduleAttr;
+
+::llvm::FailureOr<::std::tuple<IREE::VectorExt::VectorLayoutInterface,
+                               IREE::VectorExt::VectorLayoutInterface,
+                               IREE::VectorExt::VectorLayoutInterface>>
+getContractionLayout(IREE::GPU::MMAScheduleAttr scheduleAttr,
+                     VectorContractOpInfo &opInfo, linalg::LinalgOp contractOp);
+
+::llvm::FailureOr<::std::tuple<IREE::VectorExt::VectorLayoutInterface,
+                               IREE::VectorExt::VectorLayoutInterface,
+                               IREE::VectorExt::VectorLayoutInterface>>
+getContractionLayout(IREE::GPU::MMAScheduleAttr scheduleAttr,
+                     VectorContractOpInfo &opInfo,
+                     vector::ContractionOp contractOp);
+} // namespace GPU
+} // namespace IREE
 } // namespace iree_compiler
 } // namespace mlir
 
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/configure_tensor_layout.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/configure_tensor_layout.mlir
index 074664c..a419f4a 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/configure_tensor_layout.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/configure_tensor_layout.mlir
@@ -207,6 +207,57 @@
                                               workgroup_size = [64, 1, 1]
                                               subgroup_size = 64>
 
+#maps = [
+  affine_map<(bm, bn, m, n, k) -> (bm, m, k)>,
+  affine_map<(bm, bn, m, n, k) -> (bn, n, k)>,
+  affine_map<(bm, bn, m, n, k) -> (bm, m, bn, n)>
+]
+
+#traits = {
+  indexing_maps = #maps,
+  iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction"],
+  lowering_config = #iree_gpu.lowering_config<{promote_operands = [0],
+                                              mma_kind = #iree_gpu.mma_layout<MFMA_F32_16x16x16_F16>,
+                                              subgroup_m_count = 2, subgroup_n_count = 2}>
+}
+
+func.func @packed_matmul_128x128x128(%lhs: tensor<8x16x16xf16>,
+                                     %rhs: tensor<8x16x16xf16>,
+                                     %init: tensor<8x16x8x16xf32>)
+                                     -> tensor<8x16x8x16xf32>
+                           attributes { translation_info = #translation } {
+  %out = linalg.generic #traits
+                        ins(%lhs, %rhs: tensor<8x16x16xf16>, tensor<8x16x16xf16>)
+                        outs(%init: tensor<8x16x8x16xf32>) {
+    ^bb0(%in: f16, %in_1: f16, %out: f32):
+      %ex   = arith.extf %in   : f16 to f32
+      %ex_1 = arith.extf %in_1 : f16 to f32
+      %mul  = arith.mulf %ex, %ex_1 : f32
+      %sum  = arith.addf %out, %mul : f32
+      linalg.yield %sum : f32
+  } -> tensor<8x16x8x16xf32>
+  return %out : tensor<8x16x8x16xf32>
+}
+
+
+// CHECK-DAG: #[[$NESTED:.+]] = #iree_vector_ext.nested_layout<subgroup_tile = [2, 1, 1], batch_tile = [4, 1, 1], outer_tile = [1, 1, 1], thread_tile = [1, 16, 4], element_tile = [1, 1, 4], subgroup_strides = [2, 0, 0], thread_strides = [0, 1, 16]>
+// CHECK-DAG: #[[$NESTED1:.+]] = #iree_vector_ext.nested_layout<subgroup_tile = [2, 1, 1], batch_tile = [4, 1, 1], outer_tile = [1, 1, 1], thread_tile = [1, 16, 4], element_tile = [1, 1, 4], subgroup_strides = [1, 0, 0], thread_strides = [0, 1, 16]>
+// CHECK-DAG: #[[$NESTED2:.+]] = #iree_vector_ext.nested_layout<subgroup_tile = [2, 1, 2, 1], batch_tile = [4, 1, 4, 1], outer_tile = [1, 1, 1, 1], thread_tile = [1, 4, 1, 16], element_tile = [1, 4, 1, 1], subgroup_strides = [2, 0, 1, 0], thread_strides = [0, 16, 0, 1]>
+// CHECK-LABEL: func.func @packed_matmul_128x128x128
+
+// CHECK-DAG: %[[LHS:.+]] = iree_vector_ext.to_layout %{{.*}} to layout(#[[$NESTED]]) {mma_kind = #iree_gpu.mma_layout<MFMA_F32_16x16x16_F16>, shared_memory_conversion}
+// CHECK-DAG: %[[RHS:.+]] = iree_vector_ext.to_layout %{{.*}} to layout(#[[$NESTED1]]) {mma_kind = #iree_gpu.mma_layout<MFMA_F32_16x16x16_F16>}
+// CHECK-DAG: %[[ACC:.+]] = iree_vector_ext.to_layout %{{.*}} to layout(#[[$NESTED2]]) {mma_kind = #iree_gpu.mma_layout<MFMA_F32_16x16x16_F16>}
+// CHECK: linalg.generic
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]]
+// CHECK-SAME: outs(%[[ACC]]
+
+// -----
+
+#translation = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute
+                                              workgroup_size = [64, 1, 1]
+                                              subgroup_size = 64>
+
 func.func @linalg_copy(%in : tensor<16x16x16xf16>) -> tensor<16x16x16xf16>
                       attributes { translation_info = #translation } {
   %empty = tensor.empty() : tensor<16x16x16xf16>
diff --git a/tests/external/iree-test-suites/test_suite_files/attention_and_matmul_spec_unet_fp16_mi308.mlir b/tests/external/iree-test-suites/test_suite_files/attention_and_matmul_spec_unet_fp16_mi308.mlir
index b35fc42..3d14b1d 100644
--- a/tests/external/iree-test-suites/test_suite_files/attention_and_matmul_spec_unet_fp16_mi308.mlir
+++ b/tests/external/iree-test-suites/test_suite_files/attention_and_matmul_spec_unet_fp16_mi308.mlir
@@ -27,7 +27,7 @@
   transform.iree.match.cast_compatible_type %in0 = tensor<?x?x?x?xf16> : !transform.any_value
 
   %config = transform.param.constant #iree_codegen.compilation_info<
-          lowering_config = #iree_gpu.lowering_config<{workgroup = [1, 1, 128, 0, 0, 0], reduction=[0, 0, 0, 0, 0, 64], promote_operands = [1, 2]}>,
+          lowering_config = #iree_gpu.lowering_config<{workgroup = [1, 1, 64, 0, 0, 0], reduction=[0, 0, 0, 0, 0, 64], promote_operands = [1, 2]}>,
           translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute
                                                             workgroup_size = [64, 4]
                                                             subgroup_size = 64 ,