diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
index 0a27437..19a7b3a 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
@@ -12,7 +12,6 @@
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h"
-#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
 #include "iree/compiler/Codegen/Utils/VectorOpUtils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/STLForwardCompat.h"
@@ -45,169 +44,34 @@
 #define GET_ATTRDEF_CLASSES
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp.inc"
 
-using LayoutDimension = mlir::iree_compiler::IREE::VectorExt::LayoutDimension;
-using LayoutDimensionAttr =
-    mlir::iree_compiler::IREE::VectorExt::LayoutDimensionAttr;
-using VectorLayoutInterface =
-    mlir::iree_compiler::IREE::VectorExt::VectorLayoutInterface;
-using PerDimLayoutAttr = mlir::iree_compiler::IREE::VectorExt::PerDimLayoutAttr;
-using LayoutAttr = mlir::iree_compiler::IREE::VectorExt::LayoutAttr;
-using NestedLayoutAttr = mlir::iree_compiler::IREE::VectorExt::NestedLayoutAttr;
-
 namespace mlir::iree_compiler::IREE::GPU {
 
-namespace {
-// Struct containing abstract MMA shape and type information.
-struct OpaqueMmaLayout {
-  int64_t mSize = 0;
-  int64_t nSize = 0;
-  int64_t kSize = 0;
-  Type aType;
-  Type bType;
-  Type cType;
-};
-
-// Struct containing concrete MMA shape, type, and layout information.
-struct ConcreteMmaLayout {
-  OpaqueMmaLayout base;
-  PerDimLayoutAttr aMLayout;
-  PerDimLayoutAttr aKLayout;
-  PerDimLayoutAttr bKLayout;
-  PerDimLayoutAttr bNLayout;
-  PerDimLayoutAttr cMLayout;
-  PerDimLayoutAttr cNLayout;
-};
-} // namespace
-
 //===----------------------------------------------------------------------===//
-// #iree_gpu.mma_vector_layout
+// MMA intrinsics semantics: shapes, layouts, operand element types.
 //===----------------------------------------------------------------------===//
 
-static PerDimLayoutAttr getBatchedPerDimLayoutAttr(LayoutDimensionAttr batchDim,
-                                                   PerDimLayoutAttr baseLayout,
-                                                   int64_t problemSize,
-                                                   int64_t fragmentDimSize) {
-  assert(problemSize % fragmentDimSize == 0 &&
-         "invalid layout fragment for problem size");
-
-  SmallVector<LayoutDimensionAttr, 3> dimAttrs(baseLayout.getLabels());
-  dimAttrs.insert(dimAttrs.begin(), batchDim);
-
-  SmallVector<int64_t, 3> shapes(baseLayout.getShapes());
-  shapes.insert(shapes.begin(), problemSize / fragmentDimSize);
-  auto layout =
-      PerDimLayoutAttr::get(baseLayout.getContext(), dimAttrs, shapes);
-  return layout;
+static int getBlockSize(MMAIntrinsic /*intrinsic*/) {
+  // Not supporting any block size other than 1 at the moment.
+  return 1;
 }
 
-// Get the batched layout attributes for the given fragment layouts, indexing
-// map, and problem shape. The canonical fragment map is used to compare against
-// the problem map |indexingMap|. For example, for mma fragment B (RHS):
-//
-// indexingMap = affine_map<(d0, d1, d2) -> (d1, d2) # Transposed B
-// fragmentMap = affine_map<(d0, d1, d2) -> (d2, d1)
-// problemShape = [32, 64]
-// fragmentSize = [16, 8]
-// fragmentLayouts = [kLayout, nLayout]
-//
-// Gives batched layout
-//
-// Dim0 Layout = [BATCHX, nLayoutLabels], [8, nLayoutShape]
-// Dim1 Layout = [BATCHY, kLayoutLabels], [2, kLayoutShape]
-static LayoutAttr
-getBatchedLayoutAttr(AffineMap indexingMap, AffineMap fragmentMap,
-                     ArrayRef<int64_t> problemShape,
-                     ArrayRef<int64_t> fragmentSize,
-                     ArrayRef<PerDimLayoutAttr> fragmentLayouts) {
-  // Current distribution to MFMA operations does not support batched
-  // contractions so that is reflected here.
-  assert(indexingMap.getNumResults() == 2 &&
-         "invalid indexing map to non-batched simple contraction");
-
-  LayoutDimensionAttr batchX = LayoutDimensionAttr::get(
-      indexingMap.getContext(), LayoutDimension::BATCHX);
-  LayoutDimensionAttr batchY = LayoutDimensionAttr::get(
-      indexingMap.getContext(), LayoutDimension::BATCHY);
-
-  SmallVector<PerDimLayoutAttr, 2> perDimAttrs;
-  for (auto [expr, batchType] :
-       llvm::zip_equal(indexingMap.getResults(),
-                       SmallVector<LayoutDimensionAttr, 2>{batchX, batchY})) {
-    auto maybeResultPosition = fragmentMap.getResultPosition(expr);
-    assert(maybeResultPosition && "fragment map and problem map mismatch");
-    int64_t idx = *maybeResultPosition;
-    perDimAttrs.push_back(getBatchedPerDimLayoutAttr(
-        batchType, fragmentLayouts[idx], problemShape[idx], fragmentSize[idx]));
-  }
-
-  return LayoutAttr::get(indexingMap.getContext(), perDimAttrs);
+static uint32_t getArchID(MMAIntrinsic intrinsic) {
+  return static_cast<int>(intrinsic) & 0xFF00;
 }
 
-static FailureOr<std::tuple<VectorLayoutInterface, VectorLayoutInterface,
-                            VectorLayoutInterface>>
-getContractionLayout(vector::ContractionOp contract, ConcreteMmaLayout layout) {
-  MLIRContext *context = contract.getContext();
-  FailureOr<linalg::ContractionDimensions> maybeContractionDims =
-      linalg::inferContractionDims(contract.getIndexingMapsArray());
-  if (failed(maybeContractionDims)) {
-    return failure();
-  }
-  auto contractionDims = *maybeContractionDims;
-  // TODO: Relax this condition to strictly alignment requirements.
-  if (contractionDims.k.size() != 1 || contractionDims.m.size() != 1 ||
-      contractionDims.n.size() != 1) {
-    return failure();
-  }
-  // TODO: Support batched contractions.
-  if (contractionDims.batch.size() > 0) {
-    return failure();
-  }
-  unsigned mDim = contractionDims.m[0];
-  unsigned nDim = contractionDims.n[0];
-  unsigned kDim = contractionDims.k[0];
-
-  SmallVector<int64_t> iterationBounds;
-  contract.getIterationBounds(iterationBounds);
-
-  int64_t problemMSize = iterationBounds[mDim];
-  int64_t problemNSize = iterationBounds[nDim];
-  int64_t problemKSize = iterationBounds[kDim];
-
-  int64_t mSize = layout.base.mSize;
-  int64_t nSize = layout.base.nSize;
-  int64_t kSize = layout.base.kSize;
-
-  // The problem size currently must be strictly aligned to the size of the mma.
-  // This is expected to succeed assuming the correct [masked] vector size was
-  // set at strategy configuration time (for this mma).
-  if (problemMSize % mSize != 0 || problemNSize % nSize ||
-      problemKSize % kSize) {
-    return failure();
-  }
-
-  LayoutAttr aLayout = getBatchedLayoutAttr(
-      contract.getIndexingMapsArray()[0],
-      AffineMap::getMultiDimMapWithTargets(3, {mDim, kDim}, context),
-      {problemMSize, problemKSize}, {mSize, kSize},
-      {layout.aMLayout, layout.aKLayout});
-  LayoutAttr bLayout = getBatchedLayoutAttr(
-      contract.getIndexingMapsArray()[1],
-      AffineMap::getMultiDimMapWithTargets(3, {kDim, nDim}, context),
-      {problemKSize, problemNSize}, {kSize, nSize},
-      {layout.bKLayout, layout.bNLayout});
-  LayoutAttr cLayout = getBatchedLayoutAttr(
-      contract.getIndexingMapsArray()[2],
-      AffineMap::getMultiDimMapWithTargets(3, {mDim, nDim}, context),
-      {problemMSize, problemNSize}, {mSize, nSize},
-      {layout.cMLayout, layout.cNLayout});
-
-  return std::make_tuple<VectorLayoutInterface, VectorLayoutInterface,
-                         VectorLayoutInterface>(aLayout, bLayout, cLayout);
+static bool is_AMD_MFMA(MMAIntrinsic intrinsic) {
+  return getArchID(intrinsic) >= 0x1000 && getArchID(intrinsic) <= 0x17FF;
 }
 
-//===----------------------------------------------------------------------===//
-// Layout Attribute Building Helpers
-//===----------------------------------------------------------------------===//
+static bool is_AMD_WMMA(MMAIntrinsic intrinsic) {
+  return getArchID(intrinsic) >= 0x1800 && getArchID(intrinsic) <= 0x1FFF;
+}
+
+static int64_t getIntrinsicSubgroupSize(MMAIntrinsic intrinsic) {
+  // Not using Wave64 at all at the moment, so the only place where the
+  // subgroup size is CDNA* architectures.
+  return is_AMD_MFMA(intrinsic) ? 64 : 32;
+}
 
 static std::tuple<Type, Type, Type> getABCElementTypes(MLIRContext *context,
                                                        MMAIntrinsic intrinsic) {
@@ -263,233 +127,6 @@
   return {};
 }
 
-template <typename MMAIntrinsicType>
-static OpaqueMmaLayout getOpaqueMMALayout(MLIRContext *context,
-                                          MMAIntrinsicType intrinsic) {
-  OpaqueMmaLayout o;
-  std::tie(o.aType, o.bType, o.cType) = getABCElementTypes(context, intrinsic);
-  auto lhs = getSingleSubgroupLayout(intrinsic, MMAFragment::Lhs);
-  auto rhs = getSingleSubgroupLayout(intrinsic, MMAFragment::Rhs);
-  o.mSize = lhs.outer[0] * lhs.thread[0] * lhs.element[0];
-  o.kSize = lhs.outer[1] * lhs.thread[1] * lhs.element[1];
-  o.nSize = rhs.outer[1] * rhs.thread[1] * rhs.element[1];
-  return o;
-}
-
-static std::tuple<PerDimLayoutAttr, PerDimLayoutAttr>
-getPerDimLayoutAttrs(MLIRContext *context, TileSwizzle swizzle) {
-  // Step 1: obtain the swizzled tile shape, but keeping track of the source
-  // dimension indices.
-  struct SrcIndexAndSwizzleDim {
-    size_t srcIndex;
-    TileSwizzle::Dim dim;
-  };
-  SmallVector<SrcIndexAndSwizzleDim> swizzledShape;
-  for (auto [i, e] : llvm::enumerate(swizzle.expandShape)) {
-    for (TileSwizzle::Dim d : e) {
-      swizzledShape.push_back(SrcIndexAndSwizzleDim{i, d});
-    }
-  }
-  applyPermutationToVector(swizzledShape, swizzle.permutation);
-
-  // Step 2: collect the appropriate labels to use for the swizzled dims.
-  LayoutDimension internalLabels[] = {LayoutDimension::VECTORZ,
-                                      LayoutDimension::VECTORY,
-                                      LayoutDimension::VECTORX};
-  LayoutDimension crossThreadLabels[] = {
-      LayoutDimension::LANEZ, LayoutDimension::LANEY, LayoutDimension::LANEX};
-  auto internalLabelIter = std::end(internalLabels);
-  auto crossThreadLabelIter = std::end(crossThreadLabels);
-  for (SrcIndexAndSwizzleDim d : swizzledShape) {
-    if (d.dim.kind == TileSwizzle::Dim::Kind::Internal) {
-      assert(internalLabelIter != std::begin(internalLabels));
-      --internalLabelIter;
-    } else if (d.dim.kind == TileSwizzle::Dim::Kind::CrossThread) {
-      assert(crossThreadLabelIter != std::begin(crossThreadLabels));
-      --crossThreadLabelIter;
-    } else {
-      assert(false && "unexpected dimension kind in intrinsic swizzle");
-    }
-  }
-
-  // Step 3: put together the result PerDimLayoutAttr'd for the two source dims.
-  SmallVector<LayoutDimensionAttr> labels[2];
-  SmallVector<int64_t> shape[2];
-  for (SrcIndexAndSwizzleDim d : swizzledShape) {
-    shape[d.srcIndex].push_back(d.dim.size);
-    auto &labelIterRef = (d.dim.kind == TileSwizzle::Dim::Kind::Internal)
-                             ? internalLabelIter
-                             : crossThreadLabelIter;
-    labels[d.srcIndex].push_back(LayoutDimensionAttr::get(
-        context, static_cast<LayoutDimension>(*labelIterRef++)));
-  }
-  return {PerDimLayoutAttr::get(context, labels[0], shape[0]),
-          PerDimLayoutAttr::get(context, labels[1], shape[1])};
-};
-
-static ConcreteMmaLayout getConcreteMMALayout(MLIRContext *context,
-                                              MMAIntrinsic intrinsic) {
-  auto opaque = getOpaqueMMALayout(context, intrinsic);
-  ConcreteMmaLayout concreteLayout;
-  concreteLayout.base = opaque;
-  auto lhsSwizzle = getIntrinsicSwizzle(intrinsic, MMAFragment::Lhs);
-  auto rhsSwizzle = getIntrinsicSwizzle(intrinsic, MMAFragment::Rhs);
-  auto accSwizzle = getIntrinsicSwizzle(intrinsic, MMAFragment::Acc);
-  std::tie(concreteLayout.aMLayout, concreteLayout.aKLayout) =
-      getPerDimLayoutAttrs(context, lhsSwizzle);
-  std::tie(concreteLayout.bNLayout, concreteLayout.bKLayout) =
-      getPerDimLayoutAttrs(context, rhsSwizzle);
-  std::tie(concreteLayout.cMLayout, concreteLayout.cNLayout) =
-      getPerDimLayoutAttrs(context, accSwizzle);
-  return concreteLayout;
-}
-
-//===----------------------------------------------------------------------===//
-// MmaInterface Attribute Helper Functions
-//===----------------------------------------------------------------------===//
-
-MMASingleSubgroupLayout getASingleSubgroupLayout(MmaInterfaceAttr mmaKind) {
-  if (auto mmaAttr = dyn_cast<MMAAttr>(mmaKind)) {
-    return mmaAttr.getASingleSubgroupLayout();
-  } else if (auto vmmaAttr = dyn_cast<VirtualMMAAttr>(mmaKind)) {
-    return vmmaAttr.getASingleSubgroupLayout();
-  } else {
-    assert(false && "unhandled MMA Interface type.");
-    return {};
-  }
-}
-
-MMASingleSubgroupLayout getBSingleSubgroupLayout(MmaInterfaceAttr mmaKind) {
-  if (auto mmaAttr = dyn_cast<MMAAttr>(mmaKind)) {
-    return mmaAttr.getBSingleSubgroupLayout();
-  } else if (auto vmmaAttr = dyn_cast<VirtualMMAAttr>(mmaKind)) {
-    return vmmaAttr.getBSingleSubgroupLayout();
-  } else {
-    assert(false && "unhandled MMA Interface type.");
-    return {};
-  }
-}
-
-MMASingleSubgroupLayout getCSingleSubgroupLayout(MmaInterfaceAttr mmaKind) {
-  if (auto mmaAttr = dyn_cast<MMAAttr>(mmaKind)) {
-    return mmaAttr.getCSingleSubgroupLayout();
-  } else if (auto vmmaAttr = dyn_cast<VirtualMMAAttr>(mmaKind)) {
-    return vmmaAttr.getCSingleSubgroupLayout();
-  } else {
-    assert(false && "unhandled MMA Interface type.");
-    return {};
-  }
-}
-
-//===----------------------------------------------------------------------===//
-// MFMA Attributes
-//===----------------------------------------------------------------------===//
-
-Attribute MMAAttr::parse(AsmParser &p, Type type) {
-  if (failed(p.parseLess()))
-    return {};
-
-  FailureOr<MMAIntrinsicAttr> mmaIntrinsic =
-      FieldParser<MMAIntrinsicAttr>::parse(p);
-  if (failed(mmaIntrinsic)) {
-    p.emitError(p.getCurrentLocation(), "failed to parse mfma type identifier");
-    return {};
-  }
-
-  if (failed(p.parseGreater()))
-    return {};
-
-  return get(p.getContext(), mmaIntrinsic->getValue());
-}
-
-void MMAAttr::print(AsmPrinter &p) const {
-  auto &os = p.getStream();
-  os << "<";
-  os << stringifyMMAIntrinsic(getIntrinsic().getValue());
-  os << ">";
-}
-
-MMAAttr MMAAttr::get(MLIRContext *context, MMAIntrinsic type) {
-  auto layout = getOpaqueMMALayout(context, type);
-  return Base::get(context, MMAIntrinsicAttr::get(context, type), layout.mSize,
-                   layout.nSize, layout.kSize, layout.aType, layout.bType,
-                   layout.cType);
-}
-
-std::tuple<Type, Type, Type> MMAAttr::getABCElementTypes() const {
-  return {getAType(), getBType(), getCType()};
-}
-
-std::tuple<int64_t, int64_t, int64_t> MMAAttr::getMNKShape() const {
-  return {getMSize(), getNSize(), getKSize()};
-}
-
-template <typename MMAIntrinsicType>
-static VectorType getVectorType(MLIRContext *context,
-                                MMAIntrinsicType intrinsic,
-                                MMAFragment fragment) {
-  auto o = getOpaqueMMALayout(context, intrinsic);
-  auto s = getSingleSubgroupLayout(intrinsic, fragment);
-  Type elemType = (fragment == MMAFragment::Lhs)   ? o.aType
-                  : (fragment == MMAFragment::Rhs) ? o.bType
-                                                   : o.cType;
-  return VectorType::get(
-      {s.element[0] * s.element[1] * s.outer[0] * s.outer[1]}, elemType);
-}
-
-std::tuple<VectorType, VectorType, VectorType>
-MMAAttr::getABCVectorTypes() const {
-  MLIRContext *context = getContext();
-  MMAIntrinsic intrinsic = getIntrinsic().getValue();
-  VectorType aVecType = getVectorType(context, intrinsic, MMAFragment::Lhs);
-  VectorType bVecType = getVectorType(context, intrinsic, MMAFragment::Rhs);
-  VectorType cVecType = getVectorType(context, intrinsic, MMAFragment::Acc);
-  return {aVecType, bVecType, cVecType};
-}
-
-FailureOr<std::tuple<VectorLayoutInterface, VectorLayoutInterface,
-                     VectorLayoutInterface>>
-MMAAttr::getContractionLayout(vector::ContractionOp contract) const {
-  ConcreteMmaLayout layout =
-      getConcreteMMALayout(contract->getContext(), getIntrinsic().getValue());
-  return IREE::GPU::getContractionLayout(contract, layout);
-}
-
-static int getBlockSize(MMAIntrinsic /*intrinsic*/) {
-  // Not supporting any block size other than 1 at the moment.
-  return 1;
-}
-
-int64_t MMAAttr::getBlockSize() const {
-  return IREE::GPU::getBlockSize(getIntrinsic().getValue());
-}
-
-static uint32_t getArchID(MMAIntrinsic intrinsic) {
-  return static_cast<int>(intrinsic) & 0xFF00;
-}
-
-static bool is_AMD_MFMA(MMAIntrinsic intrinsic) {
-  return getArchID(intrinsic) >= 0x1000 && getArchID(intrinsic) <= 0x17FF;
-}
-
-static bool is_AMD_WMMA(MMAIntrinsic intrinsic) {
-  return getArchID(intrinsic) >= 0x1800 && getArchID(intrinsic) <= 0x1FFF;
-}
-
-static int64_t getIntrinsicSubgroupSize(MMAIntrinsic intrinsic) {
-  // Not using Wave64 at all at the moment, so the only place where the
-  // subgroup size is CDNA* architectures.
-  return is_AMD_MFMA(intrinsic) ? 64 : 32;
-}
-
-int64_t MMAAttr::getSubgroupSize() const {
-  return getIntrinsicSubgroupSize(getIntrinsic().getValue());
-}
-
-FailureOr<IREE::GPU::MMAScope> MMAAttr::getMmaScope() const {
-  return IREE::GPU::MMAScope::Subgroup;
-}
-
 MMASingleSubgroupLayout getSingleSubgroupLayout(MMAIntrinsic intrinsic,
                                                 MMAFragment fragment) {
   switch (intrinsic) {
@@ -637,6 +274,139 @@
   return {};
 }
 
+template <typename MMAIntrinsicType>
+static OpaqueMmaLayout getOpaqueMMALayout(MLIRContext *context,
+                                          MMAIntrinsicType intrinsic) {
+  OpaqueMmaLayout o;
+  std::tie(o.aType, o.bType, o.cType) = getABCElementTypes(context, intrinsic);
+  auto lhs = getSingleSubgroupLayout(intrinsic, MMAFragment::Lhs);
+  auto rhs = getSingleSubgroupLayout(intrinsic, MMAFragment::Rhs);
+  o.mSize = lhs.outer[0] * lhs.thread[0] * lhs.element[0];
+  o.kSize = lhs.outer[1] * lhs.thread[1] * lhs.element[1];
+  o.nSize = rhs.outer[1] * rhs.thread[1] * rhs.element[1];
+  return o;
+}
+
+OpaqueMmaLayout getOpaqueMMALayout(MLIRContext *context,
+                                   IREE::GPU::MMAIntrinsic intrinsic) {
+  return getOpaqueMMALayout<IREE::GPU::MMAIntrinsic>(context, intrinsic);
+}
+
+//===----------------------------------------------------------------------===//
+// MmaInterface Attribute Helper Functions
+//===----------------------------------------------------------------------===//
+
+MMASingleSubgroupLayout getASingleSubgroupLayout(MmaInterfaceAttr mmaKind) {
+  if (auto mmaAttr = dyn_cast<MMAAttr>(mmaKind)) {
+    return mmaAttr.getASingleSubgroupLayout();
+  }
+  if (auto vmmaAttr = dyn_cast<VirtualMMAAttr>(mmaKind)) {
+    return vmmaAttr.getASingleSubgroupLayout();
+  }
+  assert(false && "unhandled MMA Interface type.");
+  return {};
+}
+
+MMASingleSubgroupLayout getBSingleSubgroupLayout(MmaInterfaceAttr mmaKind) {
+  if (auto mmaAttr = dyn_cast<MMAAttr>(mmaKind)) {
+    return mmaAttr.getBSingleSubgroupLayout();
+  }
+  if (auto vmmaAttr = dyn_cast<VirtualMMAAttr>(mmaKind)) {
+    return vmmaAttr.getBSingleSubgroupLayout();
+  }
+  assert(false && "unhandled MMA Interface type.");
+  return {};
+}
+
+MMASingleSubgroupLayout getCSingleSubgroupLayout(MmaInterfaceAttr mmaKind) {
+  if (auto mmaAttr = dyn_cast<MMAAttr>(mmaKind)) {
+    return mmaAttr.getCSingleSubgroupLayout();
+  }
+  if (auto vmmaAttr = dyn_cast<VirtualMMAAttr>(mmaKind)) {
+    return vmmaAttr.getCSingleSubgroupLayout();
+  }
+  assert(false && "unhandled MMA Interface type.");
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
+// MMA Attributes
+//===----------------------------------------------------------------------===//
+
+Attribute MMAAttr::parse(AsmParser &p, Type type) {
+  if (failed(p.parseLess()))
+    return {};
+
+  FailureOr<MMAIntrinsicAttr> mmaIntrinsic =
+      FieldParser<MMAIntrinsicAttr>::parse(p);
+  if (failed(mmaIntrinsic)) {
+    p.emitError(p.getCurrentLocation(), "failed to parse mfma type identifier");
+    return {};
+  }
+
+  if (failed(p.parseGreater()))
+    return {};
+
+  return get(p.getContext(), mmaIntrinsic->getValue());
+}
+
+void MMAAttr::print(AsmPrinter &p) const {
+  auto &os = p.getStream();
+  os << "<";
+  os << stringifyMMAIntrinsic(getIntrinsic().getValue());
+  os << ">";
+}
+
+MMAAttr MMAAttr::get(MLIRContext *context, MMAIntrinsic type) {
+  auto layout = getOpaqueMMALayout(context, type);
+  return Base::get(context, MMAIntrinsicAttr::get(context, type), layout.mSize,
+                   layout.nSize, layout.kSize, layout.aType, layout.bType,
+                   layout.cType);
+}
+
+std::tuple<Type, Type, Type> MMAAttr::getABCElementTypes() const {
+  return {getAType(), getBType(), getCType()};
+}
+
+std::tuple<int64_t, int64_t, int64_t> MMAAttr::getMNKShape() const {
+  return {getMSize(), getNSize(), getKSize()};
+}
+
+template <typename MMAIntrinsicType>
+static VectorType getVectorType(MLIRContext *context,
+                                MMAIntrinsicType intrinsic,
+                                MMAFragment fragment) {
+  auto o = getOpaqueMMALayout(context, intrinsic);
+  auto s = getSingleSubgroupLayout(intrinsic, fragment);
+  Type elemType = (fragment == MMAFragment::Lhs)   ? o.aType
+                  : (fragment == MMAFragment::Rhs) ? o.bType
+                                                   : o.cType;
+  return VectorType::get(
+      {s.element[0] * s.element[1] * s.outer[0] * s.outer[1]}, elemType);
+}
+
+std::tuple<VectorType, VectorType, VectorType>
+MMAAttr::getABCVectorTypes() const {
+  MLIRContext *context = getContext();
+  MMAIntrinsic intrinsic = getIntrinsic().getValue();
+  VectorType aVecType = getVectorType(context, intrinsic, MMAFragment::Lhs);
+  VectorType bVecType = getVectorType(context, intrinsic, MMAFragment::Rhs);
+  VectorType cVecType = getVectorType(context, intrinsic, MMAFragment::Acc);
+  return {aVecType, bVecType, cVecType};
+}
+
+int64_t MMAAttr::getBlockSize() const {
+  return IREE::GPU::getBlockSize(getIntrinsic().getValue());
+}
+
+int64_t MMAAttr::getSubgroupSize() const {
+  return getIntrinsicSubgroupSize(getIntrinsic().getValue());
+}
+
+FailureOr<IREE::GPU::MMAScope> MMAAttr::getMmaScope() const {
+  return IREE::GPU::MMAScope::Subgroup;
+}
+
 MMASingleSubgroupLayout MMAAttr::getASingleSubgroupLayout() const {
   return getSingleSubgroupLayout(getIntrinsic().getValue(), MMAFragment::Lhs);
 }
@@ -781,22 +551,8 @@
     SmallVector<OpFoldResult> &offsets, SmallVector<OpFoldResult> &sizes,
     SmallVector<OpFoldResult> &strides) const {
 
-  MMASingleSubgroupLayout subgroupLayout;
-  switch (fragment) {
-  case IREE::GPU::MMAFragment::Lhs: {
-    subgroupLayout = getASingleSubgroupLayout();
-    break;
-  }
-  case IREE::GPU::MMAFragment::Rhs: {
-    subgroupLayout = getBSingleSubgroupLayout();
-    break;
-  }
-  case IREE::GPU::MMAFragment::Acc: {
-    subgroupLayout = getCSingleSubgroupLayout();
-    break;
-  }
-  }
-
+  MMASingleSubgroupLayout subgroupLayout =
+      getSingleSubgroupLayout(getIntrinsic().getValue(), fragment);
   SmallVector<OpFoldResult> canonicalOffsets;
   SmallVector<OpFoldResult> canonicalSizes;
   if (failed(populateCanonicalOffsetsSizesAndStrides(
@@ -810,82 +566,6 @@
   return success();
 }
 
-LogicalResult MMAAttr::materializeOperandConcreteShape(
-    OpBuilder &builder, IREE::GPU::MMAFragment fragment, Value operand,
-    std::optional<ArrayRef<int64_t>> permutation,
-    SmallVector<ReassociationIndices> &reassociations,
-    RankedTensorType &resultType) const {
-
-  SmallVector<int64_t, 2> outerSizes;
-  SmallVector<int64_t, 2> opaqueSizes;
-  auto [m, n, k] = getMNKShape();
-  switch (fragment) {
-  case IREE::GPU::MMAFragment::Lhs: {
-    outerSizes = getASingleSubgroupLayout().outer;
-    opaqueSizes.append({m, k});
-    break;
-  }
-  case IREE::GPU::MMAFragment::Rhs: {
-    outerSizes = getBSingleSubgroupLayout().outer;
-    opaqueSizes.append({k, n});
-    break;
-  }
-  case IREE::GPU::MMAFragment::Acc: {
-    outerSizes = getCSingleSubgroupLayout().outer;
-    opaqueSizes.append({m, n});
-    break;
-  }
-  }
-  if (permutation.has_value()) {
-    if (permutation.value().size() != outerSizes.size()) {
-      return failure();
-    }
-    applyPermutationToVector(opaqueSizes, permutation.value());
-    applyPermutationToVector(outerSizes, permutation.value());
-  }
-
-  // Inner tile must have sizes matching the opaque layout.
-  auto operandType = llvm::cast<RankedTensorType>(operand.getType());
-  ArrayRef<int64_t> operandShape = operandType.getShape();
-  SmallVector<int64_t, 2> innerShape(operandShape.end() - opaqueSizes.size(),
-                                     operandShape.end());
-  if (!llvm::equal(opaqueSizes, innerShape)) {
-    return failure();
-  }
-
-  // Expand the shape of the inner tile to reflect the MMA thread layout.
-  SmallVector<int64_t, 4> resultShape(operandShape.begin(),
-                                      operandShape.end() - 2);
-  SmallVector<ReassociationIndices> reInds =
-      llvm::map_to_vector(llvm::seq<int64_t>(resultShape.size()),
-                          [](int64_t idx) -> ReassociationIndices {
-                            return ReassociationIndices({idx});
-                          });
-  int idx = reInds.size();
-  for (auto [outer, native] : llvm::zip_equal(outerSizes, opaqueSizes)) {
-    // Skip expansion if the outer dim is unit as the SingleSubgroupLayout gives
-    // a guarantee that the |element| counts are contiguous within the layout,
-    // and a unit outer implies a single offset and size for that dimension.
-    if (outer == 1) {
-      resultShape.push_back(native);
-      reInds.push_back(ReassociationIndices({idx++}));
-      continue;
-    }
-
-    // Reshape to [outer, native / outer] == [outer, thread * element]. This
-    // corresponds to |outer| repetitions of the thread/element sublayout.
-    resultShape.push_back(outer);
-    assert(native % outer == 0 && "invalid mma layout");
-    resultShape.push_back(native / outer);
-    reInds.push_back(ReassociationIndices{idx, idx + 1});
-    idx += 2;
-  }
-
-  reassociations = reInds;
-  resultType = operandType.clone(resultShape);
-  return success();
-}
-
 //===----------------------------------------------------------------------===//
 // DataTiledMMA Attributes
 //===----------------------------------------------------------------------===//
@@ -1265,22 +945,8 @@
     SmallVector<OpFoldResult> &offsets, SmallVector<OpFoldResult> &sizes,
     SmallVector<OpFoldResult> &strides) const {
 
-  MMASingleSubgroupLayout subgroupLayout;
-  switch (fragment) {
-  case IREE::GPU::MMAFragment::Lhs: {
-    subgroupLayout = getASingleSubgroupLayout();
-    break;
-  }
-  case IREE::GPU::MMAFragment::Rhs: {
-    subgroupLayout = getBSingleSubgroupLayout();
-    break;
-  }
-  case IREE::GPU::MMAFragment::Acc: {
-    subgroupLayout = getCSingleSubgroupLayout();
-    break;
-  }
-  }
-
+  MMASingleSubgroupLayout subgroupLayout =
+      getSingleSubgroupLayout(getIntrinsic().getValue(), fragment);
   SmallVector<OpFoldResult> canonicalOffsets;
   SmallVector<OpFoldResult> canonicalSizes;
   if (failed(populateCanonicalOffsetsSizesAndStrides(
@@ -1445,348 +1111,6 @@
 }
 
 //===----------------------------------------------------------------------===//
-// MMA Schedule Attributes
-//===----------------------------------------------------------------------===//
-
-/// 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]
-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> getIdentityPerm(int64_t rank) {
-  return llvm::to_vector(llvm::seq(static_cast<int64_t>(0), rank));
-}
-
-/// Constructs an identity permutation with the given rank, except it applies
-/// the given rank-2 |perm| to the two dimensions |dim0| and |dim1|, and then
-/// swaps the positions of dim0 and dim1 in the final permutation. For example,
-///
-/// rank = 3, perm = [1, 0], dim0 = 1, dim1 = 2
-/// returns [0, 1, 2]
-///
-/// This is essentially just applying two rank-2 permutations to two particular
-/// dimensions. First it applies |perm|, which corresponds to a permutation
-/// needed by the underlying intrinsic, then it does another permutation based
-/// on the order of actual dimensions for the MMA fragment. For example, for the
-/// B matrix, dim0 = K and dim1 = N, so for the element order of an MFMA
-/// 16x16x16, perm would be `[1, 0]`, however if the actual contraction is a
-/// matmul_transpose_b, then the element order needs to be [0, 1].
-SmallVector<int64_t> getIdentityPermWithSwap(int64_t rank,
-                                             ArrayRef<int64_t> perm,
-                                             int64_t dim0, int64_t dim1) {
-  assert(perm.size() == 2 &&
-         "Unexpected non-rank 2 single subgroup dimension order");
-  SmallVector<int64_t> res = getIdentityPerm(rank);
-  if (perm[0] > perm[1]) {
-    std::swap(dim0, dim1);
-  }
-  if (dim0 > dim1) {
-    res[dim0] = dim1;
-    res[dim1] = dim0;
-  }
-  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).
-NestedLayoutAttr createNestedLayout(MLIRContext *context, int64_t rank,
-                                    int64_t outerDim, int64_t innerDim,
-                                    SmallVector<int64_t> subgroupSizes,
-                                    SmallVector<int64_t> subgroupStrides,
-                                    SmallVector<int64_t> batchCount,
-                                    MMASingleSubgroupLayout counts) {
-
-  LLVM_DEBUG({
-    llvm::errs() << "Creating Nested Layout for::";
-    llvm::errs() << "\n    outerDim = " << outerDim;
-    llvm::errs() << "\n    innerDim = " << innerDim;
-    llvm::errs() << "\n    subgroupSizes: ";
-    llvm::interleaveComma(subgroupSizes, llvm::errs());
-    llvm::errs() << "\n    subgroupStrides: ";
-    llvm::interleaveComma(subgroupStrides, llvm::errs());
-    llvm::errs() << "\n    batchCount: ";
-    llvm::interleaveComma(batchCount, llvm::errs());
-    llvm::errs() << "\n    counts.outer: ";
-    llvm::interleaveComma(counts.outer, llvm::errs());
-    llvm::errs() << "\n    counts.thread: ";
-    llvm::interleaveComma(counts.thread, llvm::errs());
-    llvm::errs() << "\n    counts.element: ";
-    llvm::interleaveComma(counts.element, llvm::errs());
-    llvm::errs() << "\n    counts.tstrides: ";
-    llvm::interleaveComma(counts.tstrides, llvm::errs());
-    llvm::errs() << "\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;
-}
-
-FailureOr<std::tuple<VectorExt::VectorLayoutInterface,
-                     VectorExt::VectorLayoutInterface,
-                     VectorExt::VectorLayoutInterface>>
-MMAScheduleAttr::getContractionLayout(VectorContractOpInfo &opInfo,
-                                      linalg::LinalgOp contractOp) const {
-  LLVM_DEBUG({
-    llvm::errs() << "Getting mma layouts for:\n" << contractOp << "\n";
-    llvm::errs() << "For schedule: " << *this << "\n";
-  });
-
-  int64_t rank = contractOp.getIteratorTypesArray().size();
-  auto mmaAttr = llvm::cast<MmaInterfaceAttr>(getIntrinsic());
-  MLIRContext *context = 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::errs() << "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 = 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 = 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];
-  }
-
-  auto cLayout = createNestedLayout(context, cRank, m, n,
-                                    /*subgroupCount=*/cSubgroupSizes,
-                                    /*subgroupStrides=*/cSubgroupStrides,
-                                    /*batchCount=*/cBatchSizes,
-                                    getCSingleSubgroupLayout(mmaAttr));
-  LLVM_DEBUG({ llvm::errs() << "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;
-
-  auto aLayout = createNestedLayout(context, aRank, afm, afk,
-                                    /*subgroupCount=*/aSubgroupSizes,
-                                    /*subgroupStrides=*/aSubgroupStrides,
-                                    /*batchCount=*/aBatchSizes,
-                                    getASingleSubgroupLayout(mmaAttr));
-  LLVM_DEBUG({ llvm::errs() << "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;
-
-  auto bLayout = createNestedLayout(context, bRank, bfk, bfn,
-                                    /*subgroupCount=*/bSubgroupSizes,
-                                    /*subgroupStrides=*/bSubgroupStrides,
-                                    /*batchCount=*/bBatchSizes,
-                                    getBSingleSubgroupLayout(mmaAttr));
-  LLVM_DEBUG({ llvm::errs() << "B layout: " << bLayout << "\n"; });
-
-  std::tuple<VectorLayoutInterface, VectorLayoutInterface,
-             VectorLayoutInterface>
-      result = {aLayout, bLayout, cLayout};
-  return result;
-}
-
-//===----------------------------------------------------------------------===//
 // Target Attributes
 //===----------------------------------------------------------------------===//
 
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h
index 7942066..92ce4f4 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h
@@ -22,12 +22,40 @@
 
 namespace mlir::iree_compiler::IREE::GPU {
 
-// Partial nested layout for an MMA intrinsic's matrix input/output inside
-// a single subgroup.
+// Struct describing the detailed subgroup-level layout of a MMA intrinsic.
+// Together with element type information and subgroup size, it completes the
+// full description of the semantics of a MMA intrinsic.
+//
+// Note: It is not possible to infer subgroup size from the information in this
+// struct. The product of the `thread` sizes here is often, but not always equal
+// to subgroup size. When the product of the `thread` sizes (call that product
+// `P`) is smaller than subgroup size, it must be a divisor of it, and the
+// semantics in that case are that threads within the subgroup whose thread-ids
+// differ by a multiple of `P`, are accessing the same elements.
+//
+// Example observed in RDNA3 WMMA Wave64 intrinsics:
+// If the subgroup size is 64 but the product `P` of `thread` sizes is 32, that
+// means that each element is being accessed by 2 threads (2 = 64/32), and the
+// threads accessing the same element are those whose tids are exactly 32 apart.
 struct MMASingleSubgroupLayout {
+  // Internal dimensions (as in TileSwizzle::Dim::Kind::Internal) that are
+  // outer-most in the layout. This happens when a MMA op, seen on a single
+  // thread, has an operand that consists of multiple elements, and these elems
+  // are NOT contiguous.
+  // This is not used by every MMA op; ops which don't use that simply have 1's.
   SmallVector<int64_t, 2> outer;
+  // Cross-thread dimensions (as in TileSwizzle::Dim::Kind::CrossThread).
+  // This is the kind of dimension that is present in all GPU MMA ops, by
+  // definition of "SIMT". It is still possible for one of the `thread` dims to
+  // be 1, but not both.
   SmallVector<int64_t, 2> thread;
+  // Strides corresponding to the cross-thread dimensions.
   SmallVector<int64_t, 2> tstrides;
+  // Internal dimensions (as in TileSwizzle::Dim::Kind::Internal) that are
+  // inner-most in the layout. This happens when a MMA op, seen on a single
+  // thread, has an operand that consists of multiple elements, and these elems
+  // are NOT contiguous.
+  // This is not used by every MMA op; ops which don't use that simply have 1's.
   SmallVector<int64_t, 2> element;
 };
 
@@ -43,6 +71,22 @@
 
 MMASingleSubgroupLayout getCSingleSubgroupLayout(MmaInterfaceAttr mmaKind);
 
+// Struct describing the shape of a MMA operation, but not the detailed layout.
+// TODO(bjacob): the only user outside of IREEGPUAttrs.cpp is
+// LLVMGPU/TransformExtensions, so maybe make that internal again if/when that
+// goes away.
+struct OpaqueMmaLayout {
+  int64_t mSize = 0;
+  int64_t nSize = 0;
+  int64_t kSize = 0;
+  Type aType;
+  Type bType;
+  Type cType;
+};
+
+OpaqueMmaLayout getOpaqueMMALayout(MLIRContext *context,
+                                   IREE::GPU::MMAIntrinsic intrinsic);
+
 } // namespace mlir::iree_compiler::IREE::GPU
 
 // clang-format off
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td
index bc6f2d3..51674db 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td
@@ -148,7 +148,6 @@
   DeclareAttrInterfaceMethods<IREEGPU_MmaInterfaceAttr, [
     "getABCElementTypes",
     "getABCVectorTypes",
-    "getContractionLayout",
     "getMNKShape",
     "getSubgroupSize",
     "getMmaScope",
@@ -157,7 +156,6 @@
     "getCSingleSubgroupLayout",
     "buildMmaOperation",
     "populateOperandOffsetsSizesStrides",
-    "materializeOperandConcreteShape",
   ]>
 ]> {
   let cppNamespace = "::mlir::iree_compiler::IREE::GPU";
@@ -259,19 +257,11 @@
     |intrinsic| field specifies which particular MMA intrinsic is targeted by
     the data-tiling.
 
-    The tile swizzling already happens, so the attribute does not need to
-    implement materializeOperandConcreteShape interface method. E.g., if the
-    target intrinsic is MFMA_F32_16x16x4_F32:
-      - The inner tile shape of LHS is 4x16.
-      - The inner tile shape of RHS is 4x16.
-      - The inner tile shape of ACC is 4x16x4.
-
-    Furthermore, the unrolling and interleaving can be represented with the
-    attribute. In the concept of data-tiling, we always unroll the parallel
-    dimensions (i.e., M, N dimensions) to be outermost, and interleave the
-    unrolled K dimension. I.e., the unrolled K dimension becomes the innermost
-    dimension. The constraint can be relaxed based on data-tiling needs. The
-    additional information can be added to `parameters`.
+    The other fields default to one, and that default results in a single
+    intrinsic equivalent to MMAAttr, while values greater than one result in
+    wider "kernels" consisting of multiple intrinsics, with the data layout
+    already swizzled into a tile layout that allows each intrinsic to access
+    data at an offset that's as simple as possible a mapping from the thread ID.
   }];
 
   let assemblyFormat = "`<` struct(params) `>`";
@@ -369,15 +359,6 @@
   );
 
   let assemblyFormat = "`<` struct(params) `>`";
-
-  let extraClassDeclaration = [{
-    // Returns the A/B/C matrix concrete layout targeting |contractOp|.
-    ::mlir::FailureOr<::std::tuple<VectorExt::VectorLayoutInterface,
-                                 VectorExt::VectorLayoutInterface,
-                                 VectorExt::VectorLayoutInterface>>
-      getContractionLayout(::mlir::iree_compiler::VectorContractOpInfo &opInfo,
-                           ::mlir::linalg::LinalgOp contractOp) const;
-  }];
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/ConcretizeMmaShapes.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/ConcretizeMmaShapes.cpp
index 94bcf3d..6b2e834 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/ConcretizeMmaShapes.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/ConcretizeMmaShapes.cpp
@@ -23,7 +23,80 @@
   using ConcretizeMmaShapesPassBase::ConcretizeMmaShapesPassBase;
   void runOnOperation() override;
 };
-} // namespace
+
+LogicalResult materializeOperandConcreteShape(
+    OpBuilder &builder, MMAAttr mma, IREE::GPU::MMAFragment fragment,
+    Value operand, std::optional<ArrayRef<int64_t>> permutation,
+    SmallVector<ReassociationIndices> &reassociations,
+    RankedTensorType &resultType) {
+
+  SmallVector<int64_t, 2> outerSizes;
+  SmallVector<int64_t, 2> opaqueSizes;
+  auto [m, n, k] = mma.getMNKShape();
+  switch (fragment) {
+  case IREE::GPU::MMAFragment::Lhs: {
+    outerSizes = mma.getASingleSubgroupLayout().outer;
+    opaqueSizes.append({m, k});
+    break;
+  }
+  case IREE::GPU::MMAFragment::Rhs: {
+    outerSizes = mma.getBSingleSubgroupLayout().outer;
+    opaqueSizes.append({k, n});
+    break;
+  }
+  case IREE::GPU::MMAFragment::Acc: {
+    outerSizes = mma.getCSingleSubgroupLayout().outer;
+    opaqueSizes.append({m, n});
+    break;
+  }
+  }
+  if (permutation.has_value()) {
+    if (permutation.value().size() != outerSizes.size()) {
+      return failure();
+    }
+    applyPermutationToVector(opaqueSizes, permutation.value());
+    applyPermutationToVector(outerSizes, permutation.value());
+  }
+
+  // Inner tile must have sizes matching the opaque layout.
+  auto operandType = llvm::cast<RankedTensorType>(operand.getType());
+  ArrayRef<int64_t> operandShape = operandType.getShape();
+  if (opaqueSizes != operandShape.take_back(opaqueSizes.size())) {
+    return failure();
+  }
+
+  // Expand the shape of the inner tile to reflect the MMA thread layout.
+  SmallVector<int64_t, 4> resultShape(operandShape.begin(),
+                                      operandShape.end() - 2);
+  SmallVector<ReassociationIndices> reInds =
+      llvm::map_to_vector(llvm::seq<int64_t>(resultShape.size()),
+                          [](int64_t idx) -> ReassociationIndices {
+                            return ReassociationIndices({idx});
+                          });
+  int idx = reInds.size();
+  for (auto [outer, native] : llvm::zip_equal(outerSizes, opaqueSizes)) {
+    // Skip expansion if the outer dim is unit as the SingleSubgroupLayout gives
+    // a guarantee that the |element| counts are contiguous within the layout,
+    // and a unit outer implies a single offset and size for that dimension.
+    if (outer == 1) {
+      resultShape.push_back(native);
+      reInds.push_back(ReassociationIndices({idx++}));
+      continue;
+    }
+
+    // Reshape to [outer, native / outer] == [outer, thread * element]. This
+    // corresponds to |outer| repetitions of the thread/element sublayout.
+    resultShape.push_back(outer);
+    assert(native % outer == 0 && "invalid mma layout");
+    resultShape.push_back(native / outer);
+    reInds.push_back(ReassociationIndices{idx, idx + 1});
+    idx += 2;
+  }
+
+  reassociations = reInds;
+  resultType = operandType.clone(resultShape);
+  return success();
+}
 
 struct ConcretizeMmaOperandShape final : OpRewritePattern<MultiMmaOp> {
   using OpRewritePattern::OpRewritePattern;
@@ -56,12 +129,15 @@
     }
 
     // Get the reassociation indices and result type of the expand_shape op.
-    MmaInterfaceAttr kind = mmaOp.getKind();
+    MMAAttr kind = dyn_cast<MMAAttr>(mmaOp.getKind());
+    if (!kind) {
+      return failure();
+    }
     SmallVector<ReassociationIndices> reassociations;
     RankedTensorType concreteType;
-    if (failed(kind.materializeOperandConcreteShape(rewriter, fragment, operand,
-                                                    permutation, reassociations,
-                                                    concreteType))) {
+    if (failed(materializeOperandConcreteShape(rewriter, kind, fragment,
+                                               operand, permutation,
+                                               reassociations, concreteType))) {
       return failure();
     }
 
@@ -140,6 +216,8 @@
   MMAFragment fragment;
 };
 
+} // namespace
+
 void ConcretizeMmaShapesPass::runOnOperation() {
   MLIRContext *context = &getContext();
   auto funcOp = getOperation();
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureTensorLayouts.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureTensorLayouts.cpp
index 69333a1..9df22c7 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureTensorLayouts.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureTensorLayouts.cpp
@@ -8,6 +8,7 @@
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
 #include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
 #include "iree/compiler/Codegen/LLVMGPU/Passes.h"
+#include "llvm/Support/Debug.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/BuiltinAttributes.h"
@@ -71,6 +72,310 @@
   return *config.getSubgroupNCount();
 }
 
+/// 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;
+}
+
+static FailureOr<std::tuple<IREE::VectorExt::VectorLayoutInterface,
+                            IREE::VectorExt::VectorLayoutInterface,
+                            IREE::VectorExt::VectorLayoutInterface>>
+getContractionLayout(IREE::GPU::MMAScheduleAttr schedule,
+                     VectorContractOpInfo &opInfo,
+                     linalg::LinalgOp 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::cast<IREE::GPU::MmaInterfaceAttr>(schedule.getIntrinsic());
+  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];
+  }
+
+  auto cLayout = createNestedLayout(context, cRank, m, n,
+                                    /*subgroupCount=*/cSubgroupSizes,
+                                    /*subgroupStrides=*/cSubgroupStrides,
+                                    /*batchCount=*/cBatchSizes,
+                                    getCSingleSubgroupLayout(mmaAttr));
+  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;
+
+  auto aLayout = createNestedLayout(context, aRank, afm, afk,
+                                    /*subgroupCount=*/aSubgroupSizes,
+                                    /*subgroupStrides=*/aSubgroupStrides,
+                                    /*batchCount=*/aBatchSizes,
+                                    getASingleSubgroupLayout(mmaAttr));
+  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;
+
+  auto bLayout = createNestedLayout(context, bRank, bfk, bfn,
+                                    /*subgroupCount=*/bSubgroupSizes,
+                                    /*subgroupStrides=*/bSubgroupStrides,
+                                    /*batchCount=*/bBatchSizes,
+                                    getBSingleSubgroupLayout(mmaAttr));
+  LLVM_DEBUG({ llvm::dbgs() << "B layout: " << bLayout << "\n"; });
+
+  std::tuple<VectorLayoutInterface, VectorLayoutInterface,
+             VectorLayoutInterface>
+      result = {aLayout, bLayout, cLayout};
+  return result;
+}
+
 static LogicalResult setContractionAnchor(IREE::GPU::MMAScheduleAttr schedule,
                                           SmallVector<bool> promotedOperands,
                                           RewriterBase &rewriter,
@@ -89,7 +394,7 @@
           contract.getIndexingMapsArray());
   assert(succeeded(opInfo) && "contraction should have been inferred");
 
-  auto layouts = schedule.getContractionLayout(opInfo.value(), contract);
+  auto layouts = getContractionLayout(schedule, opInfo.value(), contract);
   if (failed(layouts)) {
     return contract->emitError("cannot get concrete layout for contraction");
   }
@@ -176,7 +481,7 @@
   assert(succeeded(opInfo) &&
          "unit filter dim convolution should have been infered");
 
-  auto layouts = schedule.getContractionLayout(opInfo.value(), conv);
+  auto layouts = getContractionLayout(schedule, opInfo.value(), conv);
   if (failed(layouts)) {
     return conv->emitError("cannot get concrete layout for convolution");
   }
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp
index 3dd0c12..a2c2187 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp
@@ -11,6 +11,7 @@
 #include "iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.h"
 #include "iree/compiler/Codegen/Common/GPU/Passes.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
+#include "iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h"
 #include "iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUUtils.h"
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
@@ -1599,6 +1600,219 @@
   return VectorExt::LayoutAttr::get(ctx, perDimLayouts);
 }
 
+// Struct containing concrete MMA shape, type, and layout information.
+struct ConcreteMmaLayout {
+  GPU::OpaqueMmaLayout base;
+  VectorExt::PerDimLayoutAttr aMLayout;
+  VectorExt::PerDimLayoutAttr aKLayout;
+  VectorExt::PerDimLayoutAttr bKLayout;
+  VectorExt::PerDimLayoutAttr bNLayout;
+  VectorExt::PerDimLayoutAttr cMLayout;
+  VectorExt::PerDimLayoutAttr cNLayout;
+};
+
+static std::tuple<VectorExt::PerDimLayoutAttr, VectorExt::PerDimLayoutAttr>
+getPerDimLayoutAttrs(MLIRContext *context, TileSwizzle swizzle) {
+  // Step 1: obtain the swizzled tile shape, but keeping track of the source
+  // dimension indices.
+  struct SrcIndexAndSwizzleDim {
+    size_t srcIndex;
+    TileSwizzle::Dim dim;
+  };
+  SmallVector<SrcIndexAndSwizzleDim> swizzledShape;
+  for (auto [i, e] : llvm::enumerate(swizzle.expandShape)) {
+    for (TileSwizzle::Dim d : e) {
+      swizzledShape.push_back(SrcIndexAndSwizzleDim{i, d});
+    }
+  }
+  applyPermutationToVector(swizzledShape, swizzle.permutation);
+
+  // Step 2: collect the appropriate labels to use for the swizzled dims.
+  VectorExt::LayoutDimension internalLabels[] = {
+      VectorExt::LayoutDimension::VECTORZ, VectorExt::LayoutDimension::VECTORY,
+      VectorExt::LayoutDimension::VECTORX};
+  VectorExt::LayoutDimension crossThreadLabels[] = {
+      VectorExt::LayoutDimension::LANEZ, VectorExt::LayoutDimension::LANEY,
+      VectorExt::LayoutDimension::LANEX};
+  auto internalLabelIter = std::end(internalLabels);
+  auto crossThreadLabelIter = std::end(crossThreadLabels);
+  for (SrcIndexAndSwizzleDim d : swizzledShape) {
+    if (d.dim.kind == TileSwizzle::Dim::Kind::Internal) {
+      assert(internalLabelIter != std::begin(internalLabels));
+      --internalLabelIter;
+    } else if (d.dim.kind == TileSwizzle::Dim::Kind::CrossThread) {
+      assert(crossThreadLabelIter != std::begin(crossThreadLabels));
+      --crossThreadLabelIter;
+    } else {
+      assert(false && "unexpected dimension kind in intrinsic swizzle");
+    }
+  }
+
+  // Step 3: put together the result PerDimLayoutAttr'd for the two source dims.
+  SmallVector<VectorExt::LayoutDimensionAttr> labels[2];
+  SmallVector<int64_t> shape[2];
+  for (SrcIndexAndSwizzleDim d : swizzledShape) {
+    shape[d.srcIndex].push_back(d.dim.size);
+    auto &labelIterRef = (d.dim.kind == TileSwizzle::Dim::Kind::Internal)
+                             ? internalLabelIter
+                             : crossThreadLabelIter;
+    labels[d.srcIndex].push_back(VectorExt::LayoutDimensionAttr::get(
+        context, static_cast<VectorExt::LayoutDimension>(*labelIterRef++)));
+  }
+  return {VectorExt::PerDimLayoutAttr::get(context, labels[0], shape[0]),
+          VectorExt::PerDimLayoutAttr::get(context, labels[1], shape[1])};
+};
+
+static ConcreteMmaLayout getConcreteMMALayout(MLIRContext *context,
+                                              GPU::MMAIntrinsic intrinsic) {
+  auto opaque = GPU::getOpaqueMMALayout(context, intrinsic);
+  ConcreteMmaLayout concreteLayout;
+  concreteLayout.base = opaque;
+  auto lhsSwizzle = getIntrinsicSwizzle(intrinsic, GPU::MMAFragment::Lhs);
+  auto rhsSwizzle = getIntrinsicSwizzle(intrinsic, GPU::MMAFragment::Rhs);
+  auto accSwizzle = getIntrinsicSwizzle(intrinsic, GPU::MMAFragment::Acc);
+  std::tie(concreteLayout.aMLayout, concreteLayout.aKLayout) =
+      getPerDimLayoutAttrs(context, lhsSwizzle);
+  std::tie(concreteLayout.bNLayout, concreteLayout.bKLayout) =
+      getPerDimLayoutAttrs(context, rhsSwizzle);
+  std::tie(concreteLayout.cMLayout, concreteLayout.cNLayout) =
+      getPerDimLayoutAttrs(context, accSwizzle);
+  return concreteLayout;
+}
+
+static VectorExt::PerDimLayoutAttr
+getBatchedPerDimLayoutAttr(VectorExt::LayoutDimensionAttr batchDim,
+                           VectorExt::PerDimLayoutAttr baseLayout,
+                           int64_t problemSize, int64_t fragmentDimSize) {
+  assert(problemSize % fragmentDimSize == 0 &&
+         "invalid layout fragment for problem size");
+
+  SmallVector<VectorExt::LayoutDimensionAttr, 3> dimAttrs(
+      baseLayout.getLabels());
+  dimAttrs.insert(dimAttrs.begin(), batchDim);
+
+  SmallVector<int64_t, 3> shapes(baseLayout.getShapes());
+  shapes.insert(shapes.begin(), problemSize / fragmentDimSize);
+  auto layout = VectorExt::PerDimLayoutAttr::get(baseLayout.getContext(),
+                                                 dimAttrs, shapes);
+  return layout;
+}
+
+// Get the batched layout attributes for the given fragment layouts, indexing
+// map, and problem shape. The canonical fragment map is used to compare against
+// the problem map |indexingMap|. For example, for mma fragment B (RHS):
+//
+// indexingMap = affine_map<(d0, d1, d2) -> (d1, d2) # Transposed B
+// fragmentMap = affine_map<(d0, d1, d2) -> (d2, d1)
+// problemShape = [32, 64]
+// fragmentSize = [16, 8]
+// fragmentLayouts = [kLayout, nLayout]
+//
+// Gives batched layout
+//
+// Dim0 Layout = [BATCHX, nLayoutLabels], [8, nLayoutShape]
+// Dim1 Layout = [BATCHY, kLayoutLabels], [2, kLayoutShape]
+static VectorExt::LayoutAttr
+getBatchedLayoutAttr(AffineMap indexingMap, AffineMap fragmentMap,
+                     ArrayRef<int64_t> problemShape,
+                     ArrayRef<int64_t> fragmentSize,
+                     ArrayRef<VectorExt::PerDimLayoutAttr> fragmentLayouts) {
+  // Current distribution to MFMA operations does not support batched
+  // contractions so that is reflected here.
+  assert(indexingMap.getNumResults() == 2 &&
+         "invalid indexing map to non-batched simple contraction");
+
+  VectorExt::LayoutDimensionAttr batchX = VectorExt::LayoutDimensionAttr::get(
+      indexingMap.getContext(), VectorExt::LayoutDimension::BATCHX);
+  VectorExt::LayoutDimensionAttr batchY = VectorExt::LayoutDimensionAttr::get(
+      indexingMap.getContext(), VectorExt::LayoutDimension::BATCHY);
+
+  SmallVector<VectorExt::PerDimLayoutAttr, 2> perDimAttrs;
+  for (auto [expr, batchType] : llvm::zip_equal(
+           indexingMap.getResults(),
+           SmallVector<VectorExt::LayoutDimensionAttr, 2>{batchX, batchY})) {
+    auto maybeResultPosition = fragmentMap.getResultPosition(expr);
+    assert(maybeResultPosition && "fragment map and problem map mismatch");
+    int64_t idx = *maybeResultPosition;
+    perDimAttrs.push_back(getBatchedPerDimLayoutAttr(
+        batchType, fragmentLayouts[idx], problemShape[idx], fragmentSize[idx]));
+  }
+
+  return VectorExt::LayoutAttr::get(indexingMap.getContext(), perDimAttrs);
+}
+
+static FailureOr<std::tuple<VectorLayoutInterface, VectorLayoutInterface,
+                            VectorLayoutInterface>>
+getContractionLayout(vector::ContractionOp contract, ConcreteMmaLayout layout) {
+  MLIRContext *context = contract.getContext();
+  FailureOr<linalg::ContractionDimensions> maybeContractionDims =
+      linalg::inferContractionDims(contract.getIndexingMapsArray());
+  if (failed(maybeContractionDims)) {
+    return failure();
+  }
+  auto contractionDims = *maybeContractionDims;
+  // TODO: Relax this condition to strictly alignment requirements.
+  if (contractionDims.k.size() != 1 || contractionDims.m.size() != 1 ||
+      contractionDims.n.size() != 1) {
+    return failure();
+  }
+  // TODO: Support batched contractions.
+  if (contractionDims.batch.size() > 0) {
+    return failure();
+  }
+  unsigned mDim = contractionDims.m[0];
+  unsigned nDim = contractionDims.n[0];
+  unsigned kDim = contractionDims.k[0];
+
+  SmallVector<int64_t> iterationBounds;
+  contract.getIterationBounds(iterationBounds);
+
+  int64_t problemMSize = iterationBounds[mDim];
+  int64_t problemNSize = iterationBounds[nDim];
+  int64_t problemKSize = iterationBounds[kDim];
+
+  int64_t mSize = layout.base.mSize;
+  int64_t nSize = layout.base.nSize;
+  int64_t kSize = layout.base.kSize;
+
+  // The problem size currently must be strictly aligned to the size of the mma.
+  // This is expected to succeed assuming the correct [masked] vector size was
+  // set at strategy configuration time (for this mma).
+  if (problemMSize % mSize != 0 || problemNSize % nSize ||
+      problemKSize % kSize) {
+    return failure();
+  }
+
+  VectorExt::LayoutAttr aLayout = getBatchedLayoutAttr(
+      contract.getIndexingMapsArray()[0],
+      AffineMap::getMultiDimMapWithTargets(3, {mDim, kDim}, context),
+      {problemMSize, problemKSize}, {mSize, kSize},
+      {layout.aMLayout, layout.aKLayout});
+  VectorExt::LayoutAttr bLayout = getBatchedLayoutAttr(
+      contract.getIndexingMapsArray()[1],
+      AffineMap::getMultiDimMapWithTargets(3, {kDim, nDim}, context),
+      {problemKSize, problemNSize}, {kSize, nSize},
+      {layout.bKLayout, layout.bNLayout});
+  VectorExt::LayoutAttr cLayout = getBatchedLayoutAttr(
+      contract.getIndexingMapsArray()[2],
+      AffineMap::getMultiDimMapWithTargets(3, {mDim, nDim}, context),
+      {problemMSize, problemNSize}, {mSize, nSize},
+      {layout.cMLayout, layout.cNLayout});
+
+  return std::make_tuple<VectorLayoutInterface, VectorLayoutInterface,
+                         VectorLayoutInterface>(aLayout, bLayout, cLayout);
+}
+
+FailureOr<std::tuple<
+    VectorLayoutInterface, VectorLayoutInterface,
+    VectorLayoutInterface>> static getContractionLayout(GPU::MMAAttr mma,
+                                                        vector::ContractionOp
+                                                            contract) {
+  ConcreteMmaLayout layout = getConcreteMMALayout(
+      contract->getContext(), mma.getIntrinsic().getValue());
+  return getContractionLayout(contract, layout);
+}
+
 DiagnosedSilenceableFailure
 transform_dialect::SetContractionLayoutAttributes::apply(
     transform::TransformRewriter &rewriter,
@@ -1609,7 +1823,7 @@
     return emitDefiniteFailure()
            << "invalid more than one attribute for contraction annotation";
   }
-  auto mmaType = llvm::dyn_cast<IREE::GPU::MmaInterfaceAttr>(typeList.front());
+  auto mmaType = llvm::dyn_cast<GPU::MMAAttr>(typeList.front());
   if (!mmaType) {
     return emitDefiniteFailure()
            << "invalid non-mma attribute for contraction annotation "
@@ -1623,7 +1837,7 @@
              << "invalid non-contraction annotation " << payload;
     }
 
-    auto maybeLayouts = mmaType.getContractionLayout(contract);
+    auto maybeLayouts = getContractionLayout(mmaType, contract);
     if (failed(maybeLayouts)) {
       return emitDefiniteFailure()
              << "invalid opaque mma layout for annotation " << mmaType;