// Copyright 2024 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"

#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenTypes.h"
#include "iree/compiler/Codegen/Dialect/Codegen/Utils/MMAUtils.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/DerivedConfigUtils.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.h"
#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/GPU/IR/IREEGPUOps.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/MatchUtils.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "iree/compiler/Utils/EncodingUtils.h"
#include "iree/compiler/Utils/Indexing.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/STLForwardCompat.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/InterleavedRange.h"
#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/Interfaces/LoopLikeInterface.h"

#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.cpp.inc"
#define GET_ATTRDEF_CLASSES
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp.inc"

#define DEBUG_TYPE "iree-gpu-attrs"

namespace mlir::iree_compiler::IREE::GPU {

// HoistableConversionOp tag constants: definitions live in
// `Codegen/Utils/MMAUtils.h` so paired tags (which match by string) can't
// silently drift between translation units.
using ::mlir::iree_compiler::kCompilationInfoOutputName;
using ::mlir::iree_compiler::kDecompositionConfigKey;
using ::mlir::iree_compiler::IREE::Codegen::distributeMmaFragmentToIntrinsics;
using ::mlir::iree_compiler::IREE::Codegen::incrementIndices;
using ::mlir::iree_compiler::IREE::Codegen::kDataTiledAccDistribute;
using ::mlir::iree_compiler::IREE::Codegen::kDataTiledAccReassemble;
using ::mlir::iree_compiler::IREE::Codegen::kRdna3DeinterleaveAcc;
using ::mlir::iree_compiler::IREE::Codegen::kRdna3InterleaveAcc;
using ::mlir::iree_compiler::IREE::Codegen::kVDMFMADeinterleaveAcc;
using ::mlir::iree_compiler::IREE::Codegen::kVDMFMAInterleaveAcc;
using ::mlir::iree_compiler::IREE::Codegen::TileSwizzle;

constexpr StringLiteral kWorkgroupSizeKey = "workgroup_size";
constexpr StringLiteral kSubgroupSizeKey = "subgroup_size";
constexpr StringLiteral kTranslationInfoKey = "translation_info";
constexpr StringLiteral kPipelineKey = "pipeline";
constexpr StringLiteral kCodegenSpecKey = "codegen_spec";
constexpr StringLiteral kGpuPipelineOptionsKey = "gpu_pipeline_options";
constexpr StringLiteral kPrefetchNumStagesKey = "prefetch_num_stages";
constexpr StringLiteral kNoReduceSharedMemoryBankConflictsKey =
    "no_reduce_shared_memory_bank_conflicts";
constexpr StringLiteral kUseIgemmConvolutionKey = "use_igemm_convolution";
constexpr StringLiteral kReorderWorkgroupsStrategyKey =
    "reorder_workgroups_strategy";

static bool isDefaultTranslationInfoKey(StringRef key) {
  return llvm::is_contained({kWorkgroupSizeKey, kSubgroupSizeKey}, key);
}

static bool isTranslationInfoConfigKey(StringRef key) {
  return key == kGpuPipelineOptionsKey;
}

static bool isFlatMetadataKey(StringRef key) {
  return llvm::is_contained({kConfigAttrName, kTranslationInfoKey, kPipelineKey,
                             kCodegenSpecKey, kDecompositionConfigKey},
                            key);
}

static bool isCompilationInfoMaterializationInput(StringRef key) {
  return key != kDecompositionConfigKey;
}

static DictionaryAttr
getCompilationInfoMaterializationInputs(MLIRContext *ctx,
                                        DictionaryAttr knobs) {
  SmallVector<NamedAttribute> entries;
  for (NamedAttribute entry : knobs) {
    if (isCompilationInfoMaterializationInput(entry.getName().getValue())) {
      entries.push_back(entry);
    }
  }
  return DictionaryAttr::get(ctx, entries);
}

static void
appendLoweringConfigEntries(DictionaryAttr source,
                            SmallVectorImpl<NamedAttribute> &entries,
                            bool flatSource) {
  for (NamedAttribute entry : source) {
    StringRef key = entry.getName().getValue();
    if (flatSource &&
        (isFlatMetadataKey(key) || isDefaultTranslationInfoKey(key) ||
         isTranslationInfoConfigKey(key))) {
      continue;
    }
    entries.push_back(entry);
  }
}

static DictionaryAttr buildTranslationConfiguration(MLIRContext *ctx,
                                                    DictionaryAttr source) {
  SmallVector<NamedAttribute> entries;
  for (NamedAttribute entry : source) {
    StringRef key = entry.getName().getValue();
    if (isDefaultTranslationInfoKey(key) || key == kCodegenSpecKey) {
      continue;
    }
    entries.push_back(entry);
  }
  return DictionaryAttr::get(ctx, entries);
}

static FailureOr<DictionaryAttr> buildGpuPipelineOptionsConfiguration(
    MLIRContext *ctx, Attribute gpuPipelineOptions,
    function_ref<InFlightDiagnostic()> emitError) {
  auto optionsDict = dyn_cast<DictionaryAttr>(gpuPipelineOptions);
  if (!optionsDict) {
    emitError() << "expected '" << kGpuPipelineOptionsKey
                << "' to be a dictionary";
    return failure();
  }

  unsigned prefetchNumStages = 0;
  if (Attribute prefetchAttr = optionsDict.get(kPrefetchNumStagesKey)) {
    auto prefetchInt = dyn_cast<IntegerAttr>(prefetchAttr);
    if (!prefetchInt || prefetchInt.getInt() < 0) {
      emitError() << "expected '" << kPrefetchNumStagesKey
                  << "' to be a non-negative integer";
      return failure();
    }
    prefetchNumStages = static_cast<unsigned>(prefetchInt.getInt());
  }

  bool noReduceSharedMemoryBankConflicts = false;
  if (Attribute noReduceAttr =
          optionsDict.get(kNoReduceSharedMemoryBankConflictsKey)) {
    auto noReduceBool = dyn_cast<BoolAttr>(noReduceAttr);
    if (!noReduceBool) {
      emitError() << "expected '" << kNoReduceSharedMemoryBankConflictsKey
                  << "' to be a boolean";
      return failure();
    }
    noReduceSharedMemoryBankConflicts = noReduceBool.getValue();
  }

  bool useIgemmConvolution = false;
  if (Attribute useIgemmAttr = optionsDict.get(kUseIgemmConvolutionKey)) {
    auto useIgemmBool = dyn_cast<BoolAttr>(useIgemmAttr);
    if (!useIgemmBool) {
      emitError() << "expected '" << kUseIgemmConvolutionKey
                  << "' to be a boolean";
      return failure();
    }
    useIgemmConvolution = useIgemmBool.getValue();
  }

  std::optional<ReorderWorkgroupsStrategy> reorderWorkgroupsStrategy;
  if (Attribute reorderStrategyAttr =
          optionsDict.get(kReorderWorkgroupsStrategyKey)) {
    auto reorderAttr =
        dyn_cast<ReorderWorkgroupsStrategyAttr>(reorderStrategyAttr);
    if (!reorderAttr) {
      emitError() << "expected '" << kReorderWorkgroupsStrategyKey
                  << "' to be a reorder strategy attr";
      return failure();
    }
    reorderWorkgroupsStrategy = reorderAttr.getValue();
  }

  auto typedOptions = GPUPipelineOptionsAttr::get(
      ctx, prefetchNumStages, noReduceSharedMemoryBankConflicts,
      useIgemmConvolution, reorderWorkgroupsStrategy);
  return DictionaryAttr::get(
      ctx, {{StringAttr::get(ctx, kGpuPipelineOptionsKey), typedOptions}});
}

static int64_t extractI64(Attribute attr) {
  return cast<IntegerAttr>(attr).getInt();
}

static SmallVector<int64_t> extractI64Array(Attribute attr) {
  return llvm::map_to_vector(cast<ArrayAttr>(attr).getValue(), extractI64);
}

//===----------------------------------------------------------------------===//
// MMA intrinsics semantics: shapes, layouts, operand element types.
//===----------------------------------------------------------------------===//

static LogicalResult verifyMmaIndexingMaps(ArrayRef<AffineMap> maps) {
  return linalg::inferContractionDims(maps);
}

static int getBlockSize(MMAIntrinsic intrinsic) {
  switch (intrinsic) {
  case MMAIntrinsic::MFMA_F32_4x4x4x16B_F16:
  case MMAIntrinsic::MFMA_F32_4x4x4x16B_BF16:
  case MMAIntrinsic::MFMA_I32_4x4x4x16B_I8:
  case MMAIntrinsic::MFMA_F32_4x4x1x16B_F32:
    return 16;
  case MMAIntrinsic::MFMA_F32_16x16x4x4B_F16:
  case MMAIntrinsic::MFMA_F32_16x16x4x4B_BF16:
  case MMAIntrinsic::MFMA_I32_16x16x4x4B_I8:
  case MMAIntrinsic::MFMA_F32_16x16x1x4B_F32:
  case MMAIntrinsic::MFMA_F64_4x4x4x4B_F64:
    return 4;
  case MMAIntrinsic::MFMA_F32_32x32x4x2B_F16:
  case MMAIntrinsic::MFMA_F32_32x32x4x2B_BF16:
  case MMAIntrinsic::MFMA_I32_32x32x4x2B_I8:
  case MMAIntrinsic::MFMA_F32_32x32x1x2B_F32:
    return 2;
  default:
    return 1;
  }
}

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 bool is_AMD(MMAIntrinsic intrinsic) {
  return is_AMD_MFMA(intrinsic) || is_AMD_WMMA(intrinsic);
}

bool isNvMmaSync(MMAIntrinsic intrinsic) {
  return getArchID(intrinsic) == 0x2000;
}

static bool isNvWmma(MMAIntrinsic intrinsic) {
  return getArchID(intrinsic) == 0x2100;
}

int64_t getIntrinsicSubgroupSize(ScaledMMAIntrinsic intrinsic) {
  switch (intrinsic) {
  case ScaledMMAIntrinsic::MFMA_SCALE_F32_16x16x128_B32:
  case ScaledMMAIntrinsic::MFMA_SCALE_F32_32x32x64_B32:
    return 64;
  }
  assert(false &&
         "all cases should've been handled in ScaledMMA::getBlockSize()");
  return 0;
}

int64_t getIntrinsicSubgroupSize(MMAIntrinsic intrinsic) {
  // Not using Wave64 at all at the moment, so the only place where the
  // subgroup size is 64 is on CDNA* architectures.
  return is_AMD_MFMA(intrinsic) ? 64 : 32;
}

// Returns the operand and accumulator element types required by each MMA
// intrinsic.
static std::tuple<Type, Type, Type> getABCElementTypes(MLIRContext *context,
                                                       MMAIntrinsic intrinsic) {
  Type f8E4M3FNUZ = Float8E4M3FNUZType::get(context);
  Type f8E5M2FNUZ = Float8E5M2FNUZType::get(context);
  Type f8E4M3FN = Float8E4M3FNType::get(context);
  Type f8E5M2 = Float8E5M2Type::get(context);
  Type f16 = Float16Type::get(context);
  Type bf16 = BFloat16Type::get(context);
  Type f32 = Float32Type::get(context);
  Type f64 = Float64Type::get(context);
  Type i8 = IntegerType::get(context, 8);
  Type i32 = IntegerType::get(context, 32);
  switch (intrinsic) {
  case MMAIntrinsic::MFMA_F64_16x16x4_F64:
  case MMAIntrinsic::MFMA_F64_4x4x4x4B_F64:
    return {f64, f64, f64};
  case MMAIntrinsic::MFMA_F32_16x16x4_F32:
  case MMAIntrinsic::MFMA_F32_4x4x1x16B_F32:
  case MMAIntrinsic::MFMA_F32_16x16x1x4B_F32:
  case MMAIntrinsic::MFMA_F32_32x32x1x2B_F32:
  case MMAIntrinsic::WMMA_F32_16x16x4_F32:
    return {f32, f32, f32};
  case MMAIntrinsic::MFMA_F32_4x4x4x16B_F16:
  case MMAIntrinsic::MFMA_F32_16x16x4x4B_F16:
  case MMAIntrinsic::MFMA_F32_32x32x4x2B_F16:
  case MMAIntrinsic::MFMA_F32_16x16x16_F16:
  case MMAIntrinsic::MFMA_F32_32x32x8_F16:
  case MMAIntrinsic::MFMA_F32_16x16x32_F16:
  case MMAIntrinsic::MFMA_F32_32x32x16_F16:
  case MMAIntrinsic::WMMAR3_F32_16x16x16_F16:
  case MMAIntrinsic::WMMAR4_F32_16x16x16_F16:
  case MMAIntrinsic::NV_MMA_SYNC_F32_16x8x16_F16:
  case MMAIntrinsic::NV_WMMA_F32_16x16x16_F16:
  case MMAIntrinsic::WMMA_F32_16x16x32_F16:
    return {f16, f16, f32};
  case MMAIntrinsic::WMMAR3_F16_16x16x16_F16:
  case MMAIntrinsic::WMMAR4_F16_16x16x16_F16:
  case MMAIntrinsic::NV_WMMA_F16_16x16x16_F16:
  case MMAIntrinsic::NV_MMA_SYNC_F16_16x8x16_F16:
  case MMAIntrinsic::WMMA_F16_16x16x32_F16:
    return {f16, f16, f16};
  case MMAIntrinsic::MFMA_F32_4x4x4x16B_BF16:
  case MMAIntrinsic::MFMA_F32_16x16x4x4B_BF16:
  case MMAIntrinsic::MFMA_F32_32x32x4x2B_BF16:
  case MMAIntrinsic::MFMA_F32_16x16x8_BF16:
  case MMAIntrinsic::MFMA_F32_32x32x4_BF16:
  case MMAIntrinsic::MFMA_F32_16x16x16_BF16:
  case MMAIntrinsic::MFMA_F32_32x32x8_BF16:
  case MMAIntrinsic::MFMA_F32_16x16x32_BF16:
  case MMAIntrinsic::MFMA_F32_32x32x16_BF16:
  case MMAIntrinsic::WMMAR3_F32_16x16x16_BF16:
  case MMAIntrinsic::WMMAR4_F32_16x16x16_BF16:
  // NVIDIA BF16 mma.sync multiplies bf16 fragments and accumulates into f32.
  case MMAIntrinsic::NV_MMA_SYNC_F32_16x8x16_BF16:
  case MMAIntrinsic::WMMA_F32_16x16x32_BF16:
    return {bf16, bf16, f32};
  // NVIDIA FP8 mma.sync (m16n8k32, sm_89+).
  case MMAIntrinsic::NV_MMA_SYNC_F32_16x8x32_F8E4M3FN:
    return {f8E4M3FN, f8E4M3FN, f32};
  case MMAIntrinsic::NV_MMA_SYNC_F32_16x8x32_F8E5M2:
    return {f8E5M2, f8E5M2, f32};
  case MMAIntrinsic::WMMAR3_BF16_16x16x16_BF16:
  case MMAIntrinsic::WMMAR4_BF16_16x16x16_BF16:
  case MMAIntrinsic::WMMA_BF16_16x16x32_BF16:
    return {bf16, bf16, bf16};
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E4M3FNUZ:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E4M3FNUZ:
    return {f8E4M3FNUZ, f8E4M3FNUZ, f32};
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E5M2FNUZ:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E5M2FNUZ:
    return {f8E5M2FNUZ, f8E5M2FNUZ, f32};
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E4M3FNUZ_F8E5M2FNUZ:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E4M3FNUZ_F8E5M2FNUZ:
    return {f8E4M3FNUZ, f8E5M2FNUZ, f32};
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E5M2FNUZ_F8E4M3FNUZ:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E5M2FNUZ_F8E4M3FNUZ:
    return {f8E5M2FNUZ, f8E4M3FNUZ, f32};
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E5M2:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E5M2:
  case MMAIntrinsic::MFMA_F32_16x16x128_F8E5M2:
  case MMAIntrinsic::MFMA_F32_32x32x64_F8E5M2:
  case MMAIntrinsic::WMMAR4_F32_16x16x16_F8E5M2:
  case MMAIntrinsic::WMMA_F32_16x16x64_F8E5M2:
  case MMAIntrinsic::WMMA_F32_16x16x128_F8E5M2:
    return {f8E5M2, f8E5M2, f32};
  case MMAIntrinsic::WMMA_F16_16x16x64_F8E5M2:
  case MMAIntrinsic::WMMA_F16_16x16x128_F8E5M2:
    return {f8E5M2, f8E5M2, f16};
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::MFMA_F32_16x16x128_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::MFMA_F32_32x32x64_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::WMMAR4_F32_16x16x16_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::WMMA_F32_16x16x64_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::WMMA_F32_16x16x128_F8E5M2_F8E4M3FN:
    return {f8E5M2, f8E4M3FN, f32};
  case MMAIntrinsic::WMMA_F16_16x16x64_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::WMMA_F16_16x16x128_F8E5M2_F8E4M3FN:
    return {f8E5M2, f8E4M3FN, f16};
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E4M3FN:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E4M3FN:
  case MMAIntrinsic::MFMA_F32_16x16x128_F8E4M3FN:
  case MMAIntrinsic::MFMA_F32_32x32x64_F8E4M3FN:
  case MMAIntrinsic::WMMAR4_F32_16x16x16_F8E4M3FN:
  case MMAIntrinsic::WMMA_F32_16x16x64_F8E4M3FN:
  case MMAIntrinsic::WMMA_F32_16x16x128_F8E4M3FN:
    return {f8E4M3FN, f8E4M3FN, f32};
  case MMAIntrinsic::WMMA_F16_16x16x64_F8E4M3FN:
  case MMAIntrinsic::WMMA_F16_16x16x128_F8E4M3FN:
    return {f8E4M3FN, f8E4M3FN, f16};
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E4M3FN_F8E5M2:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E4M3FN_F8E5M2:
  case MMAIntrinsic::MFMA_F32_16x16x128_F8E4M3FN_F8E5M2:
  case MMAIntrinsic::MFMA_F32_32x32x64_F8E4M3FN_F8E5M2:
  case MMAIntrinsic::WMMAR4_F32_16x16x16_F8E4M3FN_F8E5M2:
  case MMAIntrinsic::WMMA_F32_16x16x64_F8E4M3FN_F8E5M2:
  case MMAIntrinsic::WMMA_F32_16x16x128_F8E4M3FN_F8E5M2:
    return {f8E4M3FN, f8E5M2, f32};
  case MMAIntrinsic::WMMA_F16_16x16x64_F8E4M3FN_F8E5M2:
  case MMAIntrinsic::WMMA_F16_16x16x128_F8E4M3FN_F8E5M2:
    return {f8E4M3FN, f8E5M2, f16};
  case MMAIntrinsic::MFMA_I32_16x16x4x4B_I8:
  case MMAIntrinsic::MFMA_I32_32x32x4x2B_I8:
  case MMAIntrinsic::MFMA_I32_4x4x4x16B_I8:
  case MMAIntrinsic::MFMA_I32_16x16x16_I8:
  case MMAIntrinsic::MFMA_I32_32x32x8_I8:
  case MMAIntrinsic::MFMA_I32_16x16x32_I8:
  case MMAIntrinsic::MFMA_I32_32x32x16_I8:
  case MMAIntrinsic::MFMA_I32_16x16x64_I8:
  case MMAIntrinsic::MFMA_I32_32x32x32_I8:
  case MMAIntrinsic::WMMAR3_I32_16x16x16_I8:
  case MMAIntrinsic::WMMAR4_I32_16x16x16_I8:
  case MMAIntrinsic::WMMA_I32_16x16x64_I8:
    return {i8, i8, i32};
  }
  assert(false && "unexpected enum value");
  return {};
}

/// Returns the MNK shape for an intrinsic without an implemented concrete
/// layout.
static std::tuple<int64_t, int64_t, int64_t>
getUnsupportedMNKShape(MMAIntrinsic intrinsic) {
  switch (intrinsic) {
  case MMAIntrinsic::NV_WMMA_F32_16x16x16_F16:
  case MMAIntrinsic::NV_WMMA_F16_16x16x16_F16:
    return {16, 16, 16};
  default:
    assert(false && "unexpected enum value");
    return {};
  }
  return {};
}

// Returns the per-lane register layout used to map each MMA operand.
MMASingleSubgroupLayout getSingleSubgroupLayout(MMAIntrinsic intrinsic,
                                                int operandIndex) {
  auto mfmaLhs16xK = [](int64_t k) -> MMASingleSubgroupLayout {
    assert(k % 4 == 0 && "doesn't support blocked MFMAs");
    return {/*outer=*/{1, 1}, /*thread=*/{16, 4}, /*tstrides=*/{1, 16},
            /*element=*/{1, k / 4}};
  };
  auto mfmaRhsKx16 = [](int64_t k) -> MMASingleSubgroupLayout {
    assert(k % 4 == 0 && "doesn't support blocked MFMAs");
    return {/*outer=*/{1, 1}, /*thread=*/{4, 16}, /*tstrides=*/{16, 1},
            /*element=*/{k / 4, 1}};
  };

  auto mfmaLhs32xK = [](int64_t k) -> MMASingleSubgroupLayout {
    assert(k % 2 == 0 && "doesn't support blocked MFMAs");
    return {/*outer=*/{1, 1}, /*thread=*/{32, 2}, /*tstrides=*/{1, 32},
            /*element=*/{1, k / 2}};
  };
  auto mfmaRhsKx32 = [](int64_t k) -> MMASingleSubgroupLayout {
    assert(k % 2 == 0 && "doesn't support blocked MFMAs");
    return {/*outer=*/{1, 1}, /*thread=*/{2, 32}, /*tstrides=*/{32, 1},
            /*element=*/{k / 2, 1}};
  };

  // For 4x4 blocked MFMAs (16 blocks).
  // Uses 3D layout with semantic dimensions [Block, M, K] for LHS.
  auto mfmaLhs4xKx16B = [](int64_t k) -> MMASingleSubgroupLayout {
    return {/*outer=*/{1, 1, 1}, /*thread=*/{16, 4, 1}, /*tstrides=*/{4, 1, 64},
            /*element=*/{1, 1, k}};
  };
  // Semantic dimensions [Block, K, N] for RHS.
  auto mfmaRhsKx4x16B = [](int64_t k) -> MMASingleSubgroupLayout {
    return {/*outer=*/{1, 1, 1}, /*thread=*/{16, 1, 4}, /*tstrides=*/{4, 64, 1},
            /*element=*/{1, k, 1}};
  };

  auto mfmaLhs16xKx4B = [](int64_t k) -> MMASingleSubgroupLayout {
    return {/*outer=*/{1, 1, 1}, /*thread=*/{4, 16, 1},
            /*tstrides=*/{16, 1, 64},
            /*element=*/{1, 1, k}};
  };

  auto mfmaRhsKx16x4B = [](int64_t k) -> MMASingleSubgroupLayout {
    return {/*outer=*/{1, 1, 1}, /*thread=*/{4, 1, 16},
            /*tstrides=*/{16, 64, 1},
            /*element=*/{1, k, 1}};
  };

  const MMASingleSubgroupLayout mfmaAcc16x16x4B = {
      /*outer=*/{1, 1, 1}, /*thread=*/{1, 4, 16}, /*tstrides=*/{64, 16, 1},
      /*element=*/{4, 4, 1}};

  auto mfmaLhs32xKx2B = [](int64_t k) -> MMASingleSubgroupLayout {
    return {/*outer=*/{1, 1, 1}, /*thread=*/{2, 32, 1},
            /*tstrides=*/{32, 1, 64},
            /*element=*/{1, 1, k}};
  };

  auto mfmaRhsKx32x2B = [](int64_t k) -> MMASingleSubgroupLayout {
    return {/*outer=*/{1, 1, 1}, /*thread=*/{2, 1, 32},
            /*tstrides=*/{32, 64, 1},
            /*element=*/{1, k, 1}};
  };

  const MMASingleSubgroupLayout mfmaAcc32x32x2B = {
      /*outer=*/{1, 4, 1}, /*thread=*/{1, 2, 32}, /*tstrides=*/{64, 32, 1},
      /*element=*/{2, 4, 1}};

  const MMASingleSubgroupLayout mfmaAcc16x16 = {
      /*outer=*/{1, 1}, /*thread=*/{4, 16}, /*tstrides=*/{16, 1},
      /*element=*/{4, 1}};
  const MMASingleSubgroupLayout mfmaAcc32x32 = {
      /*outer=*/{4, 1}, /*thread=*/{2, 32}, /*tstrides=*/{32, 1},
      /*element=*/{4, 1}};

  const MMASingleSubgroupLayout mfmaAcc4x4x16B = {
      /*outer=*/{1, 1, 1}, /*thread=*/{16, 1, 4}, /*tstrides=*/{4, 64, 1},
      /*element=*/{1, 4, 1}};

  const MMASingleSubgroupLayout mfmaF64Lhs4x4B = {
      /*outer=*/{1, 1, 1}, /*thread=*/{4, 4, 4}, /*tstrides=*/{4, 1, 16},
      /*element=*/{1, 1, 1}};

  const MMASingleSubgroupLayout mfmaF64Rhs4x4B = {
      /*outer=*/{1, 1, 1}, /*thread=*/{4, 4, 4}, /*tstrides=*/{4, 16, 1},
      /*element=*/{1, 1, 1}};

  const MMASingleSubgroupLayout mfmaF64Acc4x4B = {
      /*outer=*/{1, 1, 1}, /*thread=*/{4, 4, 4}, /*tstrides=*/{4, 16, 1},
      /*element=*/{1, 1, 1}};

  // Note: For gfx12, we specify here that, for example with K=16, lane 0 takes
  // A[0, 0..7] and that lane 16 takes A[0, 8..15]. The hardware will internally
  // bounce between the low halves and high halves of lanes every two registers
  // - that is, the value used for A[0, 4] comes out of lane 16's first register
  // in an F16 or BF16 computation. This is noted here in case someone starts
  // chasing some unusual rounding failure or is confused by why the tiling in
  // the manual doesn't *technically* match the below.
  auto gfx12Wmma16xK = [](int64_t k) -> MMASingleSubgroupLayout {
    return {/*outer=*/{1, 1}, /*thread=*/{16, 2}, /*tstrides=*/{1, 16},
            /*element=*/{1, k / 2}};
  };
  auto gfx12WmmaKx16 = [](int64_t k) -> MMASingleSubgroupLayout {
    return {/*outer=*/{1, 1}, /*thread=*/{2, 16}, /*tstrides=*/{16, 1},
            /*element=*/{k / 2, 1}};
  };
  const MMASingleSubgroupLayout gfx12WmmaAcc16x16 = gfx12WmmaKx16(16);

  switch (intrinsic) {
  case MMAIntrinsic::MFMA_F32_16x16x4_F32:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs16xK(4);
    case kMMAOperandRhs:
      return mfmaRhsKx16(4);
    case kMMAOperandAcc:
      return mfmaAcc16x16;
    }
  case MMAIntrinsic::MFMA_F64_4x4x4x4B_F64:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaF64Lhs4x4B;
    case kMMAOperandRhs:
      return mfmaF64Rhs4x4B;
    case kMMAOperandAcc:
      return mfmaF64Acc4x4B;
    }
  // Note: the returned layout for f64 differs than for other MFMAs.
  case MMAIntrinsic::MFMA_F64_16x16x4_F64:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs16xK(4);
    case kMMAOperandRhs:
      return mfmaRhsKx16(4);
    case kMMAOperandAcc:
      return {/*outer=*/{4, 1}, /*thread=*/{4, 16}, /*tstrides=*/{16, 1},
              /*element=*/{1, 1}};
    }
  case MMAIntrinsic::MFMA_F32_16x16x8_BF16: {
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs16xK(8);
    case kMMAOperandRhs:
      return mfmaRhsKx16(8);
    case kMMAOperandAcc:
      return mfmaAcc16x16;
    }
  }
  case MMAIntrinsic::MFMA_F32_32x32x4_BF16:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs32xK(4);
    case kMMAOperandRhs:
      return mfmaRhsKx32(4);
    case kMMAOperandAcc:
      return mfmaAcc32x32;
    }
  case MMAIntrinsic::MFMA_F32_4x4x4x16B_F16:
  case MMAIntrinsic::MFMA_F32_4x4x4x16B_BF16:
  case MMAIntrinsic::MFMA_I32_4x4x4x16B_I8:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs4xKx16B(4);
    case kMMAOperandRhs:
      return mfmaRhsKx4x16B(4);
    case kMMAOperandAcc:
      return mfmaAcc4x4x16B;
    }
  case MMAIntrinsic::MFMA_F32_4x4x1x16B_F32:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs4xKx16B(1);
    case kMMAOperandRhs:
      return mfmaRhsKx4x16B(1);
    case kMMAOperandAcc:
      return mfmaAcc4x4x16B;
    }
  case MMAIntrinsic::MFMA_F32_16x16x4x4B_F16:
  case MMAIntrinsic::MFMA_F32_16x16x4x4B_BF16:
  case MMAIntrinsic::MFMA_I32_16x16x4x4B_I8:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs16xKx4B(4);
    case kMMAOperandRhs:
      return mfmaRhsKx16x4B(4);
    case kMMAOperandAcc:
      return mfmaAcc16x16x4B;
    }
  case MMAIntrinsic::MFMA_F32_16x16x1x4B_F32:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs16xKx4B(1);
    case kMMAOperandRhs:
      return mfmaRhsKx16x4B(1);
    case kMMAOperandAcc:
      return mfmaAcc16x16x4B;
    }
  case MMAIntrinsic::MFMA_F32_32x32x4x2B_F16:
  case MMAIntrinsic::MFMA_F32_32x32x4x2B_BF16:
  case MMAIntrinsic::MFMA_I32_32x32x4x2B_I8:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs32xKx2B(4);
    case kMMAOperandRhs:
      return mfmaRhsKx32x2B(4);
    case kMMAOperandAcc:
      return mfmaAcc32x32x2B;
    }
  case MMAIntrinsic::MFMA_F32_32x32x1x2B_F32:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs32xKx2B(1);
    case kMMAOperandRhs:
      return mfmaRhsKx32x2B(1);
    case kMMAOperandAcc:
      return mfmaAcc32x32x2B;
    }
  case MMAIntrinsic::MFMA_I32_16x16x16_I8:
  case MMAIntrinsic::MFMA_F32_16x16x16_F16:
  case MMAIntrinsic::MFMA_F32_16x16x16_BF16:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs16xK(16);
    case kMMAOperandRhs:
      return mfmaRhsKx16(16);
    case kMMAOperandAcc:
      return mfmaAcc16x16;
    }
  case MMAIntrinsic::MFMA_I32_32x32x8_I8:
  case MMAIntrinsic::MFMA_F32_32x32x8_F16:
  case MMAIntrinsic::MFMA_F32_32x32x8_BF16:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs32xK(8);
    case kMMAOperandRhs:
      return mfmaRhsKx32(8);
    case kMMAOperandAcc:
      return mfmaAcc32x32;
    }
  case MMAIntrinsic::MFMA_F32_16x16x32_F16:
  case MMAIntrinsic::MFMA_F32_16x16x32_BF16:
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E4M3FNUZ:
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E5M2FNUZ:
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E4M3FNUZ_F8E5M2FNUZ:
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E5M2FNUZ_F8E4M3FNUZ:
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E4M3FN:
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E5M2:
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E4M3FN_F8E5M2:
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::MFMA_I32_16x16x32_I8:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs16xK(32);
    case kMMAOperandRhs:
      return mfmaRhsKx16(32);
    case kMMAOperandAcc:
      return mfmaAcc16x16;
    }
  case MMAIntrinsic::MFMA_F32_32x32x16_F16:
  case MMAIntrinsic::MFMA_F32_32x32x16_BF16:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E4M3FNUZ:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E5M2FNUZ:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E4M3FNUZ_F8E5M2FNUZ:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E5M2FNUZ_F8E4M3FNUZ:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E4M3FN:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E5M2:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E4M3FN_F8E5M2:
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::MFMA_I32_32x32x16_I8:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs32xK(16);
    case kMMAOperandRhs:
      return mfmaRhsKx32(16);
    case kMMAOperandAcc:
      return mfmaAcc32x32;
    }
  case MMAIntrinsic::MFMA_I32_16x16x64_I8:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs16xK(64);
    case kMMAOperandRhs:
      return mfmaRhsKx16(64);
    case kMMAOperandAcc:
      return mfmaAcc16x16;
    }
  case MMAIntrinsic::MFMA_I32_32x32x32_I8:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs32xK(32);
    case kMMAOperandRhs:
      return mfmaRhsKx32(32);
    case kMMAOperandAcc:
      return mfmaAcc32x32;
    }
  case MMAIntrinsic::MFMA_F32_16x16x128_F8E5M2:
  case MMAIntrinsic::MFMA_F32_16x16x128_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::MFMA_F32_16x16x128_F8E4M3FN:
  case MMAIntrinsic::MFMA_F32_16x16x128_F8E4M3FN_F8E5M2:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs16xK(128);
    case kMMAOperandRhs:
      return mfmaRhsKx16(128);
    case kMMAOperandAcc:
      return mfmaAcc16x16;
    }
  case MMAIntrinsic::MFMA_F32_32x32x64_F8E5M2:
  case MMAIntrinsic::MFMA_F32_32x32x64_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::MFMA_F32_32x32x64_F8E4M3FN:
  case MMAIntrinsic::MFMA_F32_32x32x64_F8E4M3FN_F8E5M2:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return mfmaLhs32xK(64);
    case kMMAOperandRhs:
      return mfmaRhsKx32(64);
    case kMMAOperandAcc:
      return mfmaAcc32x32;
    }

  case MMAIntrinsic::WMMAR3_F32_16x16x16_F16:
  case MMAIntrinsic::WMMAR3_F32_16x16x16_BF16:
  case MMAIntrinsic::WMMAR3_I32_16x16x16_I8:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return {/*outer=*/{1, 1}, /*thread=*/{16, 1}, /*strides=*/{1, 0},
              /*element=*/{1, 16}};
    case kMMAOperandRhs:
      return {/*outer=*/{1, 1}, /*thread=*/{1, 16}, /*tstrides=*/{0, 1},
              /*element=*/{16, 1}};
    case kMMAOperandAcc:
      return {/*outer=*/{8, 1}, /*thread=*/{2, 16}, /*tstrides=*/{16, 1},
              /*element=*/{1, 1}};
    }
  case MMAIntrinsic::WMMAR3_F16_16x16x16_F16:
  case MMAIntrinsic::WMMAR3_BF16_16x16x16_BF16:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return {/*outer=*/{1, 1}, /*thread=*/{16, 1}, /*strides=*/{1, 0},
              /*element=*/{1, 16}};
    case kMMAOperandRhs:
      return {/*outer=*/{1, 1}, /*thread=*/{1, 16}, /*tstrides=*/{0, 1},
              /*element=*/{16, 1}};
    case kMMAOperandAcc:
      // RDNA3 WMMA f16/bf16 output uses vector<16xf16> but only 8 values are
      // valid (at even indices with subwordOffset=0). The actual lane mapping
      // is the same as the f32 accumulator: 2 thread groups in M, 16 in N.
      return {/*outer=*/{8, 1}, /*thread=*/{2, 16}, /*tstrides=*/{16, 1},
              /*element=*/{1, 1}};
    }
  case MMAIntrinsic::WMMAR4_F32_16x16x16_F16:
  case MMAIntrinsic::WMMAR4_F32_16x16x16_BF16:
  case MMAIntrinsic::WMMAR4_F32_16x16x16_F8E5M2:
  case MMAIntrinsic::WMMAR4_F32_16x16x16_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::WMMAR4_F32_16x16x16_F8E4M3FN:
  case MMAIntrinsic::WMMAR4_F32_16x16x16_F8E4M3FN_F8E5M2:
  case MMAIntrinsic::WMMAR4_I32_16x16x16_I8:
  case MMAIntrinsic::WMMAR4_F16_16x16x16_F16:
  case MMAIntrinsic::WMMAR4_BF16_16x16x16_BF16:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return gfx12Wmma16xK(16);
    case kMMAOperandRhs:
      return gfx12WmmaKx16(16);
    case kMMAOperandAcc:
      return gfx12WmmaAcc16x16;
    }
  case MMAIntrinsic::WMMA_F32_16x16x4_F32:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return gfx12Wmma16xK(4);
    case kMMAOperandRhs:
      return gfx12WmmaKx16(4);
    case kMMAOperandAcc:
      return gfx12WmmaAcc16x16;
    }
  case MMAIntrinsic::WMMA_F32_16x16x32_F16:
  case MMAIntrinsic::WMMA_F32_16x16x32_BF16:
  case MMAIntrinsic::WMMA_F16_16x16x32_F16:
  case MMAIntrinsic::WMMA_BF16_16x16x32_BF16:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return gfx12Wmma16xK(32);
    case kMMAOperandRhs:
      return gfx12WmmaKx16(32);
    case kMMAOperandAcc:
      return gfx12WmmaAcc16x16;
    }
  case MMAIntrinsic::WMMA_F32_16x16x64_F8E4M3FN:
  case MMAIntrinsic::WMMA_F32_16x16x64_F8E4M3FN_F8E5M2:
  case MMAIntrinsic::WMMA_F32_16x16x64_F8E5M2:
  case MMAIntrinsic::WMMA_F32_16x16x64_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::WMMA_F16_16x16x64_F8E4M3FN:
  case MMAIntrinsic::WMMA_F16_16x16x64_F8E4M3FN_F8E5M2:
  case MMAIntrinsic::WMMA_F16_16x16x64_F8E5M2:
  case MMAIntrinsic::WMMA_F16_16x16x64_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::WMMA_I32_16x16x64_I8:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return gfx12Wmma16xK(64);
    case kMMAOperandRhs:
      return gfx12WmmaKx16(64);
    case kMMAOperandAcc:
      return gfx12WmmaAcc16x16;
    }
  case MMAIntrinsic::WMMA_F32_16x16x128_F8E5M2:
  case MMAIntrinsic::WMMA_F32_16x16x128_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::WMMA_F32_16x16x128_F8E4M3FN:
  case MMAIntrinsic::WMMA_F32_16x16x128_F8E4M3FN_F8E5M2:
  case MMAIntrinsic::WMMA_F16_16x16x128_F8E5M2:
  case MMAIntrinsic::WMMA_F16_16x16x128_F8E5M2_F8E4M3FN:
  case MMAIntrinsic::WMMA_F16_16x16x128_F8E4M3FN:
  case MMAIntrinsic::WMMA_F16_16x16x128_F8E4M3FN_F8E5M2:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return gfx12Wmma16xK(128);
    case kMMAOperandRhs:
      return gfx12WmmaKx16(128);
    case kMMAOperandAcc:
      return gfx12WmmaAcc16x16;
    }
  case MMAIntrinsic::NV_MMA_SYNC_F32_16x8x16_F16:
  case MMAIntrinsic::NV_MMA_SYNC_F16_16x8x16_F16:
  // BF16 mma.sync uses the same 16x8x16 register layout as the f16 variants.
  case MMAIntrinsic::NV_MMA_SYNC_F32_16x8x16_BF16:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return {/*outer=*/{2, 2}, /*thread=*/{8, 4}, /*strides=*/{4, 1},
              /*element=*/{1, 2}};
    case kMMAOperandRhs:
      return {/*outer=*/{2, 1}, /*thread=*/{4, 8}, /*strides=*/{1, 4},
              /*element=*/{2, 1}};
    case kMMAOperandAcc:
      return {/*outer=*/{2, 1}, /*thread=*/{8, 4}, /*strides=*/{4, 1},
              /*element=*/{1, 2}};
    default:
      return {};
    }
  // FP8 mma.sync (m16n8k32, sm_89+). Same thread arrangement as the f16
  // m16n8k16 variants; K doubles so element[K] doubles from 2 to 4 (4 FP8
  // values packed per 32-bit register vs 2 FP16). The accumulator shape is
  // unchanged (same M×N output). Both supported FP8 types share the same
  // register layout.
  case MMAIntrinsic::NV_MMA_SYNC_F32_16x8x32_F8E4M3FN:
  case MMAIntrinsic::NV_MMA_SYNC_F32_16x8x32_F8E5M2:
    switch (operandIndex) {
    case kMMAOperandLhs: // A [M=16, K=32]: 16 elements/thread
      return {/*outer=*/{2, 2}, /*thread=*/{8, 4}, /*strides=*/{4, 1},
              /*element=*/{1, 4}};
    case kMMAOperandRhs: // B [K=32, N=8]: 8 elements/thread
      return {/*outer=*/{2, 1}, /*thread=*/{4, 8}, /*strides=*/{1, 4},
              /*element=*/{4, 1}};
    case kMMAOperandAcc: // C [M=16, N=8]: 4 elements/thread (same as f16)
      return {/*outer=*/{2, 1}, /*thread=*/{8, 4}, /*strides=*/{4, 1},
              /*element=*/{1, 2}};
    default:
      return {};
    }
  case MMAIntrinsic::NV_WMMA_F32_16x16x16_F16:
  case MMAIntrinsic::NV_WMMA_F16_16x16x16_F16:
    return {};
  }
  assert(false && "unexpected enum value");
  return {};
}

MMASingleSubgroupLayout getSingleSubgroupLayout(MMAIntrinsic intrinsic,
                                                int operandIndex,
                                                bool colMajor) {
  MMASingleSubgroupLayout baseLayout =
      getSingleSubgroupLayout(intrinsic, operandIndex);
  assert((baseLayout.element.size() == 2 || baseLayout.element.size() == 3) &&
         "expected 2d or 3d layout");
  if (colMajor) {
    std::swap(baseLayout.element[0], baseLayout.element[1]);
    std::swap(baseLayout.thread[0], baseLayout.thread[1]);
    std::swap(baseLayout.outer[0], baseLayout.outer[1]);
    std::swap(baseLayout.tstrides[0], baseLayout.tstrides[1]);
  }
  return baseLayout;
}

MMASingleSubgroupLayout
getSingleSubgroupLayout(VirtualMMAIntrinsic virtualIntrinsic, int operandIndex,
                        bool colMajor) {
  MMASingleSubgroupLayout baseLayout =
      getSingleSubgroupLayout(virtualIntrinsic, operandIndex);
  assert(baseLayout.element.size() == 2 && "expected 2d layout");
  if (colMajor) {
    std::swap(baseLayout.element[0], baseLayout.element[1]);
    std::swap(baseLayout.thread[0], baseLayout.thread[1]);
    std::swap(baseLayout.outer[0], baseLayout.outer[1]);
    std::swap(baseLayout.tstrides[0], baseLayout.tstrides[1]);
  }
  return baseLayout;
}

// Struct describing the shape of a MMA operation, but not the detailed layout.
struct OpaqueMmaLayout {
  int64_t bSize = 0; // 0 for non-block intrinsics
  int64_t mSize = 0;
  int64_t nSize = 0;
  int64_t kSize = 0;
  Type aType;
  Type bType;
  Type cType;
};

static std::tuple<int64_t, int64_t, int64_t>
getMNKShapeFromIntrinsic(MMAIntrinsic intrinsic) {
  if (intrinsic == MMAIntrinsic::NV_WMMA_F32_16x16x16_F16 ||
      intrinsic == MMAIntrinsic::NV_WMMA_F16_16x16x16_F16) {
    return getUnsupportedMNKShape(intrinsic);
  }
  auto lhs = getSingleSubgroupLayout(intrinsic, kMMAOperandLhs);
  auto rhs = getSingleSubgroupLayout(intrinsic, kMMAOperandRhs);

  // For 3D layouts with batch/block dimension, skip index 0 (batch)
  // and use indices 1,2 for M/N/K computation
  bool has3DLayout = lhs.outer.size() == 3;
  if (has3DLayout) {
    return {lhs.outer[1] * lhs.thread[1] * lhs.element[1],
            rhs.outer[2] * rhs.thread[2] * rhs.element[2],
            lhs.outer[2] * lhs.thread[2] * lhs.element[2]};
  }

  return {lhs.outer[0] * lhs.thread[0] * lhs.element[0],
          rhs.outer[1] * rhs.thread[1] * rhs.element[1],
          lhs.outer[1] * lhs.thread[1] * lhs.element[1]};
}

int64_t getMSize(MMAIntrinsic intrinsic) {
  return std::get<0>(getMNKShapeFromIntrinsic(intrinsic));
}
int64_t getNSize(MMAIntrinsic intrinsic) {
  return std::get<1>(getMNKShapeFromIntrinsic(intrinsic));
}
int64_t getKSize(MMAIntrinsic intrinsic) {
  return std::get<2>(getMNKShapeFromIntrinsic(intrinsic));
}

static OpaqueMmaLayout getOpaqueMMALayout(MLIRContext *context,
                                          MMAIntrinsic intrinsic) {
  OpaqueMmaLayout o;
  std::tie(o.aType, o.bType, o.cType) = getABCElementTypes(context, intrinsic);
  std::tie(o.mSize, o.nSize, o.kSize) = getMNKShapeFromIntrinsic(intrinsic);
  auto lhs = getSingleSubgroupLayout(intrinsic, kMMAOperandLhs);
  if (lhs.outer.size() == 3) {
    o.bSize = lhs.outer[0] * lhs.thread[0] * lhs.element[0];
  }
  return o;
}

MMASingleSubgroupLayout
getSingleSubgroupLayout(IREE::Codegen::InnerTileDescAttrInterface mmaKind,
                        int operandIndex) {
  if (auto mmaAttr = dyn_cast<MMAAttr>(mmaKind)) {
    // |colMajor| indicates that the accumulator layout should be returned
    // column major.
    return IREE::GPU::getSingleSubgroupLayout(
        mmaAttr.getIntrinsic(), operandIndex,
        operandIndex == kMMAOperandAcc && mmaAttr.getColMajor());
  }
  if (auto vmmaAttr = dyn_cast<VirtualMMAAttr>(mmaKind)) {
    return IREE::GPU::getSingleSubgroupLayout(
        vmmaAttr.getIntrinsic(), operandIndex,
        operandIndex == kMMAOperandAcc && vmmaAttr.getColMajor());
  }
  if (auto smmaAttr = dyn_cast<ScaledMMAAttr>(mmaKind)) {
    return IREE::GPU::getSingleSubgroupLayout(
        smmaAttr.getIntrinsic(), operandIndex,
        operandIndex == kScaledMMAOperandAcc && smmaAttr.getColMajor());
  }
  assert(false && "unhandled MMA Interface type.");
  return {};
}

void IREE::GPU::InnerTiledSemanticsAttr::getTileTypes(
    IREE::Codegen::InnerTileDescAttrInterface kind,
    llvm::SmallVectorImpl<mlir::VectorType> &result) const {
  if (getDistributed()) {
    kind.getDistributedTileTypes(result);
  } else {
    kind.getUndistributedTileTypes(result);
  }
}

//===----------------------------------------------------------------------===//
// MMA Attributes
//===----------------------------------------------------------------------===//

MMAAttr MMAAttr::get(MLIRContext *context, MMAIntrinsic type) {
  return Base::get(context, type, /*colMajor=*/false);
}

int64_t MMAAttr::getExpectedNumInputs() const { return 2; }

int64_t MMAAttr::getExpectedNumOutputs() const { return 1; }

LogicalResult MMAAttr::verifyIndexingMaps(ArrayRef<AffineMap> maps) const {
  return verifyMmaIndexingMaps(maps);
}

void MMAAttr::getUndistributedTileTypes(
    SmallVectorImpl<VectorType> &result) const {
  MLIRContext *ctx = getContext();
  OpaqueMmaLayout o = getOpaqueMMALayout(ctx, getIntrinsic());
  if (o.bSize) {
    result.assign({VectorType::get({o.bSize, o.mSize, o.kSize}, o.aType),
                   VectorType::get({o.bSize, o.kSize, o.nSize}, o.bType),
                   VectorType::get({o.bSize, o.mSize, o.nSize}, o.cType)});
  } else {
    result.assign({VectorType::get({o.mSize, o.kSize}, o.aType),
                   VectorType::get({o.kSize, o.nSize}, o.bType),
                   VectorType::get({o.mSize, o.nSize}, o.cType)});
  }
}

template <typename MMAIntrinsicType>
static VectorType getThreadVectorType(MLIRContext *context,
                                      MMAIntrinsicType intrinsic,
                                      int operandIndex) {
  auto o = getOpaqueMMALayout(context, intrinsic);
  auto s = getSingleSubgroupLayout(intrinsic, operandIndex);
  Type elemType = isIntrinsicLhs<MMAIntrinsicType>(operandIndex)   ? o.aType
                  : isIntrinsicRhs<MMAIntrinsicType>(operandIndex) ? o.bType
                                                                   : o.cType;
  if constexpr (std::is_same_v<MMAIntrinsicType, MMAIntrinsic>) {
    if (isNvMmaSync(intrinsic)) {
      return VectorType::get(
          {s.outer[0] * s.outer[1], s.element[0] * s.element[1]}, elemType);
    }
  }
  if (s.outer.size() == 3) {
    return VectorType::get({s.outer[0] * s.element[0] * s.outer[1] *
                            s.element[1] * s.outer[2] * s.element[2]},
                           elemType);
  }
  return VectorType::get(
      {s.element[0] * s.element[1] * s.outer[0] * s.outer[1]}, elemType);
}

void MMAAttr::getDistributedTileTypes(
    SmallVectorImpl<VectorType> &result) const {
  MLIRContext *context = getContext();
  MMAIntrinsic intrinsic = getIntrinsic();
  result.assign({getThreadVectorType(context, intrinsic, kMMAOperandLhs),
                 getThreadVectorType(context, intrinsic, kMMAOperandRhs),
                 getThreadVectorType(context, intrinsic, kMMAOperandAcc)});
}

std::optional<SmallVector<int64_t, 2>>
MMAAttr::getUndistributedTileDimExpansion(int64_t operandIndex,
                                          int64_t dim) const {
  assert(operandIndex <= 2 && "invalid operand index");
  assert(dim < 3 && "invalid inner tile dim");
  MMASingleSubgroupLayout layout =
      getSingleSubgroupLayout(*this, static_cast<int>(operandIndex));
  if (layout.outer[dim] > 1) {
    return SmallVector<int64_t, 2>{layout.outer[dim],
                                   layout.element[dim] * layout.thread[dim]};
  }
  return std::nullopt;
}

int64_t MMAAttr::getBlockSize() const {
  return IREE::GPU::getBlockSize(getIntrinsic());
}

bool MMAAttr::isBlockIntrinsic() const {
  MMASingleSubgroupLayout lhs =
      getSingleSubgroupLayout(getIntrinsic(), kMMAOperandLhs);
  return lhs.outer.size() == 3;
}

int64_t MMAAttr::getSubgroupSize() const {
  return getIntrinsicSubgroupSize(getIntrinsic());
}

SmallVector<SmallVector<utils::IteratorType>>
MMAAttr::getOperandIteratorTypes() const {
  return {{utils::IteratorType::parallel, utils::IteratorType::reduction},
          {utils::IteratorType::reduction, utils::IteratorType::parallel},
          {utils::IteratorType::parallel, utils::IteratorType::parallel}};
}

Attribute MMAAttr::getDistributionMappingKind() const {
  // Explicit distribution currently unsupported for NV WMMA intrinsics.
  MMAIntrinsic intrinsic = getIntrinsic();
  if (intrinsic == MMAIntrinsic::NV_WMMA_F16_16x16x16_F16 ||
      intrinsic == MMAIntrinsic::NV_WMMA_F32_16x16x16_F16) {
    return Attribute();
  }
  return IREE::GPU::LaneIdAttr::get(getContext(), 0);
}

OpFoldResult MMAAttr::getDistributionWorkerCount(OpBuilder &, Location,
                                                 Operation *) const {
  if (!getDistributionMappingKind()) {
    return OpFoldResult();
  }
  return getAsIndexOpFoldResult(getContext(), getSubgroupSize());
}

// Returns virtual intrinsics that are composed from this concrete MMA op.
SmallVector<VirtualMMAIntrinsic> MMAAttr::getVirtualIntrinsics() const {
  switch (getIntrinsic()) {
  case MMAIntrinsic::MFMA_F32_16x16x16_F16:
    return {VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F16,
            VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_F16};
  case MMAIntrinsic::MFMA_F32_32x32x8_F16:
    return {VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F16};
  case MMAIntrinsic::MFMA_F32_16x16x16_BF16:
    return {VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_BF16};
  case MMAIntrinsic::MFMA_F32_16x16x32_F8E4M3FNUZ:
    return {VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F8E4M3FNUZ};
  case MMAIntrinsic::MFMA_F32_32x32x16_F8E4M3FNUZ:
    return {VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F8E4M3FNUZ};
  default:
    return {};
  }
}

// Creates the target-specific MMA operation for the selected intrinsic layout.
static Value createMmaOp(OpBuilder &builder, Location loc,
                         MMAIntrinsic intrinsic, Type resultType, Value lhs,
                         Value rhs, Value acc, bool colMajor = false) {
  auto getVecOrSingleElem = [&](Value vec) -> Value {
    bool one = cast<VectorType>(vec.getType()).getNumElements() == 1;
    return one ? vector::ExtractOp::create(builder, loc, vec, 0) : vec;
  };
  auto layout = getOpaqueMMALayout(builder.getContext(), intrinsic);
  if (is_AMD_MFMA(intrinsic)) {
    // MFMA intrinsics want single-element operands of element type, not vector.
    lhs = getVecOrSingleElem(lhs);
    rhs = getVecOrSingleElem(rhs);

    // Because the thread layout of the lhs and rhs are transpositions of one
    // another for all MFMA variants, to produce a column major result we can
    // simply swap the operands to the MFMA.
    if (colMajor) {
      std::swap(lhs, rhs);
    }
    acc = getVecOrSingleElem(acc);
    Type vecOrSingleAccType = acc.getType();
    Value mfmaResult =
        amdgpu::MFMAOp::create(builder, loc, vecOrSingleAccType, layout.mSize,
                               layout.nSize, layout.kSize,
                               getBlockSize(intrinsic), lhs, rhs, acc)
            .getResult();
    if (vecOrSingleAccType != resultType) {
      // If the mfma produced a scalar (single-element accumulator), broadcast
      // it to the original type.
      mfmaResult =
          vector::BroadcastOp::create(builder, loc, resultType, mfmaResult);
    }
    return mfmaResult;
  }
  if (is_AMD_WMMA(intrinsic)) {
    // As with MFMA, the thread layouts of the lhs and rhs are transpositions
    // of one another for all WMMA variants, so swapping produces a column
    // major result.
    if (colMajor) {
      std::swap(lhs, rhs);
    }
    // RDNA3 WMMA f16/bf16: the accumulator layout is outer={8,1} (8 logical
    // elements per lane), but the hardware instruction operates on 8 VGPRs
    // (vector<16xf16>) with valid values at even indices (opsel=0). We widen
    // the accumulator to vector<16xf16> here so the LLVM backend gets the
    // native type and generates correct code in reduction loops.

    // TODO: even though this interleave/deinterleave is needed for correctness,
    // we're wasting half the available registers by using the instruction this
    // way. Currently, we're only placing accumulator results in an even index,
    // keeping the high halves set to 0. However, the WMMA instruction allows us
    // to write to both the even or odd halves of each register, so we could
    // combine intrinsics into the even/odd halves of a single accumulator and
    // then deinterleave after the loop (and maybe after all reductions).
    bool isRDNA3HalfAcc =
        (intrinsic == MMAIntrinsic::WMMAR3_F16_16x16x16_F16 ||
         intrinsic == MMAIntrinsic::WMMAR3_BF16_16x16x16_BF16);
    if (isRDNA3HalfAcc) {
      Type elemTy = cast<VectorType>(acc.getType()).getElementType();
      auto halfVecTy = VectorType::get({8}, elemTy);
      auto hwVecTy = VectorType::get({16}, elemTy);
      Value hwAcc =
          // Interleave the half-width accumulator with zeros to fill the
          // full-width hardware register; the deinterleave will be hoisted
          // out of the reduction loop.
          IREE::Util::HoistableConversionOp::create(
              builder, loc, /*tag=*/kRdna3InterleaveAcc,
              /*inverseTag=*/kRdna3DeinterleaveAcc, acc,
              [halfVecTy, elemTy](OpBuilder &b, Location loc, ValueRange args) {
                Value zero = arith::ConstantOp::create(
                    b, loc,
                    SplatElementsAttr::get(halfVecTy, b.getZeroAttr(elemTy)));
                return SmallVector<Value>{
                    vector::InterleaveOp::create(b, loc, args[0], zero)};
              })
              .getResult(0);
      Value hwResult =
          amdgpu::WMMAOp::create(builder, loc, hwVecTy, layout.mSize,
                                 layout.nSize, layout.kSize, lhs, rhs, hwAcc)
              .getResult();
      Value result =
          IREE::Util::HoistableConversionOp::create(
              builder, loc, /*tag=*/kRdna3DeinterleaveAcc,
              /*inverseTag=*/kRdna3InterleaveAcc, hwResult,
              [](OpBuilder &b, Location loc, ValueRange args) {
                return SmallVector<Value>{
                    vector::DeinterleaveOp::create(b, loc, args[0]).getRes1()};
              })
              .getResult(0);
      return result;
    }
    return amdgpu::WMMAOp::create(builder, loc, resultType, layout.mSize,
                                  layout.nSize, layout.kSize, lhs, rhs, acc)
        .getResult();
  }
  if (isNvMmaSync(intrinsic)) {
    // Transpose the two outer dimensions to model the column-major register
    // ordering expected by mma.sync. The input shape differs between pipelines:
    // VectorDistribute produces 2x2x1x2, TileAndFuse produces 2x1x2x2.
    // Remove the unit dimension to simplify the transpose.
    // Preserve the operand element type so BF16 fragments are not reshaped as
    // f16.
    Type elementType = cast<VectorType>(lhs.getType()).getElementType();
    auto nonUnitVecType = VectorType::get({2, 2, 2}, elementType);
    auto reshaped =
        vector::ShapeCastOp::create(builder, loc, nonUnitVecType, lhs);
    auto permAttr = builder.getDenseI64ArrayAttr({1, 0, 2});
    auto transposed = vector::TransposeOp::create(builder, loc, nonUnitVecType,
                                                  reshaped, permAttr);
    lhs = vector::ShapeCastOp::create(builder, loc, lhs.getType(), transposed);
    Attribute mmaShape[] = {builder.getI64IntegerAttr(layout.mSize),
                            builder.getI64IntegerAttr(layout.nSize),
                            builder.getI64IntegerAttr(layout.kSize)};
    ArrayAttr mmShapeAttr = builder.getArrayAttr(mmaShape);

    return nvgpu::MmaSyncOp::create(builder, loc, lhs, rhs, acc, mmShapeAttr)
        .getResult();
  }
  return {};
}

// Generates amdgpu.mfma/wmma operation on the given inputs for this attribute
// type.
LogicalResult
MMAAttr::buildUnderlyingOperations(OpBuilder &builder, Location loc,
                                   ValueRange inputs, ValueRange outputs,
                                   SmallVectorImpl<Value> &results) const {
  if (inputs.size() != 2) {
    return failure();
  }
  if (outputs.size() != 1) {
    return failure();
  }
  SmallVector<VectorType> threadTypes;
  getDistributedTileTypes(threadTypes);
  if (!llvm::equal(threadTypes,
                   llvm::concat<Type>(inputs.getTypes(), outputs.getTypes()))) {
    return failure();
  }

  if (Value value =
          createMmaOp(builder, loc, getIntrinsic(), outputs[0].getType(),
                      inputs[0], inputs[1], outputs[0], getColMajor())) {
    results.push_back(value);
    return success();
  }
  return failure();
}

/// Creates index_hint ops wrapping delinearized lane ID values.
/// The `delinearizedLaneId` values come from delinearizing the lane ID using
/// `basis`, with the innermost/fastest-varying dimension last.
///
/// Non-final indices get lane_constant hints (uniform across lane groups).
/// The final index gets lane_increment hint (increments within lane group).
/// The group size is derived from the innermost basis element.
/// Indices with a unit basis are ignored, and given a lane_constant hint.
static SmallVector<Value>
createTransposeLoadIndexHint(OpBuilder &builder, Location loc,
                             ValueRange delinearizedLaneId,
                             ArrayRef<int64_t> basis) {
  // Need at least 2 dimensions for transpose load pattern.
  if (delinearizedLaneId.size() < 2) {
    return SmallVector<Value>(delinearizedLaneId.begin(),
                              delinearizedLaneId.end());
  }

  // Find the index of the innermost non-unit (> 1) basis element.
  // This determines which result gets the lane-increment hint.
  // Size-1 dimensions produce constant 0 outputs regardless of lane ID,
  // so they don't contribute to the meaningful group structure.
  int64_t groupSize = 1;
  size_t incrementResultIdx = delinearizedLaneId.size() - 1;
  // The delinearized indices could have N or N + 1 results, and the basis
  // elements are aligned with the last N results, so iterate backwards
  // together.
  for (size_t i = 1; i <= basis.size(); ++i) {
    groupSize = basis[basis.size() - i];
    incrementResultIdx = delinearizedLaneId.size() - i;
    if (groupSize > 1) {
      break;
    }
  }

  auto laneConstantAttr =
      IREE::GPU::LaneConstantAttr::get(builder.getContext(), groupSize);
  auto laneIncrementAttr = IREE::GPU::LaneIncrementAttr::get(
      builder.getContext(), groupSize, /*step=*/1, /*aligned=*/true);

  SmallVector<Value> results;
  for (auto [i, value] : llvm::enumerate(delinearizedLaneId)) {
    // The result corresponding to innermost non-unit basis gets lane-increment;
    // all other results get lane-constant hints.
    Attribute hint = (i == incrementResultIdx) ? Attribute(laneIncrementAttr)
                                               : Attribute(laneConstantAttr);
    auto hintOp = IREE::Codegen::IndexHintOp::create(builder, loc, value, hint);
    results.push_back(hintOp.getResult());
  }

  return results;
}

static LogicalResult populateCanonicalOffsetsSizesAndStrides(
    OpBuilder &builder, Location loc, Value laneId,
    ArrayRef<int64_t> permutation, MMASingleSubgroupLayout subgroupLayout,
    SmallVectorImpl<OpFoldResult> &canonicalOffsets,
    SmallVectorImpl<OpFoldResult> &canonicalSizes,
    SmallVectorImpl<OpFoldResult> &canonicalStrides,
    int64_t physicalLanesPerThread = 1) {
  assert(physicalLanesPerThread >= 1 &&
         "physicalLanesPerThread must be at least 1");
  SmallVector<int64_t> rankReducedShape;

  for (auto [outer, thread, element] :
       llvm::zip_equal(subgroupLayout.outer, subgroupLayout.thread,
                       subgroupLayout.element)) {
    if (outer != 1) {
      rankReducedShape.push_back(outer);
    }
    rankReducedShape.push_back(thread * element);
  }

  if (permutation.size() != rankReducedShape.size()) {
    return failure();
  }

  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  Value cZero = arith::ConstantIndexOp::create(builder, loc, 0);
  canonicalStrides.append(rankReducedShape.size(), one);

  SmallVector<Value> vtids;
  SmallVector<int64_t> vtidBasis;
  SmallVector<size_t> dimToVtid;
  if (failed(basisFromSizesStrides(subgroupLayout.thread,
                                   subgroupLayout.tstrides, vtidBasis,
                                   dimToVtid))) {
    return failure();
  }
  auto splitLaneId = affine::AffineDelinearizeIndexOp::create(
      builder, loc, laneId, vtidBasis, /*hasOuterBound=*/false);

  // Wrap delinearize results with index_hint ops for transpose load.
  // The delinearize results are already in the correct order
  // (innermost/fastest-varying dimension is last).
  SmallVector<Value> hintedSplitLaneId = createTransposeLoadIndexHint(
      builder, loc, splitLaneId.getResults(), vtidBasis);

  // Find the unique element dimension eligible for lane differentiation:
  // exactly one dimension with element[dim] > 1 and divisible by
  // physicalLanesPerThread.
  std::optional<size_t> splitDim;
  if (physicalLanesPerThread > 1) {
    for (auto [dimIdx, element] : llvm::enumerate(subgroupLayout.element)) {
      if (element > 1 && element % physicalLanesPerThread == 0) {
        if (splitDim) {
          splitDim = std::nullopt;
          break;
        }
        splitDim = dimIdx;
      }
    }
  }

  // Build laneGroupIndex from unused delinearize results. Unused results
  // (not mapped by dimToVtid) represent grouped lanes that see duplicate
  // data. When a splitDim exists, we use the single unused component as an
  // index in [0, physicalLanesPerThread) to differentiate those lanes.
  Value laneGroupIndex;
  if (splitDim) {
    llvm::SmallDenseSet<size_t, 4> usedResults(dimToVtid.begin(),
                                               dimToVtid.end());
    // Find the single unused delinearize result representing grouped lanes.
    std::optional<size_t> unusedResultIdx;
    for (size_t i = 1, e = vtidBasis.size(); i <= e; ++i) {
      if (!usedResults.contains(i)) {
        assert(!unusedResultIdx && "expected exactly one unused basis entry");
        unusedResultIdx = i;
      }
    }
    assert(unusedResultIdx && "expected one unused basis entry for lane group");
    assert(vtidBasis[*unusedResultIdx - 1] == physicalLanesPerThread &&
           "unused basis size must equal physicalLanesPerThread");
    laneGroupIndex = splitLaneId.getResult(*unusedResultIdx);
  }

  // Each thread grabs `element` contiguous data, so the vtid needs to be
  // multiplied by `element` to get the next bunch of data.
  // vtid: virtual thread id
  // tid: lane id
  // vtid = ((tid floordiv stride_i) mod size_i) * element_i.
  //
  // Instead of computing those maps, we use one big `delinearize` expression
  // in order to prevent unwanted "simplifications" on affine maps that
  // worsen the generated code quality.
  //
  // When physicalLanesPerThread > 1, the splitDim offset also incorporates
  // laneGroupIndex so each grouped lane gets a disjoint slice instead.
  for (auto [dimIdx, vtidAndElement] :
       llvm::enumerate(llvm::zip_equal(dimToVtid, subgroupLayout.element))) {
    auto [splitResultIdx, element] = vtidAndElement;
    Value vtid = hintedSplitLaneId[splitResultIdx];
    int64_t vtidLen = vtidBasis[splitResultIdx - 1];

    if (splitDim && dimIdx == *splitDim) {
      // offset = vtid * element + laneGroupIndex * perLaneElement.
      int64_t perLaneElement = element / physicalLanesPerThread;
      vtid = affine::AffineLinearizeIndexOp::create(
          builder, loc, ValueRange{vtid, laneGroupIndex, cZero},
          ArrayRef<int64_t>{vtidLen, physicalLanesPerThread, perLaneElement},
          /*disjoint=*/true);
    } else if (element != 1) {
      vtid = affine::AffineLinearizeIndexOp::create(
          builder, loc, ValueRange{vtid, cZero},
          ArrayRef<int64_t>{vtidLen, element},
          /*disjoint=*/true);
    }
    vtids.push_back(vtid);
  }

  int64_t vtidIdx = 0;
  for (auto [dimIdx, elementAndOuter] : llvm::enumerate(
           llvm::zip_equal(subgroupLayout.element, subgroupLayout.outer))) {
    auto [element, outer] = elementAndOuter;
    int64_t perLaneElement = element;
    if (splitDim && dimIdx == *splitDim) {
      perLaneElement = element / physicalLanesPerThread;
    }
    if (outer != 1) {
      canonicalSizes.push_back(builder.getIndexAttr(outer));
      canonicalOffsets.push_back(zero);
    }
    canonicalSizes.push_back(builder.getIndexAttr(perLaneElement));
    canonicalOffsets.push_back(vtids[vtidIdx++]);
  }
  canonicalOffsets.assign(applyPermutation(canonicalOffsets, permutation));
  canonicalSizes.assign(applyPermutation(canonicalSizes, permutation));
  return success();
}

LogicalResult MMAAttr::populateOperandOffsetsSizesStrides(
    OpBuilder &builder, Location loc, uint32_t operandIndex, Value laneId,
    ArrayRef<int64_t> permutation, SmallVectorImpl<OpFoldResult> &offsets,
    SmallVectorImpl<OpFoldResult> &sizes,
    SmallVectorImpl<OpFoldResult> &strides) const {
  assert(operandIndex <= 2 && "Must index valid MMA operand");
  MMASingleSubgroupLayout subgroupLayout =
      getSingleSubgroupLayout(getIntrinsic(), operandIndex,
                              operandIndex == kMMAOperandAcc && getColMajor());
  SmallVector<OpFoldResult> canonicalOffsets;
  SmallVector<OpFoldResult> canonicalSizes;
  if (failed(populateCanonicalOffsetsSizesAndStrides(
          builder, loc, laneId, permutation, subgroupLayout, canonicalOffsets,
          canonicalSizes, strides))) {
    return failure();
  }
  offsets.append(canonicalOffsets);
  sizes.append(canonicalSizes);

  return success();
}

//===----------------------------------------------------------------------===//
// DataTiledMMA Attributes
//===----------------------------------------------------------------------===//

int64_t DataTiledMMAAttr::getExpectedNumInputs() const { return 2; }

int64_t DataTiledMMAAttr::getExpectedNumOutputs() const { return 1; }

LogicalResult DataTiledMMAAttr::populateOperandOffsetsSizesStrides(
    OpBuilder &builder, Location loc, uint32_t operandIndex, Value laneId,
    ArrayRef<int64_t> permutation, SmallVectorImpl<OpFoldResult> &offsets,
    SmallVectorImpl<OpFoldResult> &sizes,
    SmallVectorImpl<OpFoldResult> &strides) const {
  return cast<DataTiledMMAInterfaceAttr>(*this)
      .populateOperandOffsetsSizesStrides(builder, loc, operandIndex, laneId,
                                          permutation, offsets, sizes, strides);
}

LogicalResult
DataTiledMMAAttr::verifyIndexingMaps(ArrayRef<AffineMap> maps) const {
  return verifyMmaIndexingMaps(maps);
}

int64_t DataTiledMMAAttr::getSubgroupSize() const {
  return getIntrinsicSubgroupSize(getIntrinsic());
}

int64_t DataTiledMMAAttr::getFlatWorkgroupSize() const {
  return getSubgroupSize() * getSubgroupsM() * getSubgroupsN() *
         getSubgroupsK();
}

SmallVector<SmallVector<utils::IteratorType>>
DataTiledMMAAttr::getOperandIteratorTypes() const {
  return {{utils::IteratorType::parallel, utils::IteratorType::reduction},
          {utils::IteratorType::reduction, utils::IteratorType::parallel},
          {utils::IteratorType::parallel, utils::IteratorType::parallel}};
}

LogicalResult DataTiledMMAAttr::buildUnderlyingOperations(
    OpBuilder &builder, Location loc, ValueRange inputs, ValueRange outputs,
    SmallVectorImpl<Value> &results) const {
  SmallVector<VectorType> regTypes;
  getDistributedTileTypes(regTypes);
  if (inputs.size() != 2 || outputs.size() != 1 ||
      !llvm::equal(regTypes,
                   llvm::concat<Type>(inputs.getTypes(), outputs.getTypes()))) {
    return failure();
  }

  MMAIntrinsic intrinsic = getIntrinsic();
  auto emitIntrinsic = [&](OpBuilder &b, Location loc, Value lhs, Value rhs,
                           Value acc) -> Value {
    // For GPU, the per-intrinsic acc type matches `getThreadVectorType` for
    // the ACC operand; createMmaOp uses it to broadcast a scalar mfma result
    // back to the per-intrinsic vector if the acc happens to be 1-element.
    return createMmaOp(b, loc, intrinsic, acc.getType(), lhs, rhs, acc);
  };
  return Codegen::buildDataTiledMMAUnderlyingOperations(
      builder, loc, getSwizzle(*this, kMMAOperandLhs),
      getSwizzle(*this, kMMAOperandRhs), getSwizzle(*this, kMMAOperandAcc),
      getIntrinsicsM(), getIntrinsicsN(), getIntrinsicsK(), inputs, outputs,
      emitIntrinsic, results);
}

TileSwizzle DataTiledMMAAttr::getTileSwizzle(unsigned operandIndex) const {
  return getSwizzle(*this, operandIndex);
}

IREE::Codegen::TileMxNxKxKb DataTiledMMAAttr::getTileMNKKb() const {
  IREE::Codegen::TileMxNxKxKb innerTile;
  std::tie(innerTile.M, innerTile.N, innerTile.K) =
      getMNKShapeFromIntrinsic(getIntrinsic());
  innerTile.M *= getIntrinsicsM() * getSubgroupsM();
  innerTile.N *= getIntrinsicsN() * getSubgroupsN();
  innerTile.K *= getIntrinsicsK() * getSubgroupsK();
  return innerTile;
}

void DataTiledMMAAttr::getElementTypes(SmallVectorImpl<Type> &result) const {
  auto [a, b, c] = IREE::GPU::getABCElementTypes(getContext(), getIntrinsic());
  result.assign({a, b, c});
}

//===----------------------------------------------------------------------===//
// VirtualMMA Attributes
//===----------------------------------------------------------------------===//

VirtualMMAAttr VirtualMMAAttr::get(MLIRContext *context,
                                   VirtualMMAIntrinsic type) {
  return Base::get(context, type, /*colMajor=*/false);
}

static std::tuple<int64_t, int64_t, int64_t>
getMNKShape(VirtualMMAIntrinsic type) {
  // V(Virtual)MFMA instructions which have 2 mfma instructions interleaved
  // along the k dimension.
  switch (type) {
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F16:
    return {16, 16, 32};
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F16:
    return {32, 32, 16};
  // Sparse trick VDMFMAs for skinny GEMMs: semantically 8x16xK.
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_BF16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_BF16:
    return {8, 16, 64};
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x2_I8:
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x1_I8:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2_F8E4M3FN:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN:
    return {8, 16, 128};
  }
  assert(false && "unhandled virtual mma layout type.");
  return {};
}

static std::tuple<Type, Type, Type>
getABCElementTypes(MLIRContext *context, VirtualMMAIntrinsic type) {
  Type f8E4M3FNUZ = Float8E4M3FNUZType::get(context);
  Type f8E5M2FNUZ = Float8E5M2FNUZType::get(context);
  Type f8E4M3FN = Float8E4M3FNType::get(context);
  Type f8E5M2 = Float8E5M2Type::get(context);
  Type bf16 = BFloat16Type::get(context);
  Type f16 = Float16Type::get(context);
  Type f32 = Float32Type::get(context);
  Type i8 = IntegerType::get(context, 8);
  Type i32 = IntegerType::get(context, 32);

  switch (type) {
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F8E4M3FNUZ:
    return {f8E4M3FNUZ, f8E4M3FNUZ, f32};
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F8E4M3FNUZ:
    return {f8E4M3FNUZ, f8E4M3FNUZ, f32};
  // V(Virtual)MFMA instructions which have 2 mfma instructions interleaved
  // along the k dimension.
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F16:
    return {f16, f16, f32};
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F16:
    return {f16, f16, f32};
  // Sparse trick VDMFMAs for skinny GEMMs.
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_F16:
    return {f16, f16, f32};
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_BF16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_BF16:
    return {bf16, bf16, f32};
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x2_I8:
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x1_I8:
    return {i8, i8, i32};
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ:
    return {f8E5M2FNUZ, f8E5M2FNUZ, f32};
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ_F8E4M3FNUZ:
    return {f8E5M2FNUZ, f8E4M3FNUZ, f32};
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ_F8E5M2FNUZ:
    return {f8E4M3FNUZ, f8E5M2FNUZ, f32};
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ:
    return {f8E4M3FNUZ, f8E4M3FNUZ, f32};
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2:
    return {f8E5M2, f8E5M2, f32};
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2_F8E4M3FN:
    return {f8E5M2, f8E4M3FN, f32};
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN_F8E5M2:
    return {f8E4M3FN, f8E5M2, f32};
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN:
    return {f8E4M3FN, f8E4M3FN, f32};
  }
  assert(false && "unhandled virtual mma layout type.");
  return {};
}

static OpaqueMmaLayout getOpaqueMMALayout(MLIRContext *context,
                                          VirtualMMAIntrinsic intrinsic) {
  OpaqueMmaLayout o;
  std::tie(o.aType, o.bType, o.cType) = getABCElementTypes(context, intrinsic);
  auto lhs = getSingleSubgroupLayout(intrinsic, kMMAOperandLhs);
  auto rhs = getSingleSubgroupLayout(intrinsic, kMMAOperandRhs);
  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;
}

int64_t VirtualMMAAttr::getExpectedNumInputs() const { return 2; }

int64_t VirtualMMAAttr::getExpectedNumOutputs() const { return 1; }

LogicalResult
VirtualMMAAttr::verifyIndexingMaps(ArrayRef<AffineMap> maps) const {
  return verifyMmaIndexingMaps(maps);
}

void VirtualMMAAttr::getUndistributedTileTypes(
    SmallVectorImpl<VectorType> &result) const {
  MLIRContext *ctx = getContext();
  OpaqueMmaLayout o = getOpaqueMMALayout(ctx, getIntrinsic());
  result.assign({VectorType::get({o.mSize, o.kSize}, o.aType),
                 VectorType::get({o.kSize, o.nSize}, o.bType),
                 VectorType::get({o.mSize, o.nSize}, o.cType)});
}

// Returns the number of elements held per lane for a given operand layout,
// accounting for physicalLanesPerThread when threadProduct < subgroupSize.
static int64_t getPerLaneElements(MMASingleSubgroupLayout layout,
                                  int64_t subgroupSize) {
  int64_t threadProduct = llvm::product_of(layout.thread);
  assert(subgroupSize % threadProduct == 0 &&
         "subgroup size must be a multiple of thread product");
  int64_t physicalLanesPerThread = subgroupSize / threadProduct;
  int64_t totalElements =
      llvm::product_of(layout.element) * llvm::product_of(layout.outer);
  assert(totalElements % physicalLanesPerThread == 0 &&
         "total elements must be divisible by physicalLanesPerThread");
  return totalElements / physicalLanesPerThread;
}

void VirtualMMAAttr::getDistributedTileTypes(
    SmallVectorImpl<VectorType> &result) const {
  MLIRContext *context = getContext();
  VirtualMMAIntrinsic intrinsic = getIntrinsic();
  auto lhsLayout = getSingleSubgroupLayout(intrinsic, kMMAOperandLhs);
  int64_t subgroupSize = getSubgroupSize();
  int64_t physicalLanesPerThread =
      subgroupSize / llvm::product_of(lhsLayout.thread);
  if (isVDMFMAIntrinsic(intrinsic) && physicalLanesPerThread > 1) {
    OpaqueMmaLayout o = getOpaqueMMALayout(context, intrinsic);
    auto rhsLayout = getSingleSubgroupLayout(intrinsic, kMMAOperandRhs);
    auto accLayout = getSingleSubgroupLayout(intrinsic, kMMAOperandAcc);
    result.assign(
        {VectorType::get({getPerLaneElements(lhsLayout, subgroupSize)},
                         o.aType),
         VectorType::get({getPerLaneElements(rhsLayout, subgroupSize)},
                         o.bType),
         VectorType::get({getPerLaneElements(accLayout, subgroupSize)},
                         o.cType)});
  } else {
    result.assign({getThreadVectorType(context, intrinsic, kMMAOperandLhs),
                   getThreadVectorType(context, intrinsic, kMMAOperandRhs),
                   getThreadVectorType(context, intrinsic, kMMAOperandAcc)});
  }
}

int64_t VirtualMMAAttr::getSubgroupSize() const {
  switch (getIntrinsic()) {
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F16:
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_BF16:
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x2_I8:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2_F8E4M3FN:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_BF16:
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x1_I8: {
    return 64;
  }
  }
  assert(false && "unhandled virtual mma layout type.");
  return 0;
}

Attribute VirtualMMAAttr::getDistributionMappingKind() const {
  return IREE::GPU::LaneIdAttr::get(getContext(), 0);
}

OpFoldResult VirtualMMAAttr::getDistributionWorkerCount(OpBuilder &, Location,
                                                        Operation *) const {
  return getAsIndexOpFoldResult(getContext(), getSubgroupSize());
}

LogicalResult VirtualMMAAttr::populateOperandOffsetsSizesStrides(
    OpBuilder &builder, Location loc, uint32_t operandIndex, Value laneId,
    ArrayRef<int64_t> permutation, SmallVectorImpl<OpFoldResult> &offsets,
    SmallVectorImpl<OpFoldResult> &sizes,
    SmallVectorImpl<OpFoldResult> &strides) const {
  assert(operandIndex <= 2 && "Must index valid MMA operand");
  MMASingleSubgroupLayout subgroupLayout =
      getSingleSubgroupLayout(getIntrinsic(), operandIndex,
                              operandIndex == kMMAOperandAcc && getColMajor());

  // When thread product < subgroup size, multiple physical lanes share the
  // same position in the thread[i] decomposition. physicalLanesPerThread
  // tells populateCanonicalOffsetsSizesAndStrides to split the element
  // dimension so each physical lane gets a unique slice.
  int64_t threadProduct = llvm::product_of(subgroupLayout.thread);
  assert(getSubgroupSize() % threadProduct == 0 &&
         "subgroup size must be a multiple of thread product");
  int64_t physicalLanesPerThread = getSubgroupSize() / threadProduct;

  SmallVector<OpFoldResult> canonicalOffsets;
  SmallVector<OpFoldResult> canonicalSizes;
  if (failed(populateCanonicalOffsetsSizesAndStrides(
          builder, loc, laneId, permutation, subgroupLayout, canonicalOffsets,
          canonicalSizes, strides, physicalLanesPerThread))) {
    return failure();
  }
  offsets.append(canonicalOffsets);
  sizes.append(canonicalSizes);
  return success();
}

// Returns true on odd lanes and false on even lanes.
static Value createLaneParityPredicate(OpBuilder &builder, Location loc) {
  Value laneId = gpu::LaneIdOp::create(builder, loc, /*upper_bound=*/nullptr);
  Value one = arith::ConstantIndexOp::create(builder, loc, 1);
  Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
  Value lowBit = arith::AndIOp::create(builder, loc, laneId, one);
  return arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::ne, lowBit,
                               zero);
}

// Creates a constant sparse index for SMFMAC operations.
//
// The sparse index encodes which 2 positions out of each group of 4
// K-elements are selected for 2:4 structured sparsity. Each 4-bit
// field within selectorBits selects positions for one K-group:
//   0x4 (0100b) -> positions {0,1};  0xE (1110b) -> positions {2,3}.
//
// gfx942 16-bit (k=32): vector<4xi8>, 2 groups per i8.
// gfx950 16-bit / gfx942 8-bit (k=64): vector<2xi16>, 4 groups per i16.
// gfx950 8-bit (k=128): i32 scalar, 8 groups packed into 32 bits.
//
// For vector types, only the first element carries active selector bits;
// remaining elements are padding zeros.
static Value createConstSparseIndex(OpBuilder &builder, Location loc,
                                    Type sparseIndexType,
                                    uint32_t selectorBits) {
  if (auto vecTy = dyn_cast<VectorType>(sparseIndexType)) {
    Type elemTy = vecTy.getElementType();
    Value zero =
        arith::ConstantOp::create(builder, loc, builder.getZeroAttr(vecTy));
    Value selector = arith::ConstantOp::create(
        builder, loc, builder.getIntegerAttr(elemTy, selectorBits));
    return vector::InsertOp::create(builder, loc, selector, zero, 0);
  }
  return arith::ConstantOp::create(
      builder, loc, builder.getIntegerAttr(sparseIndexType, selectorBits));
}

// Returns the number of native intrinsics chained along K per virtual
// intrinsic.
//
// For VDMFMA this matches the trailing `x<U>` suffix in the
// intrinsic name (e.g. `VDMFMA_F32_8x16x64x2_F16` -> 2,
// `VDMFMA_F32_8x16x64x1_F16` -> 1), such that `nativeK * U` covers the
// semantic K in the `MxNxK` triple of the name.
//
// On gfx942 (CDNA3), the native smfmac K-depth is narrower, so 2 intrinsics
// are chained per virtual intrinsic to cover the semantic K (-> `x2`
// variants). On gfx950 (CDNA4), the native smfmac K-depth is doubled, so a
// single intrinsic already covers it (-> `x1` variants).
int64_t VirtualMMAAttr::getIntrinsicsK() const {
  switch (getIntrinsic()) {
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F16:
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_BF16:
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x2_I8:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ: {
    return 2;
  }
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2_F8E4M3FN:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_BF16:
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x1_I8: {
    return 1;
  }
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F8E4M3FNUZ: {
    return 1;
  }
  }
  assert(false && "unhandled virtual mma layout type.");
  return 0;
}

// Struct with consolidated info necessary for sparse trick invocation as a
// VDMFMA.
struct VDMFMAConfig {
  int64_t m, n, nativeK;
  int64_t unrollFactor;
  Type sparseIndexType;
  uint32_t evenSparseIndex;
  uint32_t oddSparseIndex;
  int64_t aSliceWidth;
};

// Virtual Dense MFMA (VDMFMA) ops represent invocations of the sparse trick
// targeting skinny GEMMs (M=8).
//
// === The sparse trick ===
//
// Sparse MFMA (V_SMFMAC) instructions perform MMA on an imbalanced pair of
// operands: a 4:2 structured-sparse matrix A and a dense matrix B. The
// instruction also takes a sparsity index that encodes which 2 of every 4
// elements along K are non-zero within the sparse matrix A. The trick exploits
// this by pairing even/odd lanes to jointly describe a full dense row.
//
// The lane-pairing layout maps each of the 8 logical M-rows to a pair of
// adjacent physical rows (row 2i and 2i+1 for logical row i). Within each pair,
// the even lane supplies positions {0,1} from each K-group of 4 and the odd
// lane supplies positions {2,3}. The hardware interprets each physical row as
// having 2:4 structured sparsity and computes a partial dot product over only
// its non-zero elements. Summing the two physical rows' results reconstructs
// the full dense dot product for the logical row. This yields a semantic M=8
// matmul from a physical 16x16 instruction.
//
// Each lane loads unique A data via physicalLanesPerThread distribution. Even
// lanes receive K[0:aSliceWidth*unrollFactor/2], odd lanes receive
// K[aSliceWidth*unrollFactor/2:aSliceWidth*unrollFactor]. A is sliced
// sequentially into per-SMFMAC chunks of aSliceWidth elements.
//
// === Accumulator expand/collapse ===
//
// Because the sparse trick maps two hardware rows to one logical row, adjacent
// register pairs in the output hold partial sums for the same dense row.
// Collapsing sums each pair (v0+v1, v2+v3) to produce the 2-element semantic
// result: one complete value per logical row.
//
// The layout and distribution infrastructure operate on the collapsed vector
// (e.g., vector<2xf32>). buildVDMFMAOps handles the translation: it expands
// a collapsed accumulator into the 4-element physical form before the smfmac
// chain, then collapses the result back afterward.

// Returns the B shuffle indices for one unrolled VDMFMA slice.
static SmallVector<int64_t, 16>
getVDMFMABInterleaveIndices(int64_t aSliceWidth, int64_t unrollFactor,
                            int64_t sliceIndex) {
  assert(aSliceWidth % 2 == 0 && "expected B slices grouped in pairs");
  assert(unrollFactor > 0 && "expected positive unroll factor");
  assert(sliceIndex >= 0 && sliceIndex < unrollFactor &&
         "slice index must be within the unrolled K slices");

  int64_t halfWidth = aSliceWidth * unrollFactor;
  int64_t sliceOffset = aSliceWidth * sliceIndex;
  SmallVector<int64_t, 16> interleaveIndices;
  for (int64_t groupOffset = 0; groupOffset < aSliceWidth; groupOffset += 2) {
    int64_t lowerBase = sliceOffset + groupOffset;
    int64_t upperBase = lowerBase + halfWidth;
    interleaveIndices.push_back(lowerBase);
    interleaveIndices.push_back(lowerBase + 1);
    interleaveIndices.push_back(upperBase);
    interleaveIndices.push_back(upperBase + 1);
  }
  return interleaveIndices;
}

bool isVDMFMAIntrinsic(VirtualMMAIntrinsic intrinsic) {
  switch (intrinsic) {
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_BF16:
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x2_I8:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2_F8E4M3FN:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_BF16:
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x1_I8:
    return true;
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F16:
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F16:
    return false;
  }
  assert(false && "unhandled virtual mma intrinsic type");
  return false;
}

static LogicalResult buildVDMFMAOps(OpBuilder &builder, Location loc,
                                    const VDMFMAConfig &config,
                                    ValueRange inputs, Value acc,
                                    SmallVectorImpl<Value> &results) {
  // Expands a collapsed 2-element ACC into the 4-element native SMFMAC form
  // by interleaving with zeros: [c0, c1] -> [c0, 0, c1, 0]. The deinterleave
  // will be hoisted out of the reduction loop.
  auto accType = cast<VectorType>(acc.getType());
  Value smfmacAcc =
      IREE::Util::HoistableConversionOp::create(
          builder, loc, /*tag=*/kVDMFMAInterleaveAcc,
          /*inverseTag=*/kVDMFMADeinterleaveAcc, acc,
          [accType](OpBuilder &builder, Location loc, ValueRange args) {
            Value zero = arith::ConstantOp::create(
                builder, loc, builder.getZeroAttr(accType));
            return SmallVector<Value>{
                vector::InterleaveOp::create(builder, loc, args[0], zero)};
          })
          .getResult(0);
  VectorType expandedAccType = cast<VectorType>(smfmacAcc.getType());

  Value isOddLane = createLaneParityPredicate(builder, loc);

  Value sparseIndex = arith::SelectOp::create(
      builder, loc, isOddLane,
      createConstSparseIndex(builder, loc, config.sparseIndexType,
                             config.oddSparseIndex),
      createConstSparseIndex(builder, loc, config.sparseIndexType,
                             config.evenSparseIndex));

  Value lhs = inputs[0];
  Value rhs = inputs[1];

  for (int64_t i = 0; i < config.unrollFactor; ++i) {
    int64_t aOffset = config.aSliceWidth * i;
    Value aSlice = vector::ExtractStridedSliceOp::create(
        builder, loc, lhs, /*offsets=*/ArrayRef<int64_t>{aOffset},
        /*sizes=*/ArrayRef<int64_t>{config.aSliceWidth},
        /*strides=*/ArrayRef<int64_t>{1});

    SmallVector<int64_t, 16> bInterleaveIndices =
        getVDMFMABInterleaveIndices(config.aSliceWidth, config.unrollFactor, i);
    Value bSlice =
        vector::ShuffleOp::create(builder, loc, rhs, rhs, bInterleaveIndices);

    smfmacAcc = amdgpu::SparseMFMAOp::create(
        builder, loc, expandedAccType,
        /*m=*/config.m, /*n=*/config.n, /*k=*/config.nativeK,
        /*sourceA=*/aSlice, /*sourceB=*/bSlice, /*destC=*/smfmacAcc,
        /*sparseIdx=*/sparseIndex, /*cbsz=*/0, /*abid=*/0);
  }

  // Collapses a 4-element native SMFMAC ACC back to the 2-element semantic
  // form. Deinterleaves into evens [d0, d2] and odds [d1, d3], then sums
  // pairwise: [d0, d1, d2, d3] -> [d0+d1, d2+d3].
  Type elementType = cast<VectorType>(acc.getType()).getElementType();
  Value result =
      IREE::Util::HoistableConversionOp::create(
          builder, loc, /*tag=*/kVDMFMADeinterleaveAcc,
          /*inverseTag=*/kVDMFMAInterleaveAcc, smfmacAcc,
          [elementType](OpBuilder &builder, Location loc, ValueRange args) {
            auto deinterleave =
                vector::DeinterleaveOp::create(builder, loc, args[0]);
            Value evens = deinterleave.getRes1();
            Value odds = deinterleave.getRes2();
            if (isa<FloatType>(elementType)) {
              return SmallVector<Value>{
                  arith::AddFOp::create(builder, loc, evens, odds)};
            }
            return SmallVector<Value>{
                arith::AddIOp::create(builder, loc, evens, odds)};
          })
          .getResult(0);
  results.push_back(result);
  return success();
}

// Generates amdgpu.mfma/wmma operation on the given inputs for this attribute
// type.
LogicalResult VirtualMMAAttr::buildUnderlyingOperations(
    OpBuilder &builder, Location loc, ValueRange inputs, ValueRange outputs,
    SmallVectorImpl<Value> &results) const {
  if (inputs.size() != 2) {
    return failure();
  }
  if (outputs.size() != 1) {
    return failure();
  }
  SmallVector<VectorType> threadTypes;
  getDistributedTileTypes(threadTypes);
  if (!llvm::equal(threadTypes,
                   llvm::concat<Type>(inputs.getTypes(), outputs.getTypes()))) {
    return failure();
  }

  switch (getIntrinsic()) {
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F16:
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F16: {
    // Generate mfma's for K with unrolled kernels.
    const int64_t unrollKFactor = getIntrinsicsK();
    auto [m, n, k] = getMNKShape();
    // Compute actual/native intrinsic's K size.
    int64_t nativeKSize = k / unrollKFactor;

    auto [aType, bType, cType] = getABCVectorTypes();
    if (aType.getShape()[0] != bType.getShape()[0]) {
      // Currently only support case where lhs and rhs
      // has same vectorWidth.
      return failure();
    }
    int64_t vectorWidth = aType.getShape()[0] / unrollKFactor;
    Value acc = outputs[0];
    for (int i = 0; i < unrollKFactor; i++) {
      int64_t offset = vectorWidth * i;
      Value slicedLhs = vector::ExtractStridedSliceOp::create(
          builder, loc, inputs[0], ArrayRef<int64_t>{offset},
          ArrayRef<int64_t>{vectorWidth}, ArrayRef<int64_t>{1});
      Value slicedRhs = vector::ExtractStridedSliceOp::create(
          builder, loc, inputs[1], ArrayRef<int64_t>{offset},
          ArrayRef<int64_t>{vectorWidth}, ArrayRef<int64_t>{1});
      if (getColMajor()) {
        std::swap(slicedLhs, slicedRhs);
      }
      acc = amdgpu::MFMAOp::create(builder, loc, outputs[0].getType(), m, n,
                                   nativeKSize, getBlockSize(), slicedLhs,
                                   slicedRhs, acc)
                .getResult();
    }
    results.push_back(acc);
    return success();
  }
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_BF16: {
    if (getColMajor()) {
      return failure();
    }
    VDMFMAConfig config{/*m=*/16,
                        /*n=*/16,
                        /*nativeK=*/32,
                        /*unrollFactor=*/getIntrinsicsK(),
                        /*sparseIndexType=*/
                        VectorType::get({4}, builder.getIntegerType(8)),
                        /*evenSparseIndex=*/0x44,
                        /*oddSparseIndex=*/0xEE,
                        /*aSliceWidth=*/4};
    return buildVDMFMAOps(builder, loc, config, inputs, outputs[0], results);
  }
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x2_I8:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ: {
    if (getColMajor()) {
      return failure();
    }
    VDMFMAConfig config{/*m=*/16,
                        /*n=*/16,
                        /*nativeK=*/64,
                        /*unrollFactor=*/getIntrinsicsK(),
                        /*sparseIndexType=*/
                        VectorType::get({2}, builder.getIntegerType(16)),
                        /*evenSparseIndex=*/0x4444,
                        /*oddSparseIndex=*/0xEEEE,
                        /*aSliceWidth=*/8};
    return buildVDMFMAOps(builder, loc, config, inputs, outputs[0], results);
  }
  // CDNA4/gfx950 F16/BF16: wider native smfmac K=64, unroll=1.
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_BF16: {
    if (getColMajor()) {
      return failure();
    }
    VDMFMAConfig config{/*m=*/16,
                        /*n=*/16,
                        /*nativeK=*/64,
                        /*unrollFactor=*/getIntrinsicsK(),
                        /*sparseIndexType=*/
                        VectorType::get({2}, builder.getIntegerType(16)),
                        /*evenSparseIndex=*/0x4444,
                        /*oddSparseIndex=*/0xEEEE,
                        /*aSliceWidth=*/8};
    return buildVDMFMAOps(builder, loc, config, inputs, outputs[0], results);
  }
  // CDNA4/gfx950 I8 + IEEE fp8: wider native smfmac K=128, unroll=1.
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x1_I8:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2_F8E4M3FN:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN: {
    if (getColMajor()) {
      return failure();
    }
    VDMFMAConfig config{/*m=*/16,
                        /*n=*/16,
                        /*nativeK=*/128,
                        /*unrollFactor=*/getIntrinsicsK(),
                        /*sparseIndexType=*/builder.getI32Type(),
                        /*evenSparseIndex=*/0x44444444,
                        /*oddSparseIndex=*/0xEEEEEEEE,
                        /*aSliceWidth=*/16};
    return buildVDMFMAOps(builder, loc, config, inputs, outputs[0], results);
  }
  }
  return failure();
}

int64_t VirtualMMAAttr::getBlockSize() const {
  switch (getIntrinsic()) {
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F16:
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_BF16:
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x2_I8:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2_F8E4M3FN:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_BF16:
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x1_I8: {
    return 1;
  }
  }
  assert(false && "unhandled virtual mma layout type.");
  return 0;
}

SmallVector<SmallVector<utils::IteratorType>>
VirtualMMAAttr::getOperandIteratorTypes() const {
  return {{utils::IteratorType::parallel, utils::IteratorType::reduction},
          {utils::IteratorType::reduction, utils::IteratorType::parallel},
          {utils::IteratorType::parallel, utils::IteratorType::parallel}};
}

MMASingleSubgroupLayout getSingleSubgroupLayout(VirtualMMAIntrinsic intrinsic,
                                                int operandIndex) {
  switch (intrinsic) {
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F16:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return {/*outer=*/{1, 1}, /*thread=*/{16, 4}, /*tstrides=*/{1, 16},
              /*element=*/{1, 8}};
    case kMMAOperandRhs:
      return {/*outer=*/{1, 1}, /*thread=*/{4, 16}, /*tstrides=*/{16, 1},
              /*element=*/{8, 1}};
    case kMMAOperandAcc:
      return {/*outer=*/{1, 1}, /*thread=*/{4, 16}, /*tstrides=*/{16, 1},
              /*element=*/{4, 1}};
    }
  case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F8E4M3FNUZ:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return {/*outer=*/{1, 2}, /*thread=*/{16, 4}, /*tstrides=*/{1, 16},
              /*element=*/{1, 4}};
    case kMMAOperandRhs:
      return {/*outer=*/{2, 1}, /*thread=*/{4, 16}, /*tstrides=*/{16, 1},
              /*element=*/{4, 1}};
    case kMMAOperandAcc:
      return {/*outer=*/{1, 1}, /*thread=*/{4, 16}, /*tstrides=*/{16, 1},
              /*element=*/{4, 1}};
    }
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F16:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return {/*outer=*/{1, 1}, /*thread=*/{32, 2}, /*tstrides=*/{1, 32},
              /*element=*/{1, 8}};
    case kMMAOperandRhs:
      return {/*outer=*/{1, 1}, /*thread=*/{2, 32}, /*tstrides=*/{32, 1},
              /*element=*/{8, 1}};
    case kMMAOperandAcc:
      return {/*outer=*/{4, 1}, /*thread=*/{2, 32}, /*tstrides=*/{32, 1},
              /*element=*/{4, 1}};
    }
  case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F8E4M3FNUZ:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return {/*outer=*/{1, 2}, /*thread=*/{32, 2}, /*tstrides=*/{1, 32},
              /*element=*/{1, 4}};
    case kMMAOperandRhs:
      return {/*outer=*/{2, 1}, /*thread=*/{2, 32}, /*tstrides=*/{32, 1},
              /*element=*/{4, 1}};
    case kMMAOperandAcc:
      return {/*outer=*/{4, 1}, /*thread=*/{2, 32}, /*tstrides=*/{32, 1},
              /*element=*/{4, 1}};
    }
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x2_BF16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_F16:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x64x1_BF16:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return {/*outer=*/{1, 1}, /*thread=*/{8, 4}, /*tstrides=*/{2, 16},
              /*element=*/{1, 16}};
    case kMMAOperandRhs:
      return {/*outer=*/{1, 1}, /*thread=*/{4, 16}, /*tstrides=*/{16, 1},
              /*element=*/{16, 1}};
    case kMMAOperandAcc:
      return {/*outer=*/{1, 1}, /*thread=*/{4, 16}, /*tstrides=*/{16, 1},
              /*element=*/{2, 1}};
    }
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x2_I8:
  case VirtualMMAIntrinsic::VDMFMA_I32_8x16x128x1_I8:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E5M2FNUZ_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ_F8E5M2FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x2_F8E4M3FNUZ:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E5M2_F8E4M3FN:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN_F8E5M2:
  case VirtualMMAIntrinsic::VDMFMA_F32_8x16x128x1_F8E4M3FN:
    switch (operandIndex) {
    case kMMAOperandLhs:
      return {/*outer=*/{1, 1}, /*thread=*/{8, 4}, /*tstrides=*/{2, 16},
              /*element=*/{1, 32}};
    case kMMAOperandRhs:
      return {/*outer=*/{1, 1}, /*thread=*/{4, 16}, /*tstrides=*/{16, 1},
              /*element=*/{32, 1}};
    case kMMAOperandAcc:
      return {/*outer=*/{1, 1}, /*thread=*/{4, 16}, /*tstrides=*/{16, 1},
              /*element=*/{2, 1}};
    }
  }
  assert(false && "unhandled virtual mma layout type.");
  return {};
}

//===----------------------------------------------------------------------===//
// ScaledMMA Attributes
//===----------------------------------------------------------------------===//

int64_t ScaledMMAAttr::getExpectedNumInputs() const { return 4; }

int64_t ScaledMMAAttr::getExpectedNumOutputs() const { return 1; }

MMASingleSubgroupLayout getSingleSubgroupLayout(ScaledMMAIntrinsic intrinsic,
                                                int64_t operandIndex) {
  const MMASingleSubgroupLayout mfmaAcc16x16 = {
      /*outer=*/{1, 1}, /*thread=*/{4, 16}, /*tstrides=*/{16, 1},
      /*element=*/{4, 1}};
  const MMASingleSubgroupLayout mfmaAcc32x32 = {
      /*outer=*/{4, 1}, /*thread=*/{2, 32}, /*tstrides=*/{32, 1},
      /*element=*/{4, 1}};

  switch (intrinsic) {
  case ScaledMMAIntrinsic::MFMA_SCALE_F32_16x16x128_B32:
    switch (operandIndex) {
    case kScaledMMAOperandLhs:
      return {/*outer=*/{1, 1, 1}, /*thread=*/{16, 4, 1},
              /*tstrides=*/{1, 16, 1},
              /*element=*/{1, 1, 32}};
    case kScaledMMAOperandRhs:
      return {/*outer=*/{1, 1, 1}, /*thread=*/{4, 1, 16},
              /*tstrides=*/{16, 1, 1},
              /*element=*/{1, 32, 1}};
    case kScaledMMAOperandLhsScale:
      return {/*outer=*/{1, 1}, /*thread=*/{16, 4}, /*tstrides=*/{1, 16},
              /*element=*/{1, 1}};
    case kScaledMMAOperandRhsScale:
      return {/*outer=*/{1, 1}, /*thread=*/{4, 16}, /*tstrides=*/{16, 1},
              /*element=*/{1, 1}};
    case kScaledMMAOperandAcc:
      return mfmaAcc16x16;
    }
  case ScaledMMAIntrinsic::MFMA_SCALE_F32_32x32x64_B32:
    switch (operandIndex) {
    case kScaledMMAOperandLhs:
      return {/*outer=*/{1, 1, 1}, /*thread=*/{32, 2, 1},
              /*tstrides=*/{1, 32, 1},
              /*element=*/{1, 1, 32}};
    case kScaledMMAOperandRhs:
      return {/*outer=*/{1, 1, 1}, /*thread=*/{2, 1, 32},
              /*tstrides=*/{32, 1, 1},
              /*element=*/{1, 32, 1}};
    case kScaledMMAOperandLhsScale:
      return {/*outer=*/{1, 1}, /*thread=*/{32, 2}, /*tstrides=*/{1, 32},
              /*element=*/{1, 1}};
    case kScaledMMAOperandRhsScale:
      return {/*outer=*/{1, 1}, /*thread=*/{2, 32}, /*tstrides=*/{32, 1},
              /*element=*/{1, 1}};
    case kScaledMMAOperandAcc:
      return mfmaAcc32x32;
    }
  }
  assert(false && "Unhandled scaled MMA intrinsic");
  return {};
}

int64_t ScaledMMAAttr::getBlockSize() const {
  switch (getIntrinsic()) {
  case ScaledMMAIntrinsic::MFMA_SCALE_F32_16x16x128_B32:
  case ScaledMMAIntrinsic::MFMA_SCALE_F32_32x32x64_B32:
    return 32;
  }
  assert(false &&
         "all cases should've been handled in ScaledMMA::getBlockSize()");
  return 0;
}

MMASingleSubgroupLayout getSingleSubgroupLayout(ScaledMMAIntrinsic intrinsic,
                                                int64_t operandIndex,
                                                bool isAccColMajor) {
  MMASingleSubgroupLayout baseLayout =
      getSingleSubgroupLayout(intrinsic, operandIndex);
  if (operandIndex == kScaledMMAOperandAcc && isAccColMajor) {
    std::swap(baseLayout.outer[0], baseLayout.outer[1]);
    std::swap(baseLayout.thread[0], baseLayout.thread[1]);
    std::swap(baseLayout.tstrides[0], baseLayout.tstrides[1]);
    std::swap(baseLayout.element[0], baseLayout.element[1]);
  }
  return baseLayout;
}

int64_t ScaledMMAAttr::getSubgroupSize() const {
  return getIntrinsicSubgroupSize(getIntrinsic());
}

SmallVector<Type> ScaledMMAAttr::getSupportedInputTypes(MLIRContext *ctx) {
  return {Float8E8M0FNUType::get(ctx),  Float8E5M2Type::get(ctx),
          Float8E5M2FNUZType::get(ctx), Float8E4M3FNType::get(ctx),
          Float8E4M3FNUZType::get(ctx), Float4E2M1FNType::get(ctx)};
}

SmallVector<Type> ScaledMMAAttr::getSupportedOutputTypes(MLIRContext *ctx) {
  return {Float32Type::get(ctx)};
}

LogicalResult
ScaledMMAAttr::verifyIndexingMaps(ArrayRef<AffineMap> maps) const {
  if (failed(verifyMmaIndexingMaps({maps[kScaledMMAOperandLhs],
                                    maps[kScaledMMAOperandRhs],
                                    maps[kScaledMMAOperandAcc]}))) {
    return failure();
  }

  SmallVector<llvm::SmallDenseSet<AffineExpr, 4>> resExprs(
      maps.size(), llvm::SmallDenseSet<AffineExpr, 4>());
  for (auto [set, map] : llvm::zip_equal(resExprs, maps)) {
    set.insert_range(map.getResults());
  }
  // Note: the below conditions are a best guess and may be too strict.

  // Check LHS scales aren't using indexes that LHS isn't.
  if (llvm::any_of(resExprs[kScaledMMAOperandLhsScale], [&](auto e) {
        return !resExprs[kScaledMMAOperandLhs].contains(e);
      })) {
    return failure();
  }
  // Check RHS scales aren't using indexes that RHS isn't.
  if (llvm::any_of(resExprs[kScaledMMAOperandRhsScale], [&](auto e) {
        return !resExprs[kScaledMMAOperandRhs].contains(e);
      })) {
    return failure();
  }
  return success();
}

void ScaledMMAAttr::getUndistributedTileTypes(
    SmallVectorImpl<VectorType> &results) const {
  MMASingleSubgroupLayout lhsLayout =
      getSingleSubgroupLayout(getIntrinsic(), kScaledMMAOperandLhs);
  MMASingleSubgroupLayout rhsLayout =
      getSingleSubgroupLayout(getIntrinsic(), kScaledMMAOperandRhs);

  int64_t blockSize = getBlockSize();
  int64_t m = lhsLayout.outer[0] * lhsLayout.thread[0] * lhsLayout.element[0];
  int64_t kScale =
      lhsLayout.outer[1] * lhsLayout.thread[1] * lhsLayout.element[1];
  [[maybe_unused]] int64_t layoutBlockSize =
      lhsLayout.outer[2] * lhsLayout.thread[2] * lhsLayout.element[2];
  assert(blockSize == layoutBlockSize &&
         "expected block size to be set up correctly");
  int64_t n = rhsLayout.outer[2] * rhsLayout.thread[2] * rhsLayout.element[2];

  Type lhsType = getLhsElemType();
  Type rhsType = getRhsElemType();
  Type accType = getAccElemType();
  Type scaleType = Float8E8M0FNUType::get(getContext());

  results.push_back(VectorType::get({m, kScale, blockSize}, lhsType));
  results.push_back(VectorType::get({kScale, blockSize, n}, rhsType));
  results.push_back(VectorType::get({m, kScale}, scaleType));
  results.push_back(VectorType::get({kScale, n}, scaleType));
  results.push_back(VectorType::get({m, n}, accType));
}

void ScaledMMAAttr::getDistributedTileTypes(
    SmallVectorImpl<VectorType> &results) const {
  Type lhsType = getLhsElemType();
  Type rhsType = getRhsElemType();
  Type accType = getAccElemType();
  Type scaleType = Float8E8M0FNUType::get(getContext());

  std::array<Type, 5> argTypes = {lhsType, rhsType, scaleType, scaleType,
                                  accType};
  for (auto [opIndex, type] : llvm::enumerate(argTypes)) {
    MMASingleSubgroupLayout layout =
        getSingleSubgroupLayout(getIntrinsic(), opIndex);
    int64_t outer = ShapedType::getNumElements(layout.outer);
    int64_t element = ShapedType::getNumElements(layout.element);
    results.push_back(VectorType::get({outer * element}, type));
  }
}

std::optional<SmallVector<int64_t, 2>>
ScaledMMAAttr::getUndistributedTileDimExpansion(int64_t operandIndex,
                                                int64_t dim) const {
  assert(operandIndex <= kScaledMMAOperandAcc && "invalid operand index");
  MMASingleSubgroupLayout layout =
      getSingleSubgroupLayout(getIntrinsic(), operandIndex, getColMajor());
  if (layout.outer[dim] > 1) {
    return SmallVector<int64_t, 2>{layout.outer[dim],
                                   layout.element[dim] * layout.thread[dim]};
  }
  return std::nullopt;
}

Attribute ScaledMMAAttr::getDistributionMappingKind() const {
  return IREE::GPU::LaneIdAttr::get(getContext(), 0);
}

OpFoldResult ScaledMMAAttr::getDistributionWorkerCount(OpBuilder &, Location,
                                                       Operation *) const {
  return getAsIndexOpFoldResult(getContext(), getSubgroupSize());
}

LogicalResult ScaledMMAAttr::populateOperandOffsetsSizesStrides(
    OpBuilder &builder, Location loc, uint32_t operandIndex, Value laneId,
    ArrayRef<int64_t> permutation, SmallVectorImpl<OpFoldResult> &offsets,
    SmallVectorImpl<OpFoldResult> &sizes,
    SmallVectorImpl<OpFoldResult> &strides) const {
  assert(operandIndex <= kScaledMMAOperandAcc && "Scaled MFMA has 5 operands");

  MMASingleSubgroupLayout subgroupLayout =
      getSingleSubgroupLayout(getIntrinsic(), operandIndex, getColMajor());

  SmallVector<OpFoldResult> canonicalOffsets;
  SmallVector<OpFoldResult> canonicalSizes;
  if (failed(populateCanonicalOffsetsSizesAndStrides(
          builder, loc, laneId, permutation, subgroupLayout, canonicalOffsets,
          canonicalSizes, strides))) {
    return failure();
  }
  offsets.append(canonicalOffsets);
  sizes.append(canonicalSizes);

  return success();
}

LogicalResult ScaledMMAAttr::buildUnderlyingOperations(
    OpBuilder &builder, Location loc, ValueRange inputs, ValueRange outputs,
    SmallVectorImpl<Value> &results) const {
  if (inputs.size() != 4) {
    return failure();
  }
  if (outputs.size() != 1) {
    return failure();
  }
  SmallVector<VectorType> threadTypes;
  getDistributedTileTypes(threadTypes);
  if (!llvm::equal(threadTypes,
                   llvm::concat<Type>(inputs.getTypes(), outputs.getTypes()))) {
    return failure();
  }

  SmallVector<VectorType> subgroupTypes;
  getUndistributedTileTypes(subgroupTypes);

  // Note: the scales argument is given as a vector of 4
  // scales + a constant selector to say which byte in the vector is to be used.
  // This'll allow better value reuse, hence why we extend to 4-vectors
  // instead of clamping to scalars here.

  FloatType f8E8M0 = builder.getF8E8M0Type();
  Value zeroScales = arith::ConstantOp::create(
      builder, loc,
      SplatElementsAttr::get(
          VectorType::get({getScalesVectorSize()}, f8E8M0),
          llvm::APFloat::getSmallest(f8E8M0.getFloatSemantics())));
  auto padScales = [&](Value scales) {
    Value scale = vector::ExtractOp::create(builder, loc, scales, 0);
    Value padded = vector::InsertOp::create(builder, loc, scale, zeroScales, 0);
    return padded;
  };

  Value lhs = inputs[kScaledMMAOperandLhs];
  Value rhs = inputs[kScaledMMAOperandRhs];
  Value lhsScales = padScales(inputs[kScaledMMAOperandLhsScale]);
  Value rhsScales = padScales(inputs[kScaledMMAOperandRhsScale]);
  Value acc = outputs[0];

  ArrayRef<int64_t> lhsShape = subgroupTypes[kScaledMMAOperandLhs].getShape();
  ArrayRef<int64_t> rhsShape = subgroupTypes[kScaledMMAOperandRhs].getShape();
  int64_t m = lhsShape[0];
  // We use m x [k / kPerBlock] x blockSize as the LHS pre-distribution shape
  // since this makes the higher-level tiling clearer.
  int64_t k = lhsShape[1] * lhsShape[2];
  int64_t n = rhsShape[2];

  // Since the LHS and RHS layouts are both {M,N}xK, we can get a column-major
  // result just by swapping the LHS and RHS.
  if (getColMajor()) {
    std::swap(lhs, rhs);
    std::swap(lhsScales, rhsScales);
    std::swap(n, m);
  }

  Value result =
      amdgpu::ScaledMFMAOp::create(builder, loc, m, n, k, lhs, rhs, acc,
                                   lhsScales, rhsScales, /*scalesIdxA=*/0,
                                   /*scalesIdxB=*/0);
  results.push_back(result);
  return success();
}

SmallVector<SmallVector<utils::IteratorType>>
ScaledMMAAttr::getOperandIteratorTypes() const {
  return {{utils::IteratorType::parallel, utils::IteratorType::reduction,
           utils::IteratorType::reduction},
          {utils::IteratorType::reduction, utils::IteratorType::reduction,
           utils::IteratorType::parallel},
          {utils::IteratorType::parallel, utils::IteratorType::reduction},
          {utils::IteratorType::reduction, utils::IteratorType::parallel},
          {utils::IteratorType::parallel, utils::IteratorType::parallel}};
}

//===----------------------------------------------------------------------===//
// DataTiledScaledMMA Attributes
//===----------------------------------------------------------------------===//

TileSwizzle
DataTiledScaledMMAAttr::getTileSwizzle(unsigned operandIndex) const {
  return getSwizzle(*this, operandIndex);
}

static std::tuple<int64_t, int64_t, int64_t, int64_t>
getMNKKbShapeFromScaledIntrinsic(ScaledMMAIntrinsic intrinsic) {
  MMASingleSubgroupLayout lhs =
      getSingleSubgroupLayout(intrinsic, kScaledMMAOperandLhs);
  MMASingleSubgroupLayout rhs =
      getSingleSubgroupLayout(intrinsic, kScaledMMAOperandRhs);
  int64_t m = lhs.outer[0] * lhs.thread[0] * lhs.element[0];
  int64_t n = rhs.outer[2] * rhs.thread[2] * rhs.element[2];
  int64_t k = lhs.outer[1] * lhs.thread[1] * lhs.element[1];
  int64_t kB = lhs.outer[2] * lhs.thread[2] * lhs.element[2];
  return {m, n, k, kB};
}

int64_t getMSize(ScaledMMAIntrinsic intrinsic) {
  return std::get<0>(getMNKKbShapeFromScaledIntrinsic(intrinsic));
}
int64_t getNSize(ScaledMMAIntrinsic intrinsic) {
  return std::get<1>(getMNKKbShapeFromScaledIntrinsic(intrinsic));
}
int64_t getKSize(ScaledMMAIntrinsic intrinsic) {
  return std::get<2>(getMNKKbShapeFromScaledIntrinsic(intrinsic));
}
int64_t getKbSize(ScaledMMAIntrinsic intrinsic) {
  return std::get<3>(getMNKKbShapeFromScaledIntrinsic(intrinsic));
}

IREE::Codegen::TileMxNxKxKb DataTiledScaledMMAAttr::getTileMNKKb() const {
  IREE::Codegen::TileMxNxKxKb innerTile;
  std::tie(innerTile.M, innerTile.N, innerTile.K, innerTile.KB) =
      getMNKKbShapeFromScaledIntrinsic(getIntrinsic());
  innerTile.M *= getIntrinsicsM() * getSubgroupsM();
  innerTile.N *= getIntrinsicsN() * getSubgroupsN();
  innerTile.K *= getIntrinsicsK() * getSubgroupsK();
  return innerTile;
}

void DataTiledScaledMMAAttr::getElementTypes(
    SmallVectorImpl<Type> &result) const {
  result.push_back(getLhsElemType());
  result.push_back(getRhsElemType());
  result.push_back(Float8E8M0FNUType::get(getContext()));
  result.push_back(Float8E8M0FNUType::get(getContext()));
  result.push_back(getAccElemType());
  return;
}

static Value createScaledMmaOp(OpBuilder &builder, Location loc,
                               ScaledMMAIntrinsic intrinsic, Type resultType,
                               Value lhs, Value rhs, Value lhsScales,
                               Value rhsScales, Value acc) {
  FloatType f8E8M0 = builder.getF8E8M0Type();
  Value zeroScales = arith::ConstantOp::create(
      builder, loc,
      SplatElementsAttr::get(
          VectorType::get({4}, f8E8M0),
          llvm::APFloat::getSmallest(f8E8M0.getFloatSemantics())));
  auto padScales = [&](Value scales) {
    Value scale = vector::ExtractOp::create(builder, loc, scales, 0);
    Value padded = vector::InsertOp::create(builder, loc, scale, zeroScales, 0);
    return padded;
  };

  lhsScales = padScales(lhsScales);
  rhsScales = padScales(rhsScales);

  int64_t m = getMSize(intrinsic);
  // We use m x [k / kPerBlock] x blockSize as the LHS pre-distribution shape
  // since this makes the higher-level tiling clearer.
  int64_t k = getKSize(intrinsic) * getKbSize(intrinsic);
  int64_t n = getNSize(intrinsic);

  return amdgpu::ScaledMFMAOp::create(builder, loc, m, n, k, lhs, rhs, acc,
                                      lhsScales, rhsScales, /*scalesIdxA=*/0,
                                      /*scalesIdxB=*/0);
}

LogicalResult DataTiledScaledMMAAttr::buildUnderlyingOperations(
    OpBuilder &builder, Location loc, ValueRange inputs, ValueRange outputs,
    SmallVectorImpl<Value> &results) const {
  // Validation. Similar to MMAAttr::buildMmaOperation.
  if (inputs.size() != 4) {
    return failure();
  }
  if (outputs.size() != 1) {
    return failure();
  }
  SmallVector<VectorType> regTypes;
  getDistributedTileTypes(regTypes);
  if (!llvm::equal(regTypes,
                   llvm::concat<Type>(inputs.getTypes(), outputs.getTypes()))) {
    return failure();
  }

  // Prepare Lhs/Rhs/Acc operand slices to feed the intrinsic.
  const unsigned lhsIdx = 0;
  const unsigned rhsIdx = 1;
  const unsigned lhsScalesIdx = 2;
  const unsigned rhsScalesIdx = 3;
  const unsigned accIdx = 4;
  TileSwizzle lhsSwizzle = getSwizzle(*this, lhsIdx);
  LDBG() << "DataTiledScaledMMAAttr::buildMmaOperation";
  LDBG() << "    lhsSwizzle: " << lhsSwizzle;
  SmallVector<Value> intrinsicsLhs =
      distributeMmaFragmentToIntrinsics(builder, loc, inputs[0], lhsSwizzle);

  TileSwizzle rhsSwizzle = getSwizzle(*this, rhsIdx);
  LDBG() << "DataTiledScaledMMAAttr::buildMmaOperation";
  LDBG() << "    rhsSwizzle: " << rhsSwizzle;
  SmallVector<Value> intrinsicsRhs =
      distributeMmaFragmentToIntrinsics(builder, loc, inputs[1], rhsSwizzle);

  TileSwizzle lhsScalesSwizzle = getSwizzle(*this, lhsScalesIdx);
  LDBG() << "DataTiledScaledMMAAttr::buildMmaOperation";
  LDBG() << "    lhsScalesSwizzle: " << lhsScalesSwizzle;
  SmallVector<Value> intrinsicsLhsScales = distributeMmaFragmentToIntrinsics(
      builder, loc, inputs[2], lhsScalesSwizzle);

  TileSwizzle rhsScalesSwizzle = getSwizzle(*this, rhsScalesIdx);
  LDBG() << "DataTiledScaledMMAAttr::buildMmaOperation";
  LDBG() << "    rhsScalesSwizzle: " << rhsScalesSwizzle;
  SmallVector<Value> intrinsicsRhsScales = distributeMmaFragmentToIntrinsics(
      builder, loc, inputs[3], rhsScalesSwizzle);

  TileSwizzle accSwizzle = getSwizzle(*this, accIdx);
  LDBG() << "DataTiledScaledMMAAttr::buildMmaOperation";
  LDBG() << "    accSwizzle: " << accSwizzle;

  // Distribute the accumulator into per-intrinsic slices; the reassembly
  // conversion will be hoisted out of the reduction loop.
  auto distributeOp = IREE::Util::HoistableConversionOp::create(
      builder, loc, /*tag=*/kDataTiledAccDistribute,
      /*inverseTag=*/kDataTiledAccReassemble, ValueRange{outputs[0]},
      [&](OpBuilder &b, Location loc, ValueRange args) -> SmallVector<Value> {
        return distributeMmaFragmentToIntrinsics(b, loc, args[0], accSwizzle);
      });
  SmallVector<Value> intrinsicsAcc(distributeOp.getResults());

  ScaledMMAIntrinsic intrinsic = getIntrinsic();
  auto intrinCType = cast<VectorType>(intrinsicsAcc.front().getType());

  // Loop over the 3 unroll_{m,n,k} dimensions to create the intrinsics.
  for (int64_t mu = 0; mu < getIntrinsicsM(); ++mu) {
    for (int64_t nu = 0; nu < getIntrinsicsN(); ++nu) {
      for (int64_t ku = 0; ku < getIntrinsicsK(); ++ku) {
        Value lhs = intrinsicsLhs[mu * getIntrinsicsK() + ku];
        Value rhs = intrinsicsRhs[nu * getIntrinsicsK() + ku];
        Value lhsScales = intrinsicsLhsScales[mu * getIntrinsicsK() + ku];
        Value rhsScales = intrinsicsRhsScales[nu * getIntrinsicsK() + ku];
        Value &acc = intrinsicsAcc[mu * getIntrinsicsN() + nu];
        acc = createScaledMmaOp(builder, loc, intrinsic, intrinCType, lhs, rhs,
                                lhsScales, rhsScales, acc);
      }
    }
  }

  // Insert the results into the destination accumulator.
  SmallVector<int64_t> accCrossIntrinsicShape =
      Codegen::sliceSwizzledShape(accSwizzle, [](TileSwizzle::Dim dim) {
        return dim.kind() == TileSwizzle::Dim::Kind::CrossIntrinsic;
      });
  SmallVector<int64_t> accInternalShape =
      Codegen::sliceSwizzledShape(accSwizzle, [](TileSwizzle::Dim dim) {
        return dim.kind() == TileSwizzle::Dim::Kind::Internal;
      });

  LDBG() << "accCrossIntrinsicShape: "
         << llvm::interleaved(accCrossIntrinsicShape);
  LDBG() << "accInternalShape: " << llvm::interleaved(accInternalShape);

  auto reassembleOp = IREE::Util::HoistableConversionOp::create(
      builder, loc, /*tag=*/kDataTiledAccReassemble,
      /*inverseTag=*/kDataTiledAccDistribute, intrinsicsAcc,
      [&](OpBuilder &b, Location loc, ValueRange args) -> SmallVector<Value> {
        int dstRank = accCrossIntrinsicShape.size();
        SmallVector<int64_t> strides(dstRank, 1);
        SmallVector<int64_t> indices(dstRank, 0);
        Value acc = arith::ConstantOp::create(
            b, loc, b.getZeroAttr(outputs[0].getType()));
        for (Value intrAcc : args) {
          auto expandedAcc = vector::ShapeCastOp::create(
              b, loc,
              VectorType::get(
                  accInternalShape,
                  cast<VectorType>(outputs[0].getType()).getElementType()),
              intrAcc);
          acc = vector::InsertStridedSliceOp::create(b, loc, expandedAcc, acc,
                                                     indices, strides);
          incrementIndices(indices, accCrossIntrinsicShape);
        }
        return {acc};
      });
  results.push_back(reassembleOp.getResult(0));
  return success();
}

int64_t DataTiledScaledMMAAttr::getExpectedNumInputs() const { return 4; }

int64_t DataTiledScaledMMAAttr::getExpectedNumOutputs() const { return 1; }

LogicalResult DataTiledScaledMMAAttr::populateOperandOffsetsSizesStrides(
    OpBuilder &builder, Location loc, uint32_t operandIndex, Value laneId,
    ArrayRef<int64_t> permutation, SmallVectorImpl<OpFoldResult> &offsets,
    SmallVectorImpl<OpFoldResult> &sizes,
    SmallVectorImpl<OpFoldResult> &strides) const {
  if (!isUnswizzledOperand(operandIndex)) {
    return cast<DataTiledMMAInterfaceAttr>(*this)
        .populateOperandOffsetsSizesStrides(builder, loc, operandIndex, laneId,
                                            permutation, offsets, sizes,
                                            strides);
  }
  // For unswizzled operands, getTileSwizzle returns an identity permutation,
  // but populateSwizzleBasedOffsetsSizesStrides uses the swizzle's permutation
  // to order the delinearization basis, it must reflect the MMA tstrides. Use
  // the non-identity-reset swizzle for both the delinearization and the output
  // reordering so the two cancel out, yielding source-order offsets that match
  // the identity layout.
  TileSwizzle dtSwizzle = getDistributionSwizzle(*this, operandIndex);
  return populateSwizzleBasedOffsetsSizesStrides(
      builder, loc, dtSwizzle, laneId, dtSwizzle.permutation(), offsets, sizes,
      strides);
}

int64_t DataTiledScaledMMAAttr::getSubgroupSize() const {
  return getIntrinsicSubgroupSize(getIntrinsic());
}

int64_t DataTiledScaledMMAAttr::getFlatWorkgroupSize() const {
  return getSubgroupSize() * getSubgroupsM() * getSubgroupsN() *
         getSubgroupsK();
}

LogicalResult
DataTiledScaledMMAAttr::verifyIndexingMaps(ArrayRef<AffineMap> maps) const {
  return IREE::LinalgExt::inferScaledContractionDims(maps);
}

SmallVector<SmallVector<utils::IteratorType>>
DataTiledScaledMMAAttr::getOperandIteratorTypes() const {
  return {{utils::IteratorType::parallel, utils::IteratorType::reduction,
           utils::IteratorType::reduction},
          {utils::IteratorType::reduction, utils::IteratorType::reduction,
           utils::IteratorType::parallel},
          {utils::IteratorType::parallel, utils::IteratorType::reduction},
          {utils::IteratorType::reduction, utils::IteratorType::parallel},
          {utils::IteratorType::parallel, utils::IteratorType::parallel}};
}

bool DataTiledScaledMMAAttr::isUnswizzledOperand(unsigned idx) const {
  assert(idx < (unsigned)(getExpectedNumInputs() + getExpectedNumOutputs()) &&
         "operand index out of range for scaled MMA");
  auto attr = getUnswizzledOperands();
  return attr && llvm::is_contained(attr.asArrayRef(), idx);
}

bool DataTiledScaledMMAAttr::hasUnswizzledOperands() const {
  auto attr = getUnswizzledOperands();
  return attr && !attr.empty();
}

//===----------------------------------------------------------------------===//
// Target Attributes
//===----------------------------------------------------------------------===//

std::optional<int> TargetAttr::getCUDAComputeCapability() const {
  StringRef arch = getArch();
  if (!arch.starts_with("sm_")) {
    return false;
  }
  APInt version;
  if (arch.substr(3).getAsInteger(10, version)) {
    return false;
  }
  return version.getZExtValue();
}

bool TargetAttr::supportsTF32InputMMAOps() const {
  // TODO: scan the list of MMA ops to decude after plumbing through support
  // for NVIDIA TensorCore MMA ops.
  if (auto cc = getCUDAComputeCapability()) {
    return cc >= 80;
  }
  return false;
}

bool TargetAttr::supportsSyncMMAOps() const {
  if (auto cc = getCUDAComputeCapability()) {
    return cc >= 80;
  }
  return false;
}

std::array<int64_t, 3> TargetAttr::getMaximumWorkgroupCount() const {
  DenseI32ArrayAttr maxWgpCount = getWgp().getMaxWorkgroupCounts();
  assert(maxWgpCount.size() <= 3 && "expected only workgroup count for x,y,z");
  std::array<int64_t, 3> maxWorkgroupCount = {
      ShapedType::kDynamic, ShapedType::kDynamic, ShapedType::kDynamic};
  for (auto [index, value] : llvm::enumerate(maxWgpCount.asArrayRef())) {
    maxWorkgroupCount[index] = value;
  }
  return maxWorkgroupCount;
}

//===----------------------------------------------------------------------===//
// Lowering Config Attributes
//===----------------------------------------------------------------------===//

constexpr StringLiteral kWorkgroupLevelName = "workgroup";
constexpr StringLiteral kPartialReductionLevelName = "partial_reduction";
constexpr StringLiteral kReductionLevelName = "reduction";
constexpr StringLiteral kSerialLevelName = "serial";
constexpr StringLiteral kThreadLevelName = "thread";
constexpr StringLiteral kSubgroupLevelName = "subgroup";
constexpr StringLiteral kLaneLevelName = "lane";

StringRef getTilingLevelName(GPU::TilingLevel level) {
  switch (level) {
  case GPU::TilingLevel::Workgroup:
    return kWorkgroupLevelName;
  case GPU::TilingLevel::PartialReduction:
    return kPartialReductionLevelName;
  case GPU::TilingLevel::Reduction:
    return kReductionLevelName;
  case GPU::TilingLevel::Serial:
    return kSerialLevelName;
  case GPU::TilingLevel::Thread:
    return kThreadLevelName;
  case GPU::TilingLevel::Subgroup:
    return kSubgroupLevelName;
  case GPU::TilingLevel::Lane:
    return kLaneLevelName;
  }
  assert(false && "Unknown tiling level");
  return StringRef();
}

static SmallVector<int64_t> getIntegerVector(ArrayAttr array) {
  if (!array || !llvm::all_of(array.getValue(), llvm::IsaPred<IntegerAttr>)) {
    return {};
  }
  return llvm::map_to_vector(array.getValue(), [](Attribute s) -> int64_t {
    return cast<IntegerAttr>(s).getInt();
  });
}

static SmallVector<int64_t> getTileSizes(DictionaryAttr config,
                                         GPU::TilingLevel level) {
  return getIntegerVector(config.getAs<ArrayAttr>(getTilingLevelName(level)));
}

SmallVector<int64_t> LoweringConfigAttr::getWorkgroupTileSizes() const {
  return getTileSizes(getAttributes(), GPU::TilingLevel::Workgroup);
}

LogicalResult
LoweringConfigAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                           DictionaryAttr attributes) {
  return verifyPromotedOperandsList(emitError, attributes);
}

SmallVector<int64_t>
LoweringConfigAttr::getStaticTilingLevelSizes(unsigned level,
                                              Operation *op) const {
  if (level > llvm::to_underlying(GPU::TilingLevel::Lane)) {
    return {};
  }
  return getTileSizes(getAttributes(), static_cast<GPU::TilingLevel>(level));
}

SmallVector<OpFoldResult>
LoweringConfigAttr::getTilingLevelSizes(OpBuilder &b, unsigned level,
                                        Operation *op) const {
  if (level > llvm::to_underlying(GPU::TilingLevel::Lane)) {
    return {};
  }
  SmallVector<int64_t> sizes =
      getTileSizes(getAttributes(), static_cast<GPU::TilingLevel>(level));
  return llvm::map_to_vector(
      sizes, [&](int64_t s) -> OpFoldResult { return b.getIndexAttr(s); });
}

bool LoweringConfigAttr::hasTilingLevel(unsigned level) const {
  if (level > llvm::to_underlying(GPU::TilingLevel::Lane)) {
    return false;
  }
  return !getTileSizes(getAttributes(), static_cast<GPU::TilingLevel>(level))
              .empty();
}

bool LoweringConfigAttr::hasWorkgroupTilingLevel() const {
  return !getWorkgroupTileSizes().empty();
}

constexpr StringLiteral kLoweringStrategyName = "lowering_strategy";

std::optional<StringRef> LoweringConfigAttr::getLoweringStrategy() const {
  if (auto name = getAttributes().getAs<StringAttr>(kLoweringStrategyName)) {
    return name.strref();
  }
  return std::nullopt;
}
constexpr StringLiteral kWorkgroupReorderingStrategyName =
    "workgroup_reordering_strategy";

::mlir::iree_compiler::IREE::Codegen::WorkgroupReorderingAttrInterface
LoweringConfigAttr::getWorkgroupReorderingStrategy() const {
  if (auto attr = getAttributes()
                      .getAs<::mlir::iree_compiler::IREE::Codegen::
                                 WorkgroupReorderingAttrInterface>(
                          kWorkgroupReorderingStrategyName)) {
    return attr;
  }
  return nullptr;
}

//===----------------------------------------------------------------------===//
// Conditional Transpose Workgroup Reordering
//===----------------------------------------------------------------------===//

static std::tuple<SmallVector<OpFoldResult>, SmallVector<OpFoldResult>,
                  SmallVector<OpFoldResult>>
getLoopBounds(ArrayRef<Range> loopRanges,
              ArrayRef<OpFoldResult> givenTileSizes) {
  SmallVector<OpFoldResult> lbs, ubs, steps;
  for (auto [loopRange, givenTileSize] :
       llvm::zip_equal(loopRanges, givenTileSizes)) {
    // No loop if the tile size is 0.
    if (isZeroInteger(givenTileSize)) {
      continue;
    }
    lbs.push_back(loopRange.offset);
    ubs.push_back(loopRange.size);
    steps.push_back(givenTileSize);
  }
  return {lbs, ubs, steps};
}

static Value computeNumTiles(OpBuilder &b, Location loc, OpFoldResult size,
                             OpFoldResult tileSize) {
  AffineExpr s0, s1;
  bindSymbols(b.getContext(), s0, s1);
  AffineExpr ceilDivExpr = s0.ceilDiv(s1);
  OpFoldResult numTiles = affine::makeComposedFoldedAffineApply(
      b, loc, ceilDivExpr, {size, tileSize});
  return getValueOrCreateConstantIntOp(b, loc, numTiles);
}

static AffineExpr getMulExpr(OpBuilder &b) {
  AffineExpr s0, s1;
  bindSymbols(b.getContext(), s0, s1);
  return s0 * s1;
}

static AffineExpr getSubExpr(OpBuilder &b) {
  AffineExpr s0, s1;
  bindSymbols(b.getContext(), s0, s1);
  return s0 - s1;
}

/// Computes the total number of different Tile loads across all XCDs per
/// iteration for the pingpong matmul kernel. For each iteration, we distribute
/// tile loads to invidiual CUs along the X axis of the RHS.
static Value computeNumTileLoads(OpBuilder &b, Location loc, OpFoldResult size,
                                 OpFoldResult tileSize, Value numXCDs,
                                 Value numCUs) {
  Value numTiles = computeNumTiles(b, loc, size, tileSize);
  AffineExpr mulExpr = getMulExpr(b);
  Value totalCUs =
      getValueOrCreateConstantIntOp(b, loc,
                                    affine::makeComposedFoldedAffineApply(
                                        b, loc, mulExpr, {numXCDs, numCUs}));
  auto numX = arith::MinUIOp::create(b, loc, numTiles, totalCUs)->getResult(0);

  AffineExpr s0, s1;
  bindSymbols(b.getContext(), s0, s1);
  AffineExpr numTilesExpr =
      s0.ceilDiv(s1) + s1; // totalCUs.ceildiv(numX) + numX
  OpFoldResult numTileLoads = affine::makeComposedFoldedAffineApply(
      b, loc, numTilesExpr, {totalCUs, numX});
  return getValueOrCreateConstantIndexOp(b, loc, numTileLoads);
}

/// Retrieves the condition for transposed ordering. Transposed order is enabled
/// if the number of tile loads is at least 20% less than in regular order.
static Value getCondition(OpBuilder &b, Location loc,
                          ArrayRef<OpFoldResult> lbs,
                          ArrayRef<OpFoldResult> ubs,
                          ArrayRef<OpFoldResult> steps, Value numXCDs,
                          Value numCUs) {
  assert(ubs.size() == 2 && steps.size() == 2 && "rank must be 2");

  AffineExpr subExpr = getSubExpr(b);
  OpFoldResult size0 =
      affine::makeComposedFoldedAffineApply(b, loc, subExpr, {ubs[0], lbs[0]});
  OpFoldResult size1 =
      affine::makeComposedFoldedAffineApply(b, loc, subExpr, {ubs[1], lbs[1]});

  Value transposedOrder =
      computeNumTileLoads(b, loc, size0, steps[0], numXCDs, numCUs);
  Value defaultOrder =
      computeNumTileLoads(b, loc, size1, steps[1], numXCDs, numCUs);

  AffineExpr s0, s1, s2;
  bindSymbols(b.getContext(), s0, s1, s2);
  AffineExpr transposedOrderBumpExpr = (s0 * s1).floorDiv(s2);
  Value c4 = arith::ConstantIntOp::create(b, loc, 4, 32);
  Value c5 = arith::ConstantIntOp::create(b, loc, 5, 32);

  Value defaultOrderTolerance = getValueOrCreateConstantIntOp(
      b, loc,
      affine::makeComposedFoldedAffineApply(b, loc, transposedOrderBumpExpr,
                                            {defaultOrder, c4, c5}));
  return mlir::arith::CmpIOp::create(b, loc, mlir::arith::CmpIPredicate::ult,
                                     transposedOrder, defaultOrderTolerance);
}
/// Swap values based on `pred`.
static void swapIf(OpBuilder &b, Location loc, OpFoldResult pred,
                   SmallVector<OpFoldResult> &values,
                   ArrayRef<size_t> ids = {0, 1}) {

  assert(ids.size() == 2 && "Can only swap between 2 indices");
  Value v0 = getValueOrCreateConstantIndexOp(b, loc, values[ids[0]]);
  Value v1 = getValueOrCreateConstantIndexOp(b, loc, values[ids[1]]);
  Value predVal = getValueOrCreateConstantIndexOp(b, loc, pred);
  values[ids[0]] =
      mlir::arith::SelectOp::create(b, loc, predVal, v1, v0).getResult();
  values[ids[1]] =
      mlir::arith::SelectOp::create(b, loc, predVal, v0, v1).getResult();
}

static std::tuple<SmallVector<OpFoldResult>, SmallVector<OpFoldResult>,
                  SmallVector<size_t>>
computeOffsetAndSize(ArrayRef<Range> loopRanges,
                     ArrayRef<OpFoldResult> givenTileSizes,
                     ArrayRef<Value> ivs) {
  SmallVector<size_t> ids;
  SmallVector<OpFoldResult> offsets, sizes;

  offsets.reserve(loopRanges.size());
  sizes.reserve(loopRanges.size());
  const Value *ivIt = ivs.begin();

  for (auto [loopRange, givenTileSize] :
       llvm::zip_equal(loopRanges, givenTileSizes)) {
    if (isZeroInteger(givenTileSize)) {
      offsets.push_back(loopRange.offset);
      sizes.push_back(loopRange.size);
    } else {
      offsets.push_back(*ivIt++);
      sizes.push_back(givenTileSize);
      ids.push_back(sizes.size() - 1);
    }
  }
  return {offsets, sizes, ids};
}

FailureOr<mlir::scf::SCFTilingOptions::CustomLoopHeaderInfo>
ConditionalTransposeAttr::generateLoopHeaderFn(
    OpBuilder &builder, Location loc, ArrayRef<Range> loopRanges,
    ArrayRef<OpFoldResult> givenTileSizes,
    ValueRange destinationTensors) const {
  scf::ForallOp forallOp;
  // Get loop bounds.
  SmallVector<OpFoldResult> lbs, ubs, steps;
  std::tie(lbs, ubs, steps) = getLoopBounds(loopRanges, givenTileSizes);

  if (lbs.size() != 2) {
    return emitError(loc) << "conditional_transpose only supports rank 2";
  }

  SmallVector<Attribute> deviceMappingAttribute;
  deviceMappingAttribute.push_back(IREE::Codegen::WorkgroupMappingAttr::get(
      builder.getContext(), IREE::Codegen::symbolizeWorkgroupId(1).value()));
  deviceMappingAttribute.push_back(IREE::Codegen::WorkgroupMappingAttr::get(
      builder.getContext(), IREE::Codegen::symbolizeWorkgroupId(0).value()));

  std::optional<ArrayAttr> mappingAttr =
      builder.getArrayAttr(deviceMappingAttribute);

  // Compute condition for transpose.
  Value numCUsVal = arith::ConstantIndexOp::create(builder, loc, getNumCUs());
  Value numXCDs = arith::ConstantIndexOp::create(builder, loc, getNumXCDs());
  Value cond = getCondition(builder, loc, lbs, ubs, steps, numXCDs, numCUsVal);

  // Apply condition on loop bounds.
  swapIf(builder, loc, cond, lbs);
  swapIf(builder, loc, cond, ubs);
  swapIf(builder, loc, cond, steps);

  forallOp = scf::ForallOp::create(builder, loc, lbs, ubs, steps,
                                   destinationTensors, mappingAttr);
  builder.setInsertionPoint(forallOp.getTerminator());

  SmallVector<Value> ivs = forallOp.getInductionVars();
  SmallVector<OpFoldResult> offsets, sizes;
  SmallVector<size_t> ids;
  std::tie(offsets, sizes, ids) =
      computeOffsetAndSize(loopRanges, givenTileSizes, ivs);

  // Apply condition on offsets.
  swapIf(builder, loc, cond, offsets, ids);

  ValueRange innerDestinationTensors = forallOp.getRegionOutArgs();
  return mlir::scf::SCFTilingOptions::CustomLoopHeaderInfo{
      {cast<LoopLikeOpInterface>(forallOp.getOperation())},
      offsets,
      sizes,
      innerDestinationTensors};
}

LogicalResult ConditionalTransposeAttr::generateLoopTerminatorFn(
    OpBuilder &builder, Location loc, ArrayRef<LoopLikeOpInterface> loops,
    ValueRange tiledResults, ArrayRef<SmallVector<OpFoldResult>> resultOffsets,
    ArrayRef<SmallVector<OpFoldResult>> resultSizes,
    ValueRange destinationTensors) const {
  assert(loops.size() == 1 && "Expected loop count should be 1.");
  LoopLikeOpInterface loop = loops.front();
  scf::ForallOp *forallOp = dyn_cast<scf::ForallOp>(&loop);
  if (!forallOp) {
    return emitError(loc) << "Only scf.forall op are supported";
  }

  builder.setInsertionPointToEnd(forallOp->getTerminator().getBody());

  for (auto [tiledValue, destinationTensor, resultOffset, resultSize] :
       llvm::zip_equal(tiledResults, destinationTensors, resultOffsets,
                       resultSizes)) {
    SmallVector<OpFoldResult> resultStride(resultOffset.size(),
                                           builder.getIndexAttr(1));

    tensor::ParallelInsertSliceOp::create(builder, loc, tiledValue,
                                          destinationTensor, resultOffset,
                                          resultSize, resultStride);
  }
  return success();
}

LogicalResult
ConditionalTransposeAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                                 int64_t nbXcds, int64_t nbCus) {
  if (nbXcds == 0 || nbCus == 0) {
    return emitError()
           << "conditional_transpose must have non-zeros parameters";
  }
  return success();
}

//===----------------------------------------------------------------------===//
// DerivedThreadConfigAttr
//===----------------------------------------------------------------------===//

SmallVector<int64_t>
DerivedThreadConfigAttr::getStaticTilingLevelSizes(unsigned level,
                                                   Operation *op) const {
  if (level != llvm::to_underlying(GPU::TilingLevel::Thread)) {
    return {};
  }
  return deriveThreadTileSizes(op);
}

SmallVector<OpFoldResult>
DerivedThreadConfigAttr::getTilingLevelSizes(OpBuilder &b, unsigned level,
                                             Operation *op) const {
  if (level > llvm::to_underlying(GPU::TilingLevel::Thread)) {
    return {};
  }
  SmallVector<int64_t> sizes = deriveThreadTileSizes(op);
  return llvm::map_to_vector(
      sizes, [&](int64_t s) -> OpFoldResult { return b.getIndexAttr(s); });
}

bool DerivedThreadConfigAttr::hasTilingLevel(unsigned level) const {
  return level == llvm::to_underlying(GPU::TilingLevel::Thread);
}

//===----------------------------------------------------------------------===//
// UseGlobalLoadDMAAttr
//===----------------------------------------------------------------------===//

SmallVector<int64_t>
UseGlobalLoadDMAAttr::getStaticTilingLevelSizes(unsigned level,
                                                Operation *op) const {
  if (level == llvm::to_underlying(GPU::TilingLevel::Subgroup)) {
    // Subgroup tile sizes are derived from translation_info, not stored here.
    return {};
  }
  if (level == llvm::to_underlying(GPU::TilingLevel::Thread)) {
    return globalLoadDMATileSizes(op);
  }
  return {};
}

SmallVector<OpFoldResult>
UseGlobalLoadDMAAttr::getTilingLevelSizes(OpBuilder &b, unsigned level,
                                          Operation *op) const {
  if (level == llvm::to_underlying(GPU::TilingLevel::Subgroup)) {
    // Subgroup tile sizes are derived from translation_info, not stored here.
    return {};
  }
  if (level == llvm::to_underlying(GPU::TilingLevel::Thread)) {
    SmallVector<int64_t> sizes = globalLoadDMATileSizes(op);
    return getAsIndexOpFoldResult(b.getContext(), sizes);
  }
  return {};
}

bool UseGlobalLoadDMAAttr::hasTilingLevel(unsigned level) const {
  // Subgroup level is not stored in this attribute anymore.
  return level == llvm::to_underlying(GPU::TilingLevel::Thread);
}

//===----------------------------------------------------------------------===//
// PromoteWithCacheSwizzleAttr
//===----------------------------------------------------------------------===//

Value PromoteWithCacheSwizzleAttr::promoteOperand(
    mlir::OpBuilder &builder, mlir::OpOperand &operand) const {
  return cacheSwizzlePromotionImpl(builder, operand, getCopyConfig());
}

//===----------------------------------------------------------------------===//
// SwizzleOperandAttr
//===----------------------------------------------------------------------===//

Value SwizzleOperandAttr::promoteOperand(mlir::OpBuilder &builder,
                                         mlir::OpOperand &operand) const {
  return swizzlePromotionImpl(builder, operand, getCopyConfig(), getSwizzle());
}

//===----------------------------------------------------------------------===//
// LaneIdAttr
//===----------------------------------------------------------------------===//

int64_t LaneIdAttr::getMappingId() const { return getDim(); }

bool LaneIdAttr::isLinearMapping() const { return true; }

int64_t LaneIdAttr::getRelativeIndex() const { return getDim(); }

//===----------------------------------------------------------------------===//
// GPU Pipeline Attribute
//===----------------------------------------------------------------------===//

// Returns pipelines available for this target. SPIRV targets (features
// starting with "spirv:") get SPIRVPipelineAttr pipelines; all others
// get LLVMGPU PipelineAttr pipelines.
SmallVector<Attribute> TargetAttr::getAvailablePipelines() const {
  MLIRContext *ctx = getContext();
  SmallVector<Attribute> pipelines;
  if (getFeatures().starts_with("spirv:")) {
    for (unsigned i = 0, e = getMaxEnumValForSPIRVLoweringPipeline() + 1; i < e;
         ++i) {
      if (std::optional<SPIRVLoweringPipeline> pipeline =
              symbolizeSPIRVLoweringPipeline(i)) {
        pipelines.push_back(SPIRVPipelineAttr::get(ctx, *pipeline));
      }
    }
    return pipelines;
  }
  for (unsigned i = 0, e = getMaxEnumValForLoweringPipeline() + 1; i < e; ++i) {
    if (std::optional<LoweringPipeline> pipeline =
            symbolizeLoweringPipeline(i)) {
      pipelines.push_back(PipelineAttr::get(ctx, *pipeline));
    }
  }
  return pipelines;
}

//===----------------------------------------------------------------------===//
// GPU Pipeline Options
//===----------------------------------------------------------------------===//

GPUPipelineOptionsAttr GPUPipelineOptionsAttr::get(
    MLIRContext *context, unsigned prefetchNumStages,
    bool noReduceSharedMemoryBankConflicts, bool useIgemmConvolution,
    std::optional<ReorderWorkgroupsStrategy> reorderWorkgroupsStrategy) {
  auto strategyAttr = ReorderWorkgroupsStrategyAttr();
  if (reorderWorkgroupsStrategy) {
    strategyAttr =
        ReorderWorkgroupsStrategyAttr::get(context, *reorderWorkgroupsStrategy);
  }
  Builder b(context);
  std::optional<int64_t> prefetchOpt;
  if (prefetchNumStages > 0) {
    prefetchOpt = prefetchNumStages;
  }
  return Base::get(context, prefetchOpt,
                   b.getBoolAttr(noReduceSharedMemoryBankConflicts),
                   b.getBoolAttr(useIgemmConvolution), strategyAttr);
}

// Index Hint Attributes
//===----------------------------------------------------------------------===//

// Custom parser/printer to make the step and aligned fields optional.
// When step is 1 and aligned is false, these fields will be omitted.
// Format: <group_size[, step = N][, aligned]>
Attribute IREE::GPU::LaneIncrementAttr::parse(AsmParser &parser, Type) {
  int64_t groupSize = 0;
  if (failed(parser.parseLess()) || failed(parser.parseInteger(groupSize))) {
    return {};
  }
  int64_t step = 1;
  bool aligned = false;
  bool parsedStep = false;
  bool parsedAligned = false;
  // Parse optional ", step = N" and/or ", aligned"
  while (succeeded(parser.parseOptionalComma())) {
    if (succeeded(parser.parseOptionalKeyword("step"))) {
      if (parsedStep) {
        parser.emitError(parser.getCurrentLocation(),
                         "'step' specified more than once");
        return {};
      }
      if (failed(parser.parseEqual()) || failed(parser.parseInteger(step))) {
        return {};
      }
      parsedStep = true;
    } else if (succeeded(parser.parseOptionalKeyword("aligned"))) {
      if (parsedAligned) {
        parser.emitError(parser.getCurrentLocation(),
                         "'aligned' specified more than once");
        return {};
      }
      aligned = true;
      parsedAligned = true;
    } else {
      parser.emitError(parser.getCurrentLocation(),
                       "expected 'step' or 'aligned'");
      return {};
    }
  }
  if (failed(parser.parseGreater())) {
    return {};
  }
  return LaneIncrementAttr::get(parser.getContext(), groupSize, step, aligned);
}

void IREE::GPU::LaneIncrementAttr::print(AsmPrinter &printer) const {
  printer << "<" << getGroupSize();
  if (getStep() != 1) {
    printer << ", step = " << getStep();
  }
  if (getAligned()) {
    printer << ", aligned";
  }
  printer << ">";
}

//===----------------------------------------------------------------------===//
// GPU Pipeline Attribute
//===----------------------------------------------------------------------===//

static GPUPipelineBuilder &getGPUPipelineBuilderStorage() {
  static GPUPipelineBuilder builder = nullptr;
  return builder;
}

static GPUConstraintEmitter &getGPUConstraintEmitterStorage() {
  static GPUConstraintEmitter emitter = nullptr;
  return emitter;
}

void registerGPUPipelineCallbacks(GPUPipelineBuilder builder,
                                  GPUConstraintEmitter constraintEmitter) {
  // Called exactly once during global pass registration
  // (ireeCompilerGlobalInitialize -> registerCodegenLLVMGPUPasses).
  if (builder) {
    getGPUPipelineBuilderStorage() = builder;
  }
  if (constraintEmitter) {
    getGPUConstraintEmitterStorage() = constraintEmitter;
  }
}

LogicalResult
PipelineAttr::buildPipeline(OpPassManager &pm,
                            const CodegenPipelineOptions *options) const {
  GPUPipelineBuilder builder = getGPUPipelineBuilderStorage();
  assert(builder && "no GPU pipeline builder registered; ensure "
                    "registerCodegenLLVMGPUPasses() was called");
  return builder(*this, pm, options);
}

LogicalResult
PipelineAttr::emitConstraints(ArrayRef<Operation *> rootOps) const {
  GPUConstraintEmitter emitter = getGPUConstraintEmitterStorage();
  if (!emitter || rootOps.empty()) {
    return success();
  }
  return emitter(*this, rootOps);
}

FailureOr<Attribute> PipelineAttr::materializeConfigurationAttr(
    StringRef attrName, DictionaryAttr knobs,
    function_ref<FailureOr<Attribute>(Attribute)> materializeAttr,
    function_ref<InFlightDiagnostic()> emitError) const {
  if (attrName != kCompilationInfoOutputName) {
    if (Attribute attr = knobs.get(attrName)) {
      return materializeAttr(attr);
    }
    emitError() << "constraints op has no '" << attrName
                << "' entry in its knobs dictionary";
    return failure();
  }

  MLIRContext *ctx = getContext();
  FailureOr<Attribute> materializedAttr =
      materializeAttr(getCompilationInfoMaterializationInputs(ctx, knobs));
  if (failed(materializedAttr)) {
    return failure();
  }

  auto materializedKnobs = cast<DictionaryAttr>(*materializedAttr);
  DictionaryAttr loweringSource = materializedKnobs;
  bool flatLoweringSource = true;
  if (auto nestedLowering = dyn_cast_if_present<DictionaryAttr>(
          materializedKnobs.get(kConfigAttrName))) {
    loweringSource = nestedLowering;
    flatLoweringSource = false;
  }

  SmallVector<NamedAttribute> loweringEntries;
  appendLoweringConfigEntries(loweringSource, loweringEntries,
                              flatLoweringSource);
  auto loweringConfig =
      LoweringConfigAttr::get(ctx, DictionaryAttr::get(ctx, loweringEntries));

  DictionaryAttr translationSource = materializedKnobs;
  bool nestedTranslationSource = false;
  if (auto nestedTranslation = dyn_cast_if_present<DictionaryAttr>(
          materializedKnobs.get(kTranslationInfoKey))) {
    translationSource = nestedTranslation;
    nestedTranslationSource = true;
  }

  SmallVector<int64_t> workgroupSize;
  if (Attribute workgroupSizeAttr = translationSource.get(kWorkgroupSizeKey)) {
    workgroupSize = extractI64Array(workgroupSizeAttr);
  }

  int64_t subgroupSize = 0;
  if (Attribute subgroupSizeAttr = translationSource.get(kSubgroupSizeKey)) {
    subgroupSize = extractI64(subgroupSizeAttr);
  }

  SymbolRefAttr codegenSpec;
  if (Attribute codegenSpecAttr = translationSource.get(kCodegenSpecKey)) {
    codegenSpec = dyn_cast<SymbolRefAttr>(codegenSpecAttr);
    if (!codegenSpec) {
      emitError() << "expected '" << kCodegenSpecKey
                  << "' to be a symbol reference";
      return failure();
    }
  }

  DictionaryAttr configuration;
  if (nestedTranslationSource) {
    configuration = buildTranslationConfiguration(ctx, translationSource);
  } else if (Attribute gpuPipelineOptions =
                 materializedKnobs.get(kGpuPipelineOptionsKey)) {
    FailureOr<DictionaryAttr> optionsConfig =
        buildGpuPipelineOptionsConfiguration(ctx, gpuPipelineOptions,
                                             emitError);
    if (failed(optionsConfig)) {
      return failure();
    }
    configuration = *optionsConfig;
  }

  Attribute pipeline = *this;
  auto translationInfo = IREE::Codegen::TranslationInfoAttr::getChecked(
      emitError, ctx, pipeline, codegenSpec, ArrayRef<int64_t>(workgroupSize),
      subgroupSize, configuration);
  if (!translationInfo) {
    return failure();
  }
  return IREE::Codegen::CompilationInfoAttr::get(ctx, loweringConfig,
                                                 translationInfo);
}

//===----------------------------------------------------------------------===//
// SPIRV Pipeline Attribute
//===----------------------------------------------------------------------===//

static SPIRVPipelineBuilder &getSPIRVPipelineBuilderStorage() {
  static SPIRVPipelineBuilder builder = nullptr;
  return builder;
}

void registerSPIRVPipelineBuilder(SPIRVPipelineBuilder builder) {
  // Called exactly once during global pass registration
  // (ireeCompilerGlobalInitialize -> registerCodegenSPIRVPasses).
  getSPIRVPipelineBuilderStorage() = builder;
}

LogicalResult
SPIRVPipelineAttr::buildPipeline(OpPassManager &pm,
                                 const CodegenPipelineOptions *options) const {
  SPIRVPipelineBuilder builder = getSPIRVPipelineBuilderStorage();
  assert(builder && "no SPIRV pipeline builder registered; ensure "
                    "registerCodegenSPIRVPasses() was called");
  return builder(*this, pm, options);
}

//===----------------------------------------------------------------------===//
// Attribute Registration
//===----------------------------------------------------------------------===//

void IREEGPUDialect::registerAttributes() {
  addAttributes<
#define GET_ATTRDEF_LIST
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp.inc" // IWYU pragma: keep
      >();
}

} // namespace mlir::iree_compiler::IREE::GPU
