// 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
//===- CPUEncodingExternalModels.cpp --------------------------------------===//
//
// This file implements the following interfaces for CPU backends and the VMVX
// backend:
//
// - IREE::Encoding::LayoutResolverAttr
// - IREE::Encoding::SerializableAttr
// - IREE::Encoding::LayoutMaterializerAttr
// - IREE::Codegen::PackedLayoutMaterializerAttr
// - VerifiableTensorEncoding
//
// In these backends, we transpose narrow-N into narrow-M
// for a combination of reasons:
//
//   1. As linalg.matmul materializes into linalg.mmt4d, which has a transposed
//      RHS and therefore LHS<->RHS symmetry, transposeNarrowN is easy to
//      implement at that level.
//   2. We use ukernels, and this allows writing 2x fewer narrow ukernels.
//   3. Heuristics for cache-friendly dispatch tiling can get complex on CPU,
//      so it is nice that they have fewer narrow cases to consider.
//
// The only current exception to this is Arm SVE. It currently adheres to the
// canonical form of scalable vectorisation and keeps the N dimension to be
// scalable.
// This transposition is made easier by (and was all along part of the idea in)
// the RHS-transposition in mmt4d (the t in mmt4d), as generally with matrix
// multiplication
//
//   B * Transpose(A) == Transpose( A * Transpose(B) )
//
// so in mmt4d terms
//
//   mmt4d(B, A) == Transpose(mmt4d(A, B))
//
//===---------------------------------------------------------------------===//

#include "iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.h"

#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.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/Utils.h"
#include "iree/compiler/Codegen/ExternalInterfaces/Utils.h"
#include "iree/compiler/Codegen/Utils/CPUUtils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
#include "iree/compiler/Dialect/Encoding/Utils/Utils.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/MatchUtils.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/InterleavedRange.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"

#define DEBUG_TYPE "iree-codegen-materialize-encoding"

namespace mlir::iree_compiler::IREE::CPU {

using IREE::Codegen::MaterializeEncodingInfo;
using IREE::Codegen::TileMxNxK;
using IREE::Codegen::TileOCxIC;

namespace {

//===----------------------------------------------------------------------===//
// Utilities.
//===----------------------------------------------------------------------===//

/// Determines which tile dimensions should be scalable for a given target.
/// For AArch64 SVE/SVE2 and RISC-V with the V extension (when not using static
/// vectors), the N dimension is made scalable to take advantage of scalable
/// vector registers.
/// Returns failure if scalable tiling is not applicable for the target.
static FailureOr<IREE::Codegen::ScalableTileFlags>
getScalableTileFlags(linalg::ContractionDimensions cDims,
                     IREE::Encoding::EncodingAttr encoding,
                     DictionaryAttr config) {
  // TODO(egebeysel): I think this isScalable*Enabled flag should be temporary
  // and the temporary SME flag should probably come next to it.
  if (!isScalableVectorizationEnabled()) {
    LDBG() << "Scalable vectorization is not enabled!";
    return failure();
  }

  // Check if target supports scalable vectors.
  bool hasAArch64ScalableSupport =
      isAArch64(config) &&
      (hasFeature(config, "+sve") || hasFeature(config, "+sve2"));
  bool hasRISCVScalableSupport = isRISCV64(config) && hasFeature(config, "+v");

  if (!hasAArch64ScalableSupport && !hasRISCVScalableSupport) {
    LDBG() << "Target does not support scalable vectors!";
    return failure();
  }

  std::optional<unsigned> mDim =
      cDims.m.empty() ? std::nullopt
                      : encoding.mapDimToOperandIndex(cDims.m[0]);
  std::optional<unsigned> nDim =
      cDims.n.empty() ? std::nullopt
                      : encoding.mapDimToOperandIndex(cDims.n[0]);
  std::optional<unsigned> kDim = encoding.mapDimToOperandIndex(cDims.k[0]);
  IREE::Codegen::ScalableTileFlags scalableTiles;

  // M dimension: fixed for both AArch64 SVE and RISC-V V extension.
  // TODO(egebeysel): Add logic for SME which may have scalable M.
  if (mDim.has_value()) {
    if (hasFeature(config, "+sme")) {
      LDBG() << "SME with data-tiling is not supported yet!";
      return failure();
    }
    scalableTiles.push_back(false);
  }

  // N dimension: scalable for both AArch64 SVE/SVE2 and RISC-V V extension
  // (when using scalable vectors). This allows the tile to grow with the
  // vector register length.
  if (nDim.has_value()) {
    scalableTiles.push_back(hasAArch64ScalableSupport ||
                            hasRISCVScalableSupport);
  }

  // K dimension: fixed for both targets. The reduction dimension does not
  // benefit from scalability in the same way as the N dimension.
  if (kDim.has_value()) {
    scalableTiles.push_back(false);
  }
  return scalableTiles;
}

static void transposeInPlace(MaterializeEncodingInfo &info) {
  // Vector cases: nothing to do.
  if (info.innerTileSizes.size() < 2) {
    return;
  }
  // Not a vector case, so all three arrays in `info` have size at least 2,
  // outerDimsPerm may have size 3 if there is a batch dimension, but in all
  // cases, the last 2 entries of each array are M and N, not batch.
  auto transpose = [](auto &a) { std::swap(a[a.size() - 2], a[a.size() - 1]); };
  transpose(info.innerDimsPos);
  transpose(info.innerTileSizes);
  transpose(info.outerDimsPerm);
  if (info.scalableTiles) {
    transpose(info.scalableTiles.value());
  }
}

static RankedTensorType
getExpandedType(RankedTensorType type, bool isBatched, bool isTransposed,
                SmallVectorImpl<ReassociationIndices> &ri) {
  if (!isBatched) {
    ri.assign({{0, 1}, {2, 3}});
    if (!isTransposed) {
      return RankedTensorType::get(
          {1, type.getDimSize(0), 1, type.getDimSize(1)},
          type.getElementType());
    }
    return RankedTensorType::get({type.getDimSize(0), 1, type.getDimSize(1), 1},
                                 type.getElementType());
  }

  ri.assign({{0}, {1, 2}, {3, 4}});
  if (!isTransposed) {
    return RankedTensorType::get(
        {type.getDimSize(0), 1, type.getDimSize(1), 1, type.getDimSize(2)},
        type.getElementType());
  }
  return RankedTensorType::get(
      {type.getDimSize(0), type.getDimSize(1), 1, type.getDimSize(2), 1},
      type.getElementType());
}

/// Given an input Value and a desired output element type, create and return
/// an element-wise linalg::GenericOp that extends the input Value to the
/// output element type. Returns `input` if casting is not needed.
static Value createElementWiseExtUIOp(OpBuilder &builder, Value input,
                                      Location loc, Type outElemType) {
  auto inputType = cast<RankedTensorType>(input.getType());
  if (inputType.getElementType() == outElemType) {
    return input;
  }
  SmallVector<AffineMap> maps(
      2, builder.getMultiDimIdentityMap(inputType.getRank()));
  SmallVector<utils::IteratorType> iteratorTypes(inputType.getRank(),
                                                 utils::IteratorType::parallel);
  auto castedType = inputType.clone(outElemType);
  SmallVector<OpFoldResult> inputMixedSizes =
      tensor::getMixedSizes(builder, loc, input);
  Value init =
      tensor::EmptyOp::create(builder, loc, inputMixedSizes, outElemType);
  return linalg::GenericOp::create(
             builder, loc, castedType, input, init, maps, iteratorTypes,
             [&](OpBuilder &b, Location nestedLoc, ValueRange args) {
               Value castRes =
                   arith::ExtUIOp::create(b, nestedLoc, outElemType, args[0])
                       ->getResult(0);
               linalg::YieldOp::create(b, nestedLoc, castRes);
             })
      .getResult(0);
}

/// If needed, expand and the input Value, and return the resulting input with
/// the canonical mmt4d input shape. If the input element type is unsigned,
/// create a producer Linalg::GenericOp on the input that unsigned extends the
/// input to the output element type. This extension is required to keep the
/// unsignedness information on the input for ukernels. If `transpose` is true,
/// the `linalgOp`'s indexing maps are transposed.
static Value getMmt4dOperand(Value value, linalg::LinalgOp linalgOp,
                             bool transpose, OpBuilder &builder,
                             SmallVectorImpl<ReassociationIndices> &ri,
                             ArrayRef<Type> elemTypes, int operandIdx) {
  assert(linalgOp.getNumDpsInputs() == 2);
  assert(linalgOp.getNumDpsInits() == 1);
  auto cDims = linalg::inferContractionDims(linalgOp);
  Location loc = linalgOp->getLoc();
  Value expandedValue = value;
  // If vecmat with non-rhs operandIdx or matvec with non-lhs operandIdx, the
  // operand is a vector and must be extended
  if ((cDims->m.empty() && operandIdx != 1) ||
      (cDims->n.empty() && operandIdx != 0)) {
    auto type = cast<RankedTensorType>(value.getType());
    RankedTensorType newType = getExpandedType(
        type, /*isBatched=*/!cDims->batch.empty(),
        /*isTransposed=*/operandIdx == 2 && (transpose ^ cDims->n.empty()), ri);
    expandedValue =
        tensor::ExpandShapeOp::create(builder, loc, newType, value, ri);
  }
  if (elemTypes[operandIdx].isUnsignedInteger()) {
    return createElementWiseExtUIOp(builder, expandedValue, loc,
                                    elemTypes.back());
  }
  return expandedValue;
}

/// Returns the best TileMxNxK from `enumeratedTiles` pool. If the
/// `hostDefinedUpperBound` is not empty, the chosen tile sizes can not be
/// greater than the values.
TileMxNxK chooseMatmulTile(ArrayRef<TileMxNxK> enumeratedTiles,
                           IREE::Encoding::MatmulNarrowDim narrowDim) {
  // Handle narrow-N by transposing to reduce to narrow-M. Note: the
  // enumeratedTiles currently only enumerate narrow-M cases.
  if (narrowDim.isN() && !isScalableVectorizationEnabled()) {
    narrowDim.dim = IREE::Encoding::MatmulNarrowDim::Dim::M;
    TileMxNxK tile = chooseMatmulTile(enumeratedTiles, narrowDim);
    std::swap(tile.M, tile.N);
    return tile;
  }
  // Handle kDynamic: currently this is only used with VMVX, where there is only
  // one enumerated tile and it has all three M/N/K dimensions dynamic, so for
  // now we only support that. Generalize that as needed when more dynamic tile
  // sizes are used outside of VMVX, e.g. perhaps some day with Arm SVE. Decide
  // how to incorporate the handling of kDynamic in the cost-model evaluation
  // below to decide when to prefer a dynamic vs a static tile shape.
  for (auto tile : enumeratedTiles) {
    if (ShapedType::isDynamic(tile.M) || ShapedType::isDynamic(tile.N) ||
        ShapedType::isDynamic(tile.K)) {
      assert(enumeratedTiles.size() == 1);
      assert(ShapedType::isDynamic(tile.M) && ShapedType::isDynamic(tile.N) &&
             ShapedType::isDynamic(tile.K));
      return tile;
    }
  }
  // We're going to "rate" the enumerated tiles.
  struct RatedTileMxNxK : TileMxNxK {
    RatedTileMxNxK() {}
    RatedTileMxNxK(TileMxNxK tile) : TileMxNxK(tile) {}
    // Penalize tiles that are wider in the M dimension than matmulNarrowM.
    int64_t paddingPenalty = 0;
    // Favor larger tiles, as long as they still minimize paddingPenalty.
    int64_t productMxNxK = 0;
  };
  SmallVector<RatedTileMxNxK> ratedTiles;
  ratedTiles.reserve(enumeratedTiles.size());
  int64_t bestPaddingPenalty = INT64_MAX;
  for (auto tile : enumeratedTiles) {
    RatedTileMxNxK ratedTile(tile);
    ratedTile.paddingPenalty = 0;
    // If we are choosing a tile for a narrow-M case, we want to minimize
    // padding along the M dimension.
    // The PowerOf2Ceil is so that we are OK with padding up to the next
    // power of two, we just try to avoid padding beyond that. For example,
    // if matmulNarrowM==7 and we have enumerated tiles with M=8,4,2,1, we
    // are OK with the tile that has M==8 even though it requires some padding.
    // Otherwise, we would be penalizing the tiles with M==8,4,2 and we would
    // end up selecting the vecmat tile (M==1) for that case!
    if (narrowDim) {
      ratedTile.paddingPenalty =
          std::max<int64_t>(tile.M - llvm::PowerOf2Ceil(narrowDim.size), 0);
    }
    ratedTile.productMxNxK = tile.M * tile.N * tile.K;
    ratedTiles.push_back(ratedTile);
    LDBG() << "candidate: "
           << llvm::interleaved(ArrayRef{tile.M, tile.N, tile.K})
           << " penalty:" << ratedTile.paddingPenalty;
    bestPaddingPenalty = std::min(bestPaddingPenalty, ratedTile.paddingPenalty);
  }
  RatedTileMxNxK bestRatedTile;
  for (auto ratedTile : ratedTiles) {
    // Choose only among tiles that minimize paddingPenalty. Among those,
    // maximize productMxNxK.
    if (ratedTile.paddingPenalty == bestPaddingPenalty &&
        bestRatedTile.productMxNxK < ratedTile.productMxNxK) {
      bestRatedTile = ratedTile;
    }
  }
  // Sanity check. This assert can only fail if there's a programming mistake
  // locally here.
  assert(bestRatedTile.paddingPenalty == bestPaddingPenalty);
  LDBG() << "bestRatedTile: "
         << llvm::interleaved(
                ArrayRef{bestRatedTile.M, bestRatedTile.N, bestRatedTile.K})
         << " penalty:" << bestRatedTile.paddingPenalty;
  return bestRatedTile;
}

static bool getEnableInnerTiledFromConfig(DictionaryAttr config) {
  Attribute attr = config.get("enable_inner_tiled");
  if (auto battr = dyn_cast_if_present<BoolAttr>(attr)) {
    return battr.getValue();
  }
  return false;
}

/// Returns the matmul {M, N, K} tile shape covered by a CPU DataTiledMMAAttr.
/// The swizzle's `expandShape` is partitioned per its convention (see
/// `getIntrinsicSwizzle`): LHS is `[M dims, K dims]`, RHS is `[N dims, K
/// dims]`. So M is the product of LHS group 0, K is the product of LHS
/// group 1, and N is the product of RHS group 0. Unlike reading from the
/// vector tile type, this works regardless of the swizzle's permutation.
static IREE::Codegen::TileMxNxK getTileMxNxK(IREE::CPU::DataTiledMMAAttr mma) {
  auto groupProduct = [](ArrayRef<IREE::Codegen::TileSwizzle::Dim> group) {
    int64_t product = 1;
    for (const auto &d : group) {
      product *= d.size();
    }
    return product;
  };
  IREE::Codegen::TileSwizzle lhsSwizzle = IREE::CPU::getSwizzle(mma, 0);
  IREE::Codegen::TileSwizzle rhsSwizzle = IREE::CPU::getSwizzle(mma, 1);
  return IREE::Codegen::TileMxNxK{groupProduct(lhsSwizzle.expandShape()[0]),
                                  groupProduct(rhsSwizzle.expandShape()[0]),
                                  groupProduct(lhsSwizzle.expandShape()[1])};
}

/// Returns the set of `+feature` strings that must all be present in the
/// target config for `intr` to be usable. Most AVX-512 intrinsics map
/// 1:1 to a single `+feature`, but the AVX2 f32 MMA intrinsics require both
/// `+avx2` and `+fma` (FMA3 is a separate ISA extension from AVX2 itself,
/// and our AVX2 intrinsics lower to `vfmadd...`).
static SmallVector<StringRef>
getMmaIntrinsicRequiredFeatures(IREE::CPU::MMAIntrinsic intr) {
  using IREE::CPU::MMAIntrinsic;
  switch (intr) {
  case MMAIntrinsic::MMA_X86_AVX2_FMA_1x8x1_F32_F32:
  case MMAIntrinsic::MMA_X86_AVX2_FMA_8x1x1_F32_F32:
    return {"+avx2", "+fma"};
  case MMAIntrinsic::MMA_X86_AVX512_1x8x1_F64_F64:
  case MMAIntrinsic::MMA_X86_AVX512_8x1x1_F64_F64:
  case MMAIntrinsic::MMA_X86_AVX512_1x16x1_F32_F32:
  case MMAIntrinsic::MMA_X86_AVX512_16x1x1_F32_F32:
  case MMAIntrinsic::MMA_X86_AVX512_1x16x1_F32_F16_CASTF32:
  case MMAIntrinsic::MMA_X86_AVX512_16x1x1_F32_F16_CASTF32:
  case MMAIntrinsic::MMA_X86_AVX512_1x16x2_I32_I16:
  case MMAIntrinsic::MMA_X86_AVX512_16x1x2_I32_I16:
  case MMAIntrinsic::MMA_X86_AVX512_1x16x2_I32_I8_CASTI16:
  case MMAIntrinsic::MMA_X86_AVX512_16x1x2_I32_I8_CASTI16:
    return {"+avx512f"};
  case MMAIntrinsic::MMA_X86_AVX512FP16_1x32x1_F16_F16:
  case MMAIntrinsic::MMA_X86_AVX512FP16_32x1x1_F16_F16:
    return {"+avx512fp16"};
  case MMAIntrinsic::MMA_X86_AVX512BF16_1x16x2_F32_BF16:
  case MMAIntrinsic::MMA_X86_AVX512BF16_16x1x2_F32_BF16:
    return {"+avx512bf16"};
  case MMAIntrinsic::MMA_X86_AVX512VNNI_1x16x2_I32_I16:
  case MMAIntrinsic::MMA_X86_AVX512VNNI_16x1x2_I32_I16:
  case MMAIntrinsic::MMA_X86_AVX512VNNI_1x16x2_I32_I8_CASTI16:
  case MMAIntrinsic::MMA_X86_AVX512VNNI_16x1x2_I32_I8_CASTI16:
  case MMAIntrinsic::MMA_X86_AVX512VNNI_1x16x4_I32_UI8_I8:
  case MMAIntrinsic::MMA_X86_AVX512VNNI_16x1x4_I32_I8_UI8:
  case MMAIntrinsic::MMA_X86_AVX512VNNI_16x16x2_I32_I8_CASTI16:
    return {"+avx512vnni"};
  default:
    return {};
  }
}

/// Returns the `MMA_GENERIC_SCALAR_1x1x1_REG*` variant for `config`. Almost
/// every 64-bit CPU has at least 16 architectural registers we could put
/// scalar tiles into (GPRs or lane 0 of a SIMD register, depending on what
/// LLVM picks), so we pick `_REG16` on 64-bit ISAs and `_REG8` on 32-bit
/// ones. This is a performance heuristic only — too high a budget would
/// spill, too low would underutilize — and the generic-scalar fallback is
/// slow either way.
static IREE::CPU::MMAIntrinsic
pickGenericScalarMMAForTarget(DictionaryAttr config) {
  using IREE::CPU::MMAIntrinsic;
  std::optional<llvm::Triple> triple = getTargetTriple(config);
  bool is64Bit = triple ? triple->isArch64Bit() : true;
  return is64Bit ? MMAIntrinsic::MMA_GENERIC_SCALAR_1x1x1_REG16
                 : MMAIntrinsic::MMA_GENERIC_SCALAR_1x1x1_REG8;
}

/// Validates `arr` per the conventions in IREECPUEnums.td. Used inside an
/// `assert()` at the end of `getMmaIntrinsicsForTargetConfig` so it only
/// runs in debug builds — it constructs `DataTiledMMAAttr` instances to
/// inspect tile shapes/types, which would be too expensive to do
/// unconditionally on every cost-model query.
///
/// Conditions:
///   1. Strictly increasing in enum value (subsumes "no duplicates" and
///      rules out misordered pairs).
///   2. By the bit-0 convention from IREECPUEnums.td, every M↔N-swapped
///      entry (bit 0 set) is immediately preceded by its natural sibling
///      (`prev == v ^ 1`).
///   3. For each (natural, swapped) sibling pair, the swapped entry's
///      tile type has M↔N exchanged and LHS↔RHS element types exchanged
///      vs the natural's; K and ACC element type are unchanged.
///   4. A natural entry without a swapped sibling (singleton) must be
///      self-symmetric: square shape AND identical LHS/RHS element types
///      — otherwise the swapped orientation would be a distinct operation
///      and the array would be incomplete.
///
/// The type-polymorphic generic-scalar 1×1×1 family is square and
/// type-polymorphic, so it satisfies (4) vacuously. We skip the
/// `DataTiledMMAAttr`-based extraction for it (passing null element types
/// to the attr would invalidate the resulting vector types).
static bool isMmaIntrinsicArrayValid(MLIRContext *ctx,
                                     ArrayRef<IREE::CPU::MMAIntrinsic> arr) {
  using IREE::CPU::DataTiledMMAAttr;
  using IREE::CPU::MMAIntrinsic;

  // Conventions (1) and (2): order, strict increase, swap-pair structure.
  if (!arr.empty() && (static_cast<uint32_t>(arr[0]) & 1) != 0) {
    return false;
  }
  for (size_t i = 1; i < arr.size(); ++i) {
    uint32_t v = static_cast<uint32_t>(arr[i]);
    uint32_t prev = static_cast<uint32_t>(arr[i - 1]);
    if (prev >= v) {
      return false;
    }
    if ((v & 1) != 0 && prev != (v ^ 1)) {
      return false;
    }
  }

  // (3) and (4) need types *with* signedness so e.g. vpdpbusd's natural
  // `ui8` LHS visibly swaps with its sibling's `ui8` RHS. Read directly
  // from `getABCElementTypes`; `getUndistributedTileTypes` would strip it.
  auto getShapeAndTypes = [&](MMAIntrinsic intr) {
    auto mnk = IREE::CPU::getRowMajorTilesMNKShape(intr);
    assert(mnk && "validator only handles row-major-tile intrinsics");
    auto [m, n, k] = *mnk;
    auto [lhs, rhs, acc] = IREE::CPU::getABCElementTypes(ctx, intr);
    return std::make_tuple(m, n, k, lhs, rhs, acc);
  };

  for (size_t i = 0; i < arr.size(); ++i) {
    if ((static_cast<uint32_t>(arr[i]) & 1) != 0) {
      continue; // swapped entry — checked when we hit its natural sibling.
    }
    if (IREE::CPU::isGenericScalar(arr[i])) {
      continue; // type-polymorphic 1×1×1: vacuously valid.
    }
    auto [mn, nn, kn, lhsn, rhsn, accn] = getShapeAndTypes(arr[i]);
    bool hasSibling =
        (i + 1 < arr.size()) && (static_cast<uint32_t>(arr[i + 1]) ==
                                 (static_cast<uint32_t>(arr[i]) | 1));
    if (hasSibling) {
      auto [ms, ns, ks, lhss, rhss, accs] = getShapeAndTypes(arr[i + 1]);
      // Condition (3): siblings have M↔N and LHS↔RHS exchanged; K and ACC
      // element type are unchanged.
      if (ms != nn || ns != mn || ks != kn) {
        return false;
      }
      if (lhss != rhsn || rhss != lhsn || accs != accn) {
        return false;
      }
    } else {
      // Condition (4): singleton — must be self-symmetric.
      if (mn != nn || lhsn != rhsn) {
        return false;
      }
    }
  }
  return true;
}

/// Returns the `MMAIntrinsic` cases potentially usable for `config`: the x86
/// architecture-specific intrinsics whose required ISA extensions are all
/// present (both natural and M↔N-swapped orientations enumerated as separate
/// values), followed by one of the architecture-agnostic type-polymorphic
/// `MMA_GENERIC_SCALAR_1x1x1_REG*` fallback variants (the one whose register
/// budget matches `config`'s target). The 1×1×1 generic intrinsic naturally
/// loses to any real intrinsic that fits, so it only wins as a fallback
/// when no real MMA covers the requested element types.
///
/// There must be no early-return path that omits the trailing
/// generic-scalar push — `out` is asserted to be non-empty / well-formed at
/// the bottom and downstream callers rely on always seeing at least the
/// generic fallback. Place the generic at the end (rather than the front)
/// because: (a) its enum value sits above the architecture-specific ones,
/// so trailing it keeps `out` strictly increasing, which the validator
/// expects, and (b) trailing it lets the cost-model's tie-break (first
/// wins on equal score) prefer real intrinsics over the generic fallback.
static SmallVector<IREE::CPU::MMAIntrinsic>
getMmaIntrinsicsForTargetConfig(DictionaryAttr config) {
  using IREE::CPU::MMAIntrinsic;
  if (!config) {
    return {};
  }
  SmallVector<MMAIntrinsic> out;
  if (isX86(config)) {
    static const MMAIntrinsic kAllX86[] = {
        MMAIntrinsic::MMA_X86_AVX2_FMA_1x8x1_F32_F32,
        MMAIntrinsic::MMA_X86_AVX2_FMA_8x1x1_F32_F32,
        MMAIntrinsic::MMA_X86_AVX512_1x8x1_F64_F64,
        MMAIntrinsic::MMA_X86_AVX512_8x1x1_F64_F64,
        MMAIntrinsic::MMA_X86_AVX512_1x16x1_F32_F32,
        MMAIntrinsic::MMA_X86_AVX512_16x1x1_F32_F32,
        MMAIntrinsic::MMA_X86_AVX512_1x16x1_F32_F16_CASTF32,
        MMAIntrinsic::MMA_X86_AVX512_16x1x1_F32_F16_CASTF32,
        MMAIntrinsic::MMA_X86_AVX512FP16_1x32x1_F16_F16,
        MMAIntrinsic::MMA_X86_AVX512FP16_32x1x1_F16_F16,
        MMAIntrinsic::MMA_X86_AVX512BF16_1x16x2_F32_BF16,
        MMAIntrinsic::MMA_X86_AVX512BF16_16x1x2_F32_BF16,
        MMAIntrinsic::MMA_X86_AVX512_1x16x2_I32_I16,
        MMAIntrinsic::MMA_X86_AVX512_16x1x2_I32_I16,
        MMAIntrinsic::MMA_X86_AVX512VNNI_1x16x2_I32_I16,
        MMAIntrinsic::MMA_X86_AVX512VNNI_16x1x2_I32_I16,
        MMAIntrinsic::MMA_X86_AVX512_1x16x2_I32_I8_CASTI16,
        MMAIntrinsic::MMA_X86_AVX512_16x1x2_I32_I8_CASTI16,
        MMAIntrinsic::MMA_X86_AVX512VNNI_1x16x2_I32_I8_CASTI16,
        MMAIntrinsic::MMA_X86_AVX512VNNI_16x1x2_I32_I8_CASTI16,
        MMAIntrinsic::MMA_X86_AVX512VNNI_1x16x4_I32_UI8_I8,
        MMAIntrinsic::MMA_X86_AVX512VNNI_16x1x4_I32_I8_UI8,
        MMAIntrinsic::MMA_X86_AVX512VNNI_16x16x2_I32_I8_CASTI16,
    };
    for (MMAIntrinsic intr : kAllX86) {
      SmallVector<StringRef> required = getMmaIntrinsicRequiredFeatures(intr);
      if (required.empty()) {
        continue;
      }
      if (llvm::all_of(required,
                       [&](StringRef f) { return hasFeature(config, f); })) {
        out.push_back(intr);
      }
    }
  }
  out.push_back(pickGenericScalarMMAForTarget(config));
  assert(isMmaIntrinsicArrayValid(config.getContext(), out) &&
         "getMmaIntrinsicsForTargetConfig must return a list satisfying the "
         "conventions in IREECPUEnums.td (strictly increasing, paired "
         "natural/swapped entries, etc.)");
  return out;
}

/// Shape and element-bit-width summary of one intrinsic, with unroll factors
/// all 1. The LHS is (M, K), the RHS is (N, K). Scalable dims (e.g. SVE) are
/// treated as their 128-bit-minimum static shape here; `getRegisterSpaceBytes`
/// applies the matching simplification on the capacity side.
struct IntrinsicInfo {
  int64_t intrinsicM = 0, intrinsicN = 0, intrinsicK = 0;
  int64_t lhsBits = 0, rhsBits = 0, accBits = 0;
};

/// Returns the IntrinsicInfo for `intr` if its ABC element types match
/// `elementTypes`. Returns nullopt otherwise. The
/// `MMA_GENERIC_SCALAR_1x1x1_REG*` family is type-polymorphic — those
/// values match any element-type triple by construction; their `lhsBits` /
/// `rhsBits` / `accBits` stay at the struct's default 0 since the
/// generic-scalar branch of `chooseUnrolling` doesn't read them.
static std::optional<IntrinsicInfo>
getIntrinsicInfo(MLIRContext *ctx, ArrayRef<Type> elementTypes,
                 IREE::CPU::MMAIntrinsic intr) {
  if (IREE::CPU::isGenericScalar(intr)) {
    if (elementTypes.size() != 3 || !elementTypes[0] || !elementTypes[1] ||
        !elementTypes[2]) {
      return std::nullopt;
    }
    IntrinsicInfo info;
    info.intrinsicM = info.intrinsicN = info.intrinsicK = 1;
    return info;
  }
  // Same as the validator above: match against the encoding's
  // `element_types` *with* signedness — vpdpbusd's `ui8` is what
  // distinguishes it from CASTI16 paths.
  auto [lhsTy, rhsTy, accTy] = IREE::CPU::getABCElementTypes(ctx, intr);
  if (!lhsTy || !rhsTy || !accTy) {
    return std::nullopt;
  }
  if (lhsTy != elementTypes[0] || rhsTy != elementTypes[1] ||
      accTy != elementTypes[2]) {
    return std::nullopt;
  }
  auto mnk = IREE::CPU::getRowMajorTilesMNKShape(intr);
  if (!mnk) {
    return std::nullopt;
  }
  auto [m, n, k] = *mnk;
  IntrinsicInfo info;
  info.intrinsicM = m;
  info.intrinsicN = n;
  info.intrinsicK = k;
  info.lhsBits = lhsTy.getIntOrFloatBitWidth();
  info.rhsBits = rhsTy.getIntOrFloatBitWidth();
  info.accBits = accTy.getIntOrFloatBitWidth();
  return info;
}

/// Phase 1 of `chooseCpuInnerTiledMmaForEncoding`: pick the intrinsic. We
/// maximize `usefulOps` per invocation,
///   usefulOps = min(intrinsicM, M) * min(intrinsicN, N) * intrinsicK
/// where the `min` clamps an intrinsicM/intrinsicN that exceeds a narrow
/// static matmul dim — i.e. it charges the intrinsic for the padding it
/// would force. A dynamic matmul dim is treated as "not narrow" (the full
/// intrinsic size counts toward usefulOps). Ties are broken by the raw
/// intrinsic size `intrinsicM*intrinsicN*intrinsicK` (prefer the bigger
/// intrinsic: it gives the Phase-2 register budget more tiles to play with)
/// and, finally, by iteration order. The natural and M↔N-swapped
/// orientations of each hardware MMA appear as distinct enum values in the
/// candidate list, so the choice between them falls out of this same
/// scoring loop. `K` is never narrow for optimization purposes, so it never
/// enters the `min`.
///
/// Returns nullopt if no compatible intrinsic exists.
static IREE::CPU::MMAIntrinsic
chooseIntrinsic(MLIRContext *ctx, ArrayRef<Type> elementTypes,
                DictionaryAttr config,
                const IREE::Encoding::BxMxNxKxKb &matmulSizes) {
  auto usefulSize = [](int64_t matmulSize, int64_t intrinsicSize) -> int64_t {
    return ShapedType::isDynamic(matmulSize)
               ? intrinsicSize
               : std::min(intrinsicSize, matmulSize);
  };
  IREE::CPU::MMAIntrinsic best = IREE::CPU::MMAIntrinsic::None;
  std::pair<int64_t, int64_t> bestScore = {-1, -1};
  for (IREE::CPU::MMAIntrinsic intr : getMmaIntrinsicsForTargetConfig(config)) {
    std::optional<IntrinsicInfo> info =
        getIntrinsicInfo(ctx, elementTypes, intr);
    if (!info) {
      continue;
    }
    int64_t usefulOps = usefulSize(matmulSizes.M, info->intrinsicM) *
                        usefulSize(matmulSizes.N, info->intrinsicN) *
                        info->intrinsicK;
    int64_t rawOps = info->intrinsicM * info->intrinsicN * info->intrinsicK;
    std::pair<int64_t, int64_t> score = {usefulOps, rawOps};
    if (score > bestScore) {
      bestScore = score;
      best = intr;
    }
  }
  return best;
}

/// Per-dim cap on the unroll factor: how many intrinsic-sized tiles fit in
/// the matmul dim. For dynamic matmul dims we return `fallback` (the
/// register budget), which the budget check in the enumeration loop bounds
/// anyway.
static int64_t unrollCap(int64_t matmulSize, int64_t intrinsicSize,
                         int64_t fallback) {
  if (ShapedType::isDynamic(matmulSize)) {
    return fallback;
  }
  return std::max<int64_t>(1, llvm::divideCeil(matmulSize, intrinsicSize));
}

// Phase 2 of `chooseCpuInnerTiledMmaForEncoding`: for an already-chosen
// intrinsic, pick unroll factors (intrinsicsM, intrinsicsN) such that the
// three tiles (ACC + LHS + RHS) still fit in the target's register budget,
// breaking ties with arithmetic intensity (effM*effN)/(effM+effN) so
// approximately-square tiles win.
//
// "Register budget" depends on what the lowering will use:
//   * Architecture-specific SIMD intrinsics use the vector register file,
//     measured in bits, with element widths from `IntrinsicInfo` so a
//     wider element type costs proportionally more of the budget.
//   * The type-polymorphic `MMA_GENERIC_SCALAR_1x1x1_REG*` lowers to
//     scalar arithmetic, where one element occupies one register (a GPR
//     or lane 0 of a SIMD register, depending on what LLVM picks)
//     regardless of bit width. So the budget is in registers, with all
//     element "widths" treated as 1.
//
// Returns nullopt if no feasible unrolling exists. The returned tuple is
// (intrinsicsM, intrinsicsN, intrinsicsK). When the (im, in) search loop
// finds no candidate that fits the budget — possible for the generic
// intrinsic with a sub-byte LHS/RHS forcing intrinsics_k = 8 against an
// 8-register budget — we fall back to (1, 1) so we still return *some*
// valid unrolling. Some register spill is preferable to failing data-
// tiling outright.
static std::optional<std::tuple<int64_t, int64_t, int64_t>>
chooseUnrolling(MLIRContext *ctx, ArrayRef<Type> elementTypes,
                IREE::CPU::MMAIntrinsic intr,
                const IREE::Encoding::BxMxNxKxKb &matmulSizes) {
  std::optional<IntrinsicInfo> info = getIntrinsicInfo(ctx, elementTypes, intr);
  if (!info) {
    return std::nullopt;
  }
  int64_t intrinsicsK = 1;
  int64_t budget;
  int64_t accUnit, lhsUnit, rhsUnit;
  if (IREE::CPU::isGenericScalar(intr)) {
    // Real hardware MMA intrinsics bake whatever K-grouping they need into
    // the intrinsic itself. The generic-scalar fallback doesn't, so for
    // sub-byte LHS/RHS types we have to group enough K-elements per
    // contiguous block that a packed group is byte-addressable: pick the
    // smallest power of two K such that K*lhsBits and K*rhsBits are both
    // multiples of 8. K∈{1,2,4,8} is enough to cover every type from i8/f8
    // (K=1) down to i1 (K=8).
    int64_t lhsBits = elementTypes[0].getIntOrFloatBitWidth();
    int64_t rhsBits = elementTypes[1].getIntOrFloatBitWidth();
    while (intrinsicsK <= 8 && (((intrinsicsK * lhsBits) % 8) != 0 ||
                                ((intrinsicsK * rhsBits) % 8) != 0)) {
      intrinsicsK *= 2;
    }
    if (intrinsicsK > 8) {
      return std::nullopt;
    }
    // The budget — chosen at intrinsic-pick time as a function of the
    // target's pointer width — is encoded in the enum value itself.
    budget = IREE::CPU::getGenericScalarRegisterBudget(intr);
    accUnit = lhsUnit = rhsUnit = 1;
  } else {
    budget = IREE::CPU::getRegisterSpaceBytes(intr) * 8;
    accUnit = info->accBits;
    lhsUnit = info->lhsBits;
    rhsUnit = info->rhsBits;
  }
  int64_t capM = unrollCap(matmulSizes.M, info->intrinsicM, budget);
  int64_t capN = unrollCap(matmulSizes.N, info->intrinsicN, budget);
  int64_t accTerm = info->intrinsicM * info->intrinsicN * accUnit;
  int64_t lhsTerm = info->intrinsicM * info->intrinsicK * intrinsicsK * lhsUnit;
  int64_t rhsTerm = info->intrinsicN * info->intrinsicK * intrinsicsK * rhsUnit;
  std::optional<std::pair<int64_t, int64_t>> bestMN;
  double bestIntensity = -1.0;
  // Enumerate intrinsicsM; for each, pick the largest feasible intrinsicsN
  // under the budget and the static N cap. The budget bounds im on its own
  // (im*lhsTerm alone must be < budget), which terminates the loop.
  for (int64_t im = 1; im <= capM; ++im) {
    int64_t remaining = budget - im * lhsTerm;
    if (remaining <= 0) {
      break;
    }
    int64_t inMaxBudget = remaining / (im * accTerm + rhsTerm);
    int64_t in = std::min<int64_t>(capN, inMaxBudget);
    if (in < 1) {
      continue;
    }
    double effM = static_cast<double>(im) * info->intrinsicM;
    double effN = static_cast<double>(in) * info->intrinsicN;
    double intensity = effM * effN / (effM + effN);
    if (intensity > bestIntensity) {
      bestIntensity = intensity;
      bestMN = {im, in};
    }
  }
  // Fall back to (1, 1) if no (im, in) fit — see function comment.
  if (!bestMN) {
    bestMN = {1, 1};
  }
  return std::make_tuple(bestMN->first, bestMN->second, intrinsicsK);
}

// Picks a CPU `DataTiledMMAAttr` for `iree_codegen.inner_tiled` given an
// encoding and target config.
static IREE::CPU::DataTiledMMAAttr
chooseCpuInnerTiledMmaForEncoding(MLIRContext *ctx,
                                  IREE::Encoding::EncodingAttr encoding,
                                  DictionaryAttr config) {
  SmallVector<Type> elementTypes = encoding.getElementTypesArray();
  if (elementTypes.size() != 3) {
    return {};
  }
  FailureOr<IREE::Encoding::BxMxNxKxKb> matmulSizes =
      IREE::Encoding::getEncodingContractionLikeSizes(encoding);
  if (failed(matmulSizes)) {
    return {};
  }
  IREE::CPU::MMAIntrinsic intr =
      chooseIntrinsic(ctx, elementTypes, config, *matmulSizes);
  if (intr == IREE::CPU::MMAIntrinsic::None) {
    return {};
  }
  std::optional<std::tuple<int64_t, int64_t, int64_t>> unroll =
      chooseUnrolling(ctx, elementTypes, intr, *matmulSizes);
  if (!unroll) {
    return {};
  }
  auto [intrinsicsM, intrinsicsN, intrinsicsK] = *unroll;
  // The type-polymorphic generic intrinsic doesn't bake element types into
  // its enum value; we have to pin them down on the attr itself so the
  // lowering and `getABCElementTypes(ctx, attr)` can read them back.
  Type lhsType, rhsType, accType;
  if (IREE::CPU::isGenericScalar(intr)) {
    lhsType = elementTypes[0];
    rhsType = elementTypes[1];
    accType = elementTypes[2];
  }
  return IREE::CPU::DataTiledMMAAttr::get(ctx, intr, intrinsicsM, intrinsicsN,
                                          intrinsicsK, lhsType, rhsType,
                                          accType);
}

/// Lowers a contraction under a `CPUEncodingResolverAttr` with
/// `enable_inner_tiled = true` to an `iree_codegen.inner_tiled` op whose kind
/// is a CPU `data_tiled_mma_layout`. Returns nullptr if no CPU MMA intrinsic
/// is available for the encoding/target.
static Operation *lowerContractionToInnerTiled(
    OpBuilder &builder, linalg::LinalgOp linalgOp, ValueRange operands,
    IREE::Encoding::LayoutMaterializerAttr layoutAttr) {
  if (!linalgOp.hasPureTensorSemantics()) {
    return nullptr;
  }

  auto inputs = linalgOp.getDpsInputOperands();
  auto outputs = linalgOp.getDpsInits();

  auto lhsType = cast<RankedTensorType>(inputs[0]->get().getType());
  auto rhsType = cast<RankedTensorType>(inputs[1]->get().getType());
  auto resultType = cast<RankedTensorType>(outputs[0].getType());
  auto lhsEncoding = IREE::Encoding::getEncodingAttr(lhsType);
  auto rhsEncoding = IREE::Encoding::getEncodingAttr(rhsType);
  auto resultEncoding = IREE::Encoding::getEncodingAttr(resultType);
  if (!lhsEncoding || !rhsEncoding || !resultEncoding) {
    return nullptr;
  }

  if (lhsEncoding.getOperandIndex().getValue() != IREE::Encoding::MATMUL_LHS ||
      rhsEncoding.getOperandIndex().getValue() != IREE::Encoding::MATMUL_RHS ||
      resultEncoding.getOperandIndex().getValue() !=
          IREE::Encoding::MATMUL_RESULT) {
    return nullptr;
  }

  auto cDims = linalg::inferContractionDims(linalgOp);
  if (!cDims->batch.empty()) {
    LDBG() << "inner_tiled lowering: batched contraction not implemented";
    return nullptr;
  }

  MLIRContext *ctx = builder.getContext();
  Location loc = linalgOp.getLoc();
  AffineExpr d0 = builder.getAffineDimExpr(0);
  AffineExpr d1 = builder.getAffineDimExpr(1);
  AffineExpr d2 = builder.getAffineDimExpr(2);
  // RHS uses (d1, d2) i.e. (N_iter, K_iter) outer ordering — mmt4d-style — to
  // match the packed RHS operand shape produced upstream (the RHS pack uses
  // `outer_dims_perm = [1, 0]` so its outer dims come out as N_iter, K_iter,
  // not K_iter, N_iter).
  SmallVector<AffineMap> indexingMaps = {
      AffineMap::get(3, 0, {d0, d2}, ctx),
      AffineMap::get(3, 0, {d1, d2}, ctx),
      AffineMap::get(3, 0, {d0, d1}, ctx),
  };
  SmallVector<utils::IteratorType> iteratorTypes = {
      utils::IteratorType::parallel, utils::IteratorType::parallel,
      utils::IteratorType::reduction};

  DictionaryAttr targetConfig;
  if (auto cpuResolver = dyn_cast<IREE::CPU::CPUEncodingResolverAttr>(
          cast<Attribute>(layoutAttr))) {
    targetConfig = cpuResolver.getConfiguration();
  }

  // Mirrors the GPU DataTiledMMA path: the MMA attribute is determined from
  // the encoding + target config alone (the same inputs used by the encoding
  // materialization side), not reverse-engineered from the already-packed
  // operand tile shapes.
  IREE::CPU::DataTiledMMAAttr chosenKind =
      chooseCpuInnerTiledMmaForEncoding(ctx, resultEncoding, targetConfig);
  if (!chosenKind) {
    LDBG() << "inner_tiled lowering: no CPU DataTiledMMA kind available for "
              "encoding/target";
    return nullptr;
  }

  auto semanticsAttr = IREE::CPU::InnerTiledSemanticsAttr::get(ctx);
  Operation *inner = IREE::Codegen::InnerTiledOp::create(
      builder, loc, ValueRange{operands[0], operands[1]},
      ValueRange{operands[2]}, indexingMaps, iteratorTypes, chosenKind,
      semanticsAttr);
  return inner;
}

Operation *lowerContractionOpWithEncoding(
    OpBuilder &builder, linalg::LinalgOp linalgOp, ValueRange operands,
    IREE::Encoding::LayoutMaterializerAttr layoutAttr) {
  if (!linalgOp.hasPureTensorSemantics()) {
    return nullptr;
  }

  auto inputs = linalgOp.getDpsInputOperands();
  auto outputs = linalgOp.getDpsInits();

  auto lhsType = cast<RankedTensorType>(inputs[0]->get().getType());
  auto rhsType = cast<RankedTensorType>(inputs[1]->get().getType());
  auto resultType = cast<RankedTensorType>(outputs[0].getType());
  auto lhsEncoding = IREE::Encoding::getEncodingAttr(lhsType);
  auto rhsEncoding = IREE::Encoding::getEncodingAttr(rhsType);
  auto resultEncoding = IREE::Encoding::getEncodingAttr(resultType);
  if (!lhsEncoding || !rhsEncoding || !resultEncoding) {
    return nullptr;
  }

  if (lhsEncoding.getOperandIndex().getValue() != IREE::Encoding::MATMUL_LHS ||
      rhsEncoding.getOperandIndex().getValue() != IREE::Encoding::MATMUL_RHS ||
      resultEncoding.getOperandIndex().getValue() !=
          IREE::Encoding::MATMUL_RESULT) {
    return nullptr;
  }

  MaterializeEncodingInfo encodingInfo = {};
  if (auto packedLayoutAttr =
          dyn_cast<IREE::Codegen::PackedLayoutMaterializerAttr>(layoutAttr)) {
    encodingInfo = packedLayoutAttr.getEncodingInfo(
        cast<RankedTensorType>(linalgOp->getResultTypes()[0]));
  }

  if (isIdentityLayout(encodingInfo)) {
    return dropEncodingAndCloneOp(builder, linalgOp,
                                  operands.take_front(inputs.size()),
                                  operands.drop_front(inputs.size()));
  }

  bool transpose = isNarrowNResult(resultEncoding);
  // Do not transpose in case we have scalable tiles.
  transpose &= llvm::none_of(
      encodingInfo.scalableTiles.value_or(IREE::Codegen::ScalableTileFlags{}),
      [](bool flag) { return flag; });
  SmallVector<Type> elemTypes = lhsEncoding.getElementTypesArray();
  SmallVector<ReassociationIndices> ri;
  Value newLhs = getMmt4dOperand(operands[0], linalgOp, transpose, builder, ri,
                                 elemTypes, /*operandIdx=*/0);
  Value newRhs = getMmt4dOperand(operands[1], linalgOp, transpose, builder, ri,
                                 elemTypes, /*operandIdx=*/1);
  Value newResult = getMmt4dOperand(operands[2], linalgOp, transpose, builder,
                                    ri, elemTypes, /*operandIdx=*/2);
  if (transpose) {
    std::swap(newLhs, newRhs);
  }
  Type newResultType = newResult.getType();
  auto cDims = IREE::Encoding::getEncodingContractionDims(lhsEncoding);
  Operation *result;
  if (cDims->batch.empty()) {
    result = linalg::Mmt4DOp::create(builder, linalgOp.getLoc(), newResultType,
                                     ValueRange{newLhs, newRhs},
                                     ValueRange{newResult});
  } else {
    result = linalg::BatchMmt4DOp::create(
        builder, linalgOp.getLoc(), newResultType, ValueRange{newLhs, newRhs},
        ValueRange{newResult});
  }
  if (!ri.empty()) {
    result = tensor::CollapseShapeOp::create(builder, linalgOp->getLoc(),
                                             operands[2].getType(),
                                             result->getResult(0), ri);
  }
  return result;
}

/// Builds 9D linalg.generic implementing data-tiled convolution
/// over already-packed, rank-5/6 operands.
static linalg::GenericOp
buildDataTiledConvOp(OpBuilder &builder, Location loc, Value packedInput,
                     Value packedFilter, Value packedOutput,
                     ArrayRef<int64_t> strides, ArrayRef<int64_t> dilations) {
  auto outputType = cast<RankedTensorType>(packedOutput.getType());
  IREE::Codegen::DataTiledConvIterationSpace space =
      IREE::Codegen::getDataTiledConvIterationSpace(builder.getContext(),
                                                    strides, dilations);
  return linalg::GenericOp::create(
      builder, loc, outputType,
      /*inputs=*/ValueRange{packedInput, packedFilter},
      /*outputs=*/ValueRange{packedOutput}, space.indexingMaps,
      space.iteratorTypes, [](OpBuilder &b, Location l, ValueRange args) {
        Value mul = arith::MulFOp::create(b, l, args[0], args[1]);
        Value add = arith::AddFOp::create(b, l, mul, args[2]);
        linalg::YieldOp::create(b, l, add);
      });
}

/// Lowers a data-tiled convolution whose operands carry conv encodings.
///
/// Builds a 9-D linalg.generic over the packed operands:
///   input:  [N, IC/c0, H, W, c0]
///   filter: [OC/k0, IC/c0, FH, FW, c0, k0]
///   output: [N, OC/k0, OH, OW, k0]
Operation *lowerConvolutionOpWithEncoding(
    OpBuilder &builder, linalg::LinalgOp linalgOp, ValueRange operands,
    IREE::Encoding::LayoutMaterializerAttr layoutAttr) {
  if (!linalgOp.hasPureTensorSemantics()) {
    return nullptr;
  }

  auto inputs = linalgOp.getDpsInputOperands();
  auto outputs = linalgOp.getDpsInits();

  auto encodedInputType = cast<RankedTensorType>(inputs[0]->get().getType());
  auto encodedFilterType = cast<RankedTensorType>(inputs[1]->get().getType());
  auto encodedOutputType = cast<RankedTensorType>(outputs[0].getType());
  auto inEncoding = IREE::Encoding::getEncodingAttr(encodedInputType);
  auto filterEncoding = IREE::Encoding::getEncodingAttr(encodedFilterType);
  auto outEncoding = IREE::Encoding::getEncodingAttr(encodedOutputType);
  if (!inEncoding || !filterEncoding || !outEncoding) {
    return nullptr;
  }

  if (inEncoding.getOperandIndex().getValue() != IREE::Encoding::CONV_IN ||
      filterEncoding.getOperandIndex().getValue() !=
          IREE::Encoding::CONV_FILTER ||
      outEncoding.getOperandIndex().getValue() != IREE::Encoding::CONV_OUT) {
    return nullptr;
  }

  MaterializeEncodingInfo encodingInfo = {};
  if (auto packedLayoutAttr =
          dyn_cast<IREE::Codegen::PackedLayoutMaterializerAttr>(layoutAttr)) {
    encodingInfo = packedLayoutAttr.getEncodingInfo(
        cast<RankedTensorType>(linalgOp->getResultTypes()[0]));
  }

  if (isIdentityLayout(encodingInfo)) {
    return dropEncodingAndCloneOp(builder, linalgOp,
                                  operands.take_front(inputs.size()),
                                  operands.drop_front(inputs.size()));
  }

  Value packedInput = operands[0];
  Value packedFilter = operands[1];
  Value packedOutput = operands[2];

  Location loc = linalgOp.getLoc();

  auto packedInputType = cast<RankedTensorType>(packedInput.getType());
  auto packedOutputType = cast<RankedTensorType>(packedOutput.getType());

  // Expected packed ranks for a 2D conv: input/output rank-5 ([N, IC/c0, H, W,
  // c0] and [N, OC/k0, OH, OW, k0]), or rank-4 if the N=1 batch dim has been
  // folded away upstream, and filter rank-6 ([OC/k0, IC/c0, FH, FW, c0, k0]).
  const int64_t inRank = packedInputType.getRank();
  const int64_t outRank = packedOutputType.getRank();
  const int64_t filterRank =
      cast<RankedTensorType>(packedFilter.getType()).getRank();
  if ((inRank != 4 && inRank != 5) || inRank != outRank || filterRank != 6) {
    return nullptr;
  }

  // WORKAROUND: When N=1, upstream folding may remove the batch dimension.
  // Restore it so the remainder of this lowering can assume rank-5 packed
  // tensors. If batch folding before materialization is eventually prevented,
  // this expand/collapse workaround can be dropped.
  auto prependUnitBatchDim = [&](Value src) -> Value {
    auto srcType = cast<RankedTensorType>(src.getType());

    SmallVector<int64_t, 5> expandedShape = {1};
    expandedShape.append(srcType.getShape().begin(), srcType.getShape().end());

    SmallVector<ReassociationIndices> ri = {{0, 1}, {2}, {3}, {4}};

    return tensor::ExpandShapeOp::create(
        builder, loc,
        RankedTensorType::get(expandedShape, srcType.getElementType()), src,
        ri);
  };

  bool batchDimFolded = (inRank == 4);
  if (batchDimFolded) {
    // [IC/c0, H, W, c0] -> [1, IC/c0, H, W, c0]
    packedInput = prependUnitBatchDim(packedInput);
    // [OC/k0, OH, OW, k0] -> [1, OC/k0, OH, OW, k0]
    packedOutput = prependUnitBatchDim(packedOutput);
  }

  auto cDims = linalg::inferConvolutionDims(linalgOp);
  if (failed(cDims) || cDims->strides.size() != 2 ||
      cDims->dilations.size() != 2) {
    return nullptr;
  }

  Operation *convOp =
      buildDataTiledConvOp(builder, loc, packedInput, packedFilter,
                           packedOutput, cDims->strides, cDims->dilations);

  // [1, OC/k0, OH, OW, k0] -> [OC/k0, OH, OW, k0]
  if (batchDimFolded) {
    SmallVector<ReassociationIndices> ri = {{0, 1}, {2}, {3}, {4}};
    return tensor::CollapseShapeOp::create(builder, loc, packedOutputType,
                                           convOp->getResult(0), ri);
  }
  return convOp;
}

//===----------------------------------------------------------------------===//
// Interface methods implementation for iree_cpu.cpu_encoding_resolver.
//===----------------------------------------------------------------------===//

// Enumerate tile sizes to choose from on riscv32.
// For narrow-{M,N} cases, this only enumerates on narrow M. The narrow-N cases
// are handled by transposition in chooseMatmulTile.
static SmallVector<TileMxNxK>
enumerateMatmulTileRiscv32(DictionaryAttr config) {
  if (hasUkernel(config)) {
    return {
        TileMxNxK{8, 8, 4}, // Some reasonable tile shape.
        TileMxNxK{4, 8, 4}, // Truncation of the above.
        TileMxNxK{2, 8, 4}, // Truncation of the above.
        TileMxNxK{1, 8, 4}, // Truncation of the above.
    };
  }
  // Fallback - no architecture-optimized tile size for this case.
  return {};
}
// RISC-V has vector register length extensions: zvl128b, zvl256b etc.
// If these extension are specified in target cpu feature,
// they can be used to determine VLEN. This function assumes that
// 'v' feature is present
size_t getRISCVVVlenFromCPUFeatures(DictionaryAttr config) {
  // If +zvl* feature is not explicitly specified,
  // fallback to +zvl128b, as spec specifies minimum VLEN
  // of 128b for the V extension: https://rb.gy/p8rbzv
  size_t vlen;
  if (hasFeature(config, "+zvl65536b")) {
    vlen = 65536;
  } else if (hasFeature(config, "+zvl32768b")) {
    vlen = 32768;
  } else if (hasFeature(config, "+zvl16384b")) {
    vlen = 16384;
  } else if (hasFeature(config, "+zvl8192b")) {
    vlen = 8192;
  } else if (hasFeature(config, "+zvl4096b")) {
    vlen = 4096;
  } else if (hasFeature(config, "+zvl2048b")) {
    vlen = 2048;
  } else if (hasFeature(config, "+zvl1024b")) {
    vlen = 1024;
  } else if (hasFeature(config, "+zvl512b")) {
    vlen = 512;
  } else if (hasFeature(config, "+zvl256b")) {
    vlen = 256;
  } else {
    vlen = 128;
  }
  return vlen;
}
// Enumerate tile sizes to choose from on riscv64.
// For narrow-{M,N} cases, this only enumerates on narrow M. The narrow-N cases
// are handled by transposition in chooseMatmulTile.
//
// When using static VLEN-based tile sizes, the N dimension is computed from the
// VLEN using zvl* features.
//
// When using scalable vectors
// (iree-llvmcpu-enable-scalable-vectorization=true), we use smaller base tile
// sizes that scale with vscale (VLEN/64). The base tile sizes assume VLEN=64
// bits, meaning vscale=1 at the minimum. When scalable vectorization is
// enabled, these base sizes will be multiplied by vscale at runtime to utilize
// the full vector register width.
//
// Base tile calculation for VLEN=64:
//   - f32 (32 bits): 64/32 = 2 elements per vscale
//   - f16 (16 bits): 64/16 = 4 elements per vscale
static SmallVector<TileMxNxK>
enumerateMatmulTileRiscv64(TypeRange elementTypes, DictionaryAttr config) {
  // Data-tiling is only implemented for the V extension.
  if (!hasFeature(config, "+v")) {
    return {};
  }

  assert(elementTypes.size() == 3);
  Type lhs = elementTypes[0];
  Type rhs = elementTypes[1];
  Type out = elementTypes[2];

  // RVV has both static and scalable vector modes.
  // Static mode: use VLEN-based tile sizes from zvl* features.
  // Scalable mode: use base VLEN=64 bits as the base size, the vscale
  // multiplier and the correct scalable data-tiled layout will be emitted
  // accordingly in `getScalableTileFlags`.
  constexpr size_t kBaseVlen = 64;
  size_t vlen = isScalableVectorizationEnabled()
                    ? kBaseVlen
                    : getRISCVVVlenFromCPUFeatures(config);

  if (lhs.isF32() && rhs.isF32() && out.isF32()) {
    // VLEN-aware Tile size selection
    // One concern that needs to be addressed here is that
    // for larger VLENs tile sizes would be very large
    // leading to a very high padding overhead
    int N0 = vlen / 8;
    return {
        TileMxNxK{7, N0, 1}, // Aim to use vfmacc, 100% register utilization.
        TileMxNxK{4, N0, 1}, // Truncation of the above.
        TileMxNxK{2, N0, 1}, // Truncation of the above.
        TileMxNxK{1, N0, 1}, // Truncation of the above.
    };
  }
  if (lhs.isF16() && rhs.isF16()) {
    // TODO(egebeysel): currently, the F16 path uses LMUL=2 and therefore 14
    // registers for accumulators. This has to be investigated and the
    // corresponding microkernels have to be adjusted if need be.
    int N0 = vlen / 8;
    if (hasFeature(config, "+zvfh")) {
      return {
          TileMxNxK{7, N0, 1}, // Aim to use vfmacc.
          TileMxNxK{4, N0, 1}, // Truncation of the above.
          TileMxNxK{2, N0, 1}, // Truncation of the above.
          TileMxNxK{1, N0, 1}, // Truncation of the above.
      };
    }
    if (hasFeature(config, "+zvfhmin")) {
      return {
          TileMxNxK{6, N0, 1}, // Aim to use vfmacc.
          TileMxNxK{4, N0, 1}, // Truncation of the above.
          TileMxNxK{2, N0, 1}, // Truncation of the above.
          TileMxNxK{1, N0, 1}, // Truncation of the above.
      };
    }
  }
  // BF16 path: RVV bf16 widening multiply-accumulate via Zvfbfwma.
  // Same formula as the f32 path: N0 = VLEN/8, so the bf16 operands are LMUL=2
  // (m2) and the f32 accumulators EMUL=4 (m4). vfwmaccbf16 fuses the bf16 ->
  // f32 widening into the MAC, so the accumulator is f32 and there is no
  // bf16-accumulator variant.
  // M0=7 with EMUL=4 accumulators: 7*4 + 2(rhs) = 30 regs.
  if (lhs.isBF16() && rhs.isBF16() && out.isF32()) {
    int N0 = vlen / 8;
    if (hasFeature(config, "+zvfbfwma")) {
      return {
          TileMxNxK{7, N0, 1}, // Aim to use VFWMACCBF16.VF.
          TileMxNxK{4, N0, 1}, // Truncation of the above.
          TileMxNxK{2, N0, 1}, // Truncation of the above.
          TileMxNxK{1, N0, 1}, // Truncation of the above.
      };
    }
  }
  // Integer 8 path.
  if (lhs.isSignlessInteger(8) && rhs.isSignlessInteger(8) &&
      out.isSignlessInteger(32)) {
    int N0 = vlen / 8;
    // SpaceMiT IME (vmadot): fixed hardware tiles for VLEN ∈ {256, 1024,
    // 4096}. Tile is a 3×4 grid of the SEW=8 MAC atom (4×4×8 / 8×8×16 /
    // 16×16×32). Other VLENs / scalable mode fall through to standard RVV
    // tiles.
    if (hasFeature(config, "+xsmtvdot") && !isScalableVectorizationEnabled()) {
      int atom = 0;
      if (vlen == 256) {
        atom = 4;
      } else if (vlen == 1024) {
        atom = 8;
      } else if (vlen == 4096) {
        atom = 16;
      }
      if (atom != 0) {
        return {TileMxNxK{3 * atom, 4 * atom, 2 * atom}};
      }
    }
    // Standard RVV widening multiply-accumulate.
    // Widening chain: i8(m1) -> i16(m2) -> i32(m4).
    // M0=7 with LMUL=4 accumulators: 7*4 + 2(prod) + 1(rhs) = 31 regs.
    return {
        TileMxNxK{7, N0, 1}, // Primary shape for RVV widening int8 kernels.
        TileMxNxK{4, N0, 1}, // Truncation of the above.
        TileMxNxK{2, N0, 1}, // Truncation of the above.
        TileMxNxK{1, N0, 1}, // Truncation of the above.
    };
  }
  // Fallback - no architecture-optimized tile size for this case.
  return {};
}

// Enumerate tile sizes to choose from on arm64.
// For narrow-{M,N} cases, this only enumerates on narrow M. The narrow-N cases
// are handled by transposition in chooseMatmulTile.
static SmallVector<TileMxNxK> enumerateMatmulTileArm64(TypeRange elementTypes,
                                                       DictionaryAttr config) {
  // For SVE and scalable vectors, this methods selects base sizes that match
  // the NEON fixed-width sizes.
  // TODO: Add SME inner tile sizes and corresponding tests.
  assert(elementTypes.size() == 3);
  Type lhs = elementTypes[0];
  Type rhs = elementTypes[1];
  Type out = elementTypes[2];

  if (out.isF32() || out.isF16() || out.isBF16()) {
    if (lhs.isBF16() && rhs.isBF16() && (out.isBF16() || out.isF32()) &&
        hasFeature(config, "+bf16")) {
      return {
          TileMxNxK{8, 8, 4}, // Aim to use BFMMLA.
          TileMxNxK{4, 8, 4}, // Truncation of the above.
          TileMxNxK{2, 8, 4}, // Truncation of the above.
          TileMxNxK{1, 8, 4}, // Truncation of the above.
      };
    }
    if (isa<FloatType>(lhs) && isa<FloatType>(rhs)) {
      // Note: 16-bit floating point types currently use the same tile size as
      // f32. This makes sense when either (1) the accumulator is f32, or (2)
      // the arithmetic will have to expand f16 to f32 in registers. We may
      // reconsider when taking advantage of native f16/bf16 arithmetic when the
      // accumulator itself is f16/bf16, as we could typically have a 2x wider
      // tile in that case. However, on current CPUs, the existing tiles seem
      // wide enough already to approach peak performance.
      return {
          TileMxNxK{8, 8, 1}, // Aim to use FMLA or FMLAL.
          TileMxNxK{4, 8, 1}, // Truncation of the above.
          TileMxNxK{2, 8, 1}, // Truncation of the above.
          TileMxNxK{1, 8, 1}, // Truncation of the above.
      };
    }
  }
  if (lhs.isSignlessInteger(8) && rhs.isSignlessInteger(8) &&
      out.isSignlessInteger(32)) {
    if (hasFeature(config, "+i8mm")) {
      return {
          TileMxNxK{8, 8, 8}, // Aim to use SMMLA.
          TileMxNxK{4, 8, 8}, // Truncation of the above.
          TileMxNxK{2, 8, 8}, // Truncation of the above.
          TileMxNxK{1, 8, 8}, // Truncation of the above.
      };
    }
    if (hasFeature(config, "+dotprod")) {
      return {
          TileMxNxK{8, 8, 4}, // Aim to use SDOT.
          TileMxNxK{4, 8, 4}, // Truncation of the above.
          TileMxNxK{2, 8, 4}, // Truncation of the above.
          TileMxNxK{1, 8, 4}, // Truncation of the above.
      };
    }
  }
  if (lhs.isSignlessInteger(8) && rhs.isSignlessInteger(4) &&
      out.isSignlessInteger(32)) {
    if (hasFeature(config, "+i8mm")) {
      return {
          TileMxNxK{4, 8, 16},
          TileMxNxK{2, 8, 16},
          TileMxNxK{1, 8, 16},
      };
    }
    if (hasFeature(config, "+dotprod")) {
      return {
          TileMxNxK{8, 8, 8},
          TileMxNxK{4, 8, 8},
          TileMxNxK{2, 8, 8},
          TileMxNxK{1, 8, 8},
      };
    }
    return {
        TileMxNxK{4, 16, 2},
        TileMxNxK{2, 16, 2},
        TileMxNxK{1, 16, 2},
    };
  }
  // Fallback - no architecture-optimized tile size for this case.
  return {};
}

// Enumerate tile sizes to choose from on x86-64.
// For narrow-{M,N} cases, this only enumerates on narrow M. The narrow-N cases
// are handled by transposition in chooseMatmulTile.
static SmallVector<TileMxNxK> enumerateMatmulTileX86_64(TypeRange elementTypes,
                                                        DictionaryAttr config) {
  assert(elementTypes.size() == 3);
  Type lhs = elementTypes[0];
  Type rhs = elementTypes[1];
  Type out = elementTypes[2];

  if (out.isF32() || out.isF16() || out.isBF16()) {
    if (lhs.isBF16() && rhs.isBF16() && (out.isBF16() || out.isF32())) {
      if (hasFeature(config, "+avx512bf16")) {
        return {
            TileMxNxK{16, 16, 2}, // Aim to use VDPBF16PS (zmm).
            TileMxNxK{8, 16, 2},  // Truncation of the above.
            TileMxNxK{4, 16, 2},  // Truncation of the above.
            TileMxNxK{2, 16, 2},  // Truncation of the above.
            TileMxNxK{1, 16, 2},  // Truncation of the above.
        };
      }
    }
    if (isa<FloatType>(lhs) && isa<FloatType>(rhs)) {
      // Note: 16-bit floating point types currently use the same tile size as
      // f32. This makes sense when either (1) the accumulator is f32, or (2)
      // the arithmetic will have to expand f16 to f32 in registers. We may
      // reconsider when taking advantage of native f16/bf16 arithmetic when the
      // accumulator itself is f16/bf16.
      if (hasFeature(config, "+avx512f")) {
        return {
            TileMxNxK{16, 16, 1}, // Aim to use VFMADD* (zmm).
            TileMxNxK{8, 16, 1},  // Truncation of the above.
            TileMxNxK{4, 16, 1},  // Truncation of the above.
            TileMxNxK{2, 16, 1},  // Truncation of the above.
            TileMxNxK{1, 16, 1},  // Truncation of the above.
        };
      }
      if (hasFeature(config, "+avx")) {
        // Note: for good performance, most +avx users will also want to add
        // +fma, but that's a local instruction selection detail and the tile
        // layout is unaffected, as there are enough registers even with the
        // need for intermediate product registers when +fma is not used.
        return {
            TileMxNxK{8, 8, 1}, // Aim to use VFMADD* (ymm).
            TileMxNxK{4, 8, 1}, // Truncation of the above.
            TileMxNxK{2, 8, 1}, // Truncation of the above.
            TileMxNxK{1, 8, 1}, // Truncation of the above.
        };
      }
      // SSE fallback.
      return {
          TileMxNxK{8, 4, 1}, // Aim to use MULPS/ADDPS (xmm).
          TileMxNxK{4, 4, 1}, // Truncation of the above.
          TileMxNxK{2, 4, 1}, // Truncation of the above.
          TileMxNxK{1, 4, 1}, // Truncation of the above.
      };
    }
  }

  if (out.isSignlessInteger(32) &&
      ((lhs.isSignlessInteger(8) && rhs.isSignlessInteger(8)) ||
       (lhs.isSignlessInteger(16) && rhs.isSignlessInteger(16)))) {
    if (hasFeature(config, "+avx512vnni")) {
      // This is the same tile size as with VPMADDWD as the only difference
      // is that VPDPWSSD accumulates. VPDPBUSD would call for {16, 16, 4} but
      // we can't easily use it because of its unsigned*signed semantics.
      return {
          TileMxNxK{16, 16, 2}, // Aim to use VPDPWSSD (zmm).
          TileMxNxK{8, 16, 2},  // Truncation of the above.
          TileMxNxK{4, 16, 2},  // Truncation of the above.
          TileMxNxK{2, 16, 2},  // Truncation of the above.
          TileMxNxK{1, 16, 2},  // Truncation of the above.
      };
    }
    if (hasFeature(config, "+avx512bw")) {
      return {
          TileMxNxK{16, 16, 2}, // Aim to use VPMADDWD (zmm).
          TileMxNxK{8, 16, 2},  // Truncation of the above.
          TileMxNxK{4, 16, 2},  // Truncation of the above.
          TileMxNxK{2, 16, 2},  // Truncation of the above.
          TileMxNxK{1, 16, 2},  // Truncation of the above.
      };
    }
    if (hasFeature(config, "+avx2")) {
      return {
          TileMxNxK{8, 8, 2}, // Aim to use VPMADDWD (ymm).
          TileMxNxK{4, 8, 2}, // Truncation of the above.
          TileMxNxK{2, 8, 2}, // Truncation of the above.
          TileMxNxK{1, 8, 2}, // Truncation of the above.
      };
    }
    // SSE fallback.
    return {
        TileMxNxK{8, 4, 2}, // Aim to use PMADDWD (xmm).
        TileMxNxK{4, 4, 2}, // Truncation of the above.
        TileMxNxK{2, 4, 2}, // Truncation of the above.
        TileMxNxK{1, 4, 2}, // Truncation of the above.
    };
  }

  if (out.isSignlessInteger(32) && lhs.isSignlessInteger(16) &&
      rhs.isUnsignedInteger(4)) {
    // Experimental s16u4s32 case. Focusing only on the vecmat case for now.
    if (hasFeature(config, "+avx512vnni")) {
      return {
          TileMxNxK{1, 32, 8}, // Aim to use VPDPBUSD (zmm).
      };
    }
  }

  // Fallback - no architecture-optimized tile size for this case.
  return {};
}

static SmallVector<TileMxNxK>
enumerateCPUMatmulTiles(IREE::Encoding::EncodingAttr encoding,
                        DictionaryAttr config) {
  // Enumerate available tile shapes for the given encoding and config.
  SmallVector<Type> elementTypes = encoding.getElementTypesArray();
  if (isAArch64(config)) {
    return enumerateMatmulTileArm64(elementTypes, config);
  }
  if (isX86_64(config)) {
    return enumerateMatmulTileX86_64(elementTypes, config);
  }
  if (isRISCV32(config)) {
    return enumerateMatmulTileRiscv32(config);
  }
  if (isRISCV64(config)) {
    return enumerateMatmulTileRiscv64(elementTypes, config);
  }
  return {};
}

// Enumerate (OC, IC) conv inner tile sizes to choose from on arm64.
static SmallVector<TileOCxIC> enumerateConvTileArm64(TypeRange elementTypes,
                                                     DictionaryAttr config) {
  assert(elementTypes.size() == 3 && "expected input, filter, output types");
  Type inputElem = elementTypes[0];
  Type filterElem = elementTypes[1];
  Type outputElem = elementTypes[2];
  // Only fp32 is currently supported; the tile shape and
  // the ukernel are specific to 32-bit float lanes.
  if (inputElem.isF32() && filterElem.isF32() && outputElem.isF32()) {
    return {TileOCxIC{8, 8}}; // Aim to use FMLA (neon).
  }

  // Fallback - no architecture-optimized tile size for this case.
  return {};
}

// Enumerate (OC, IC) conv inner tile sizes to choose from on x86-64.
static SmallVector<TileOCxIC> enumerateConvTileX86_64(TypeRange elementTypes,
                                                      DictionaryAttr config) {
  assert(elementTypes.size() == 3 && "expected input, filter, output types");
  Type inputElem = elementTypes[0];
  Type filterElem = elementTypes[1];
  Type outputElem = elementTypes[2];
  // Only fp32 is currently supported; the tile shape and
  // the ukernel are specific to 32-bit float lanes.
  if (inputElem.isF32() && filterElem.isF32() && outputElem.isF32()) {
    if (hasFeature(config, "+avx512f")) {
      return {TileOCxIC{16, 16}}; // Aim to use VFMADD* (zmm).
    }
  }

  // Fallback - no architecture-optimized tile size for this case.
  return {};
}

static SmallVector<TileOCxIC>
enumerateCPUConvTiles(IREE::Encoding::EncodingAttr encoding,
                      DictionaryAttr config) {
  SmallVector<Type> elementTypes = encoding.getElementTypesArray();
  if (isAArch64(config)) {
    return enumerateConvTileArm64(elementTypes, config);
  }
  if (isX86_64(config)) {
    return enumerateConvTileX86_64(elementTypes, config);
  }

  // Fallback - no architecture-optimized tile size for this case.
  return {};
}

struct CPUEncodingPackedLayoutMaterializerAttr
    : PackedLayoutMaterializerAttrExternalModelBase<
          CPUEncodingPackedLayoutMaterializerAttr, CPUEncodingResolverAttr> {

  DictionaryAttr getConfiguration(Attribute attr) const {
    return cast<CPUEncodingResolverAttr>(attr).getConfiguration();
  }

  MaterializeEncodingInfo getEncodingInfoImpl(Attribute attr,
                                              RankedTensorType type) const {
    auto layoutAttr = cast<CPUEncodingResolverAttr>(attr);

    auto encoding =
        dyn_cast_if_present<IREE::Encoding::EncodingAttr>(type.getEncoding());

    MaterializeEncodingInfo info;
    if (!encoding) {
      return info;
    }

    // Handle convolution encodings.
    if (encoding.getOpType().getValue() ==
        IREE::Encoding::EncodingOpType::conv) {
      FailureOr<linalg::ConvolutionDimensions> cDims =
          linalg::inferConvolutionDims(encoding.getRootMaps());
      if (failed(cDims) || !cDims->depth.empty() ||
          cDims->outputImage.size() != 2 ||
          llvm::any_of(cDims->dilations, [](int64_t d) { return d != 1; })) {
        return info; // identity layout
      }
      SmallVector<TileOCxIC> tiles =
          enumerateCPUConvTiles(encoding, layoutAttr.getConfiguration());
      if (tiles.empty()) {
        return info;
      }
      TileOCxIC tile = tiles[0];
      FailureOr<MaterializeEncodingInfo> maybeInfo =
          IREE::Codegen::getEncodingInfoForConv(encoding, tile);
      if (failed(maybeInfo)) {
        return info;
      }
      return maybeInfo.value();
    }

    // We only know about contractions with {Batch, M, N, K} <= 1 at the moment.
    auto cDims = getEncodingContractionDims(encoding);
    if (failed(cDims) || cDims->batch.size() > 1 || cDims->m.size() > 1 ||
        cDims->n.size() > 1 || cDims->k.size() > 1) {
      return info;
    }

    DictionaryAttr config = layoutAttr.getConfiguration();
    if (getEnableInnerTiledFromConfig(config)) {
      return getInnerTiledEncodingInfo(type.getContext(), encoding, *cDims,
                                       config);
    }
    return getMmt4dEncodingInfo(encoding, *cDims, config);
  }

private:
  /// Legacy mmt4d path: enumerate candidate `TileMxNxK`s, pick one with the
  /// narrow-dim-aware scoring, and apply the narrow-N→narrow-M transpose
  /// trick required by the handwritten microkernels.
  static MaterializeEncodingInfo
  getMmt4dEncodingInfo(IREE::Encoding::EncodingAttr encoding,
                       const linalg::ContractionDimensions &cDims,
                       DictionaryAttr config) {
    MaterializeEncodingInfo info;
    SmallVector<TileMxNxK> enumeratedTileMxNxK =
        enumerateCPUMatmulTiles(encoding, config);
    if (enumeratedTileMxNxK.empty()) {
      return info;
    }
    auto narrowDim = IREE::Encoding::getPo2MatmulNarrowDim(encoding);
    // In case of scalable vectors, we explicitly use the narrow M dimension
    // config to use the M dimension for calculating padding penalties while we
    // choose the layout.
    // TODO(egebeysel): reconcile this with the static case. Currently,
    // this results in vector types with non-trailing scalable dimensions, e.g.
    // vector<[8]x8xf32>, which cannot be legalized or lowered to LLVM properly.
    // LLVM only handles vector types with only-trailing scalable dimensions
    // properly, e.g. vector<8x[8]xf32> is lowered to !llvm.array<8 x
    // vector<[8]xf32>>. The former have to be legalized, so that they can be
    // lowered to an array of scalable vectors, but the patterns are missing.
    if (isScalableVectorizationEnabled() && narrowDim.isN() &&
        !cDims.m.empty()) {
      int64_t m = getEncodingContractionLikeSizes(encoding).value().M;
      if (ShapedType::isDynamic(m)) {
        narrowDim = {};
      } else {
        narrowDim = {IREE::Encoding::MatmulNarrowDim::Dim::M, m};
      }
    }
    // Choose a final matmul TileMxNxK from the above-enumerated tile shapes,
    // taking narrow dimensions into account.
    TileMxNxK chosenTileMxNxK =
        chooseMatmulTile(enumeratedTileMxNxK, narrowDim);
    FailureOr<MaterializeEncodingInfo> maybeEncodingInfo =
        getEncodingInfoForMatmul(encoding, chosenTileMxNxK);
    if (failed(maybeEncodingInfo)) {
      return info;
    }
    info = std::move(maybeEncodingInfo.value());
    FailureOr<IREE::Codegen::ScalableTileFlags> scalableFlags =
        getScalableTileFlags(cDims, encoding, config);
    if (succeeded(scalableFlags)) {
      info.scalableTiles = std::move(scalableFlags);
    }
    if (IREE::Encoding::isNarrowNResult(encoding) &&
        llvm::none_of(info.scalableTiles.value_or(Codegen::ScalableTileFlags{}),
                      [](bool flag) { return flag; })) {
      transposeInPlace(info);
    }
    return info;
  }

  /// Intrinsic-first path for `iree_codegen.inner_tiled`: pick a
  /// `DataTiledMMAAttr` (same decision the lowering will make) and derive the
  /// packed-layout tile shape from it. No narrow-N transpose: unlike the
  /// legacy mmt4d path, the inner_tiled lowering handles narrow dims natively
  /// via `intrinsics_m` / `intrinsics_n`.
  static MaterializeEncodingInfo getInnerTiledEncodingInfo(
      MLIRContext *ctx, IREE::Encoding::EncodingAttr encoding,
      const linalg::ContractionDimensions &cDims, DictionaryAttr config) {
    MaterializeEncodingInfo info;
    IREE::CPU::DataTiledMMAAttr mma =
        chooseCpuInnerTiledMmaForEncoding(ctx, encoding, config);
    if (!mma) {
      return info;
    }
    FailureOr<MaterializeEncodingInfo> maybeEncodingInfo =
        getEncodingInfoForMatmul(encoding, getTileMxNxK(mma));
    if (failed(maybeEncodingInfo)) {
      return info;
    }
    info = std::move(maybeEncodingInfo.value());
    FailureOr<IREE::Codegen::ScalableTileFlags> scalableFlags =
        getScalableTileFlags(cDims, encoding, config);
    if (succeeded(scalableFlags)) {
      info.scalableTiles = std::move(scalableFlags);
    }
    // Attach the per-operand swizzle so the pack pipeline's
    // `getSwizzledShape` applies the `expand_shape` + `transpose` to the
    // inner tile, matching the permutation-aware tile type that
    // `DataTiledMMAAttr::getUndistributedTileTypes` returns. Mirrors GPU.
    info.swizzle =
        IREE::CPU::getSwizzle(mma, encoding.getOperandIndex().getInt());
    return info;
  }
};

struct CPUEncodingResolverMaterializerAttr final
    : EncodingLayoutMaterializerAttrExternalModelBase<
          CPUEncodingResolverMaterializerAttr, CPUEncodingResolverAttr> {

  Operation *lowerOp(Attribute attr, OpBuilder &b, Operation *op,
                     TypeRange convertedResTypes,
                     ValueRange convertedOperands) const {
    auto layoutAttr = cast<CPUEncodingResolverAttr>(attr);
    auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
    if (!linalgOp) {
      return nullptr;
    }
    if (auto fillOp = dyn_cast<linalg::FillOp>(op)) {
      return lowerFillOpWithResolvedLayouts(b, fillOp, convertedResTypes,
                                            convertedOperands);
    }
    // Scaled contraction (MX matmul) is not yet supported on CPU, so we drop
    // the encoding and clone the op as-is.
    if (IREE::LinalgExt::isaScaledContractionOpInterface(linalgOp)) {
      int64_t numInputs = linalgOp.getNumDpsInputs();
      return dropEncodingAndCloneOp(b, linalgOp,
                                    convertedOperands.take_front(numInputs),
                                    convertedOperands.drop_front(numInputs));
    }
    if (linalg::isaConvolutionOpInterface(linalgOp)) {
      return lowerConvolutionOpWithEncoding(
          b, linalgOp, convertedOperands,
          cast<IREE::Encoding::LayoutMaterializerAttr>(layoutAttr));
    }
    if (linalg::isaContractionOpInterface(linalgOp)) {
      DictionaryAttr config = layoutAttr.getConfiguration();
      if (getEnableInnerTiledFromConfig(config)) {
        return lowerContractionToInnerTiled(
            b, linalgOp, convertedOperands,
            cast<IREE::Encoding::LayoutMaterializerAttr>(layoutAttr));
      }
      return lowerContractionOpWithEncoding(
          b, linalgOp, convertedOperands,
          cast<IREE::Encoding::LayoutMaterializerAttr>(layoutAttr));
    }
    if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {
      return lowerGenericOpWithResolvedLayouts(
          b, genericOp, convertedResTypes, convertedOperands,
          cast<IREE::Encoding::LayoutMaterializerAttr>(attr));
    }
    return nullptr;
  }
};

struct CPULayoutResolverAttr final
    : IREE::Encoding::LayoutResolverAttr::ExternalModel<
          CPULayoutResolverAttr, CPUEncodingResolverAttr> {
  Attribute cloneWithSimplifiedConfig(Attribute attr,
                                      DictionaryAttr config) const {
    MLIRContext *ctx = attr.getContext();
    SmallVector<NamedAttribute> configItems;
    if (std::optional<StringRef> cpuFeatures = getConfigCpuFeatures(config)) {
      addConfigCpuFeatures(ctx, cpuFeatures.value(), configItems);
    }
    if (std::optional<StringRef> targetTriple = getConfigTargetTriple(config)) {
      addConfigTargetTriple(ctx, targetTriple.value(), configItems);
    }
    storeNamedAttrIfPresent(configItems, config, "ukernels");
    storeNamedAttrIfPresent(configItems, config, "enable_inner_tiled");
    return CPUEncodingResolverAttr::get(ctx,
                                        DictionaryAttr::get(ctx, configItems));
  }

  Attribute getLayout(Attribute attr, RankedTensorType type) const {
    MLIRContext *ctx = attr.getContext();
    return CPUEncodingResolverAttr::get(ctx, getPackedLayoutImpl(attr, type));
  }
};

struct CPUSerializableAttr final
    : IREE::Encoding::SerializableAttr::ExternalModel<CPUSerializableAttr,
                                                      CPUEncodingResolverAttr> {
  bool isSerialized(Attribute attr) const {
    auto configuration = cast<CPUEncodingResolverAttr>(attr).getConfiguration();
    return configuration && configuration.contains(kEncodingInfoAttrName);
  }

  Value calculateStorageSizeInBytes(Attribute attr, Location loc,
                                    OpBuilder &builder, RankedTensorType type,
                                    ValueRange dynamicDims) const {
    return calculatePackedStorageSizeInBytesImpl(attr, loc, builder, type,
                                                 dynamicDims);
  }
};

struct CPUEncodingResolverVerifier
    : mlir::VerifiableTensorEncoding::ExternalModel<CPUEncodingResolverVerifier,
                                                    CPUEncodingResolverAttr> {

  LogicalResult
  verifyEncoding(Attribute attr, ArrayRef<int64_t> shape, Type elementType,
                 function_ref<InFlightDiagnostic()> emitError) const {
    auto packedLayoutMaterializerAttr =
        cast<Codegen::PackedLayoutMaterializerAttr>(attr);
    return packedLayoutMaterializerAttr.verifyPackedLayoutWithType(
        shape, elementType, emitError);
  }
};

//===----------------------------------------------------------------------===//
// Interface methods implementation for iree_cpu.vmvx_encoding_resolver.
//===----------------------------------------------------------------------===//

// Enumerate tile sizes to choose from when no specific architecture is
// targeted. For narrow-{M,N} cases, this only enumerates on narrow M. The
// narrow-N cases are handled by transposition in chooseMatmulTile.
static SmallVector<TileMxNxK>
enumerateVMVXMatmulTiles(linalg::ContractionDimensions cDims,
                         IREE::Encoding::EncodingAttr encoding,
                         DictionaryAttr config) {
  bool hasUkernelSupport = hasUkernel(config);

  // TODO(hanchung): The ukernel path does not support 3d
  // codegen.query_tile_sizes op, so we disable dynamic tile shapes for
  // batch_matmul. Also, they are not set up for narrow M/N matmul, so it is
  // disabled when it is the case.
  if (!cDims.batch.empty() || getPo2MatmulNarrowDim(encoding)) {
    hasUkernelSupport = false;
  }
  if (hasUkernelSupport) {
    // VMVX+ukernel uses dynamic tile shapes.
    return {TileMxNxK{ShapedType::kDynamic, ShapedType::kDynamic,
                      ShapedType::kDynamic}};
  }

  return {
      TileMxNxK{8, 8, 4}, // Some vaguely reasonable tile shape.
      TileMxNxK{4, 8, 4}, // Truncation of the above.
      TileMxNxK{2, 8, 4}, // Truncation of the above.
      TileMxNxK{1, 8, 4}, // Truncation of the above.
  };
}

struct VMVXEncodingPackedLayoutMaterializerAttr final
    : PackedLayoutMaterializerAttrExternalModelBase<
          VMVXEncodingPackedLayoutMaterializerAttr, VMVXEncodingResolverAttr> {

  DictionaryAttr getConfiguration(Attribute attr) const {
    return cast<VMVXEncodingResolverAttr>(attr).getConfiguration();
  }

  MaterializeEncodingInfo getEncodingInfoImpl(Attribute attr,
                                              RankedTensorType type) const {
    auto layoutAttr = cast<VMVXEncodingResolverAttr>(attr);

    auto encoding =
        dyn_cast_if_present<IREE::Encoding::EncodingAttr>(type.getEncoding());

    MaterializeEncodingInfo info;
    if (!encoding) {
      return info;
    }

    // We only know about contractions with {Batch, M, N, K} <= 1 at the moment.
    auto cDims = getEncodingContractionDims(encoding);
    if (failed(cDims) || cDims->batch.size() > 1 || cDims->m.size() > 1 ||
        cDims->n.size() > 1 || cDims->k.size() > 1) {
      return info;
    }

    SmallVector<TileMxNxK> enumeratedTileMxNxK = enumerateVMVXMatmulTiles(
        cDims.value(), encoding, layoutAttr.getConfiguration());
    if (enumeratedTileMxNxK.empty()) {
      return info;
    }
    auto narrowDim = IREE::Encoding::getPo2MatmulNarrowDim(encoding);
    // Choose a final matmul TileMxNxK from the above-enumerated tile shapes,
    // taking narrow dimensions into account.
    TileMxNxK chosenTileMxNxK =
        chooseMatmulTile(enumeratedTileMxNxK, narrowDim);
    FailureOr<MaterializeEncodingInfo> maybeEncodingInfo =
        getEncodingInfoForMatmul(encoding, chosenTileMxNxK);
    if (failed(maybeEncodingInfo)) {
      return info;
    }
    info = std::move(maybeEncodingInfo.value());
    if (IREE::Encoding::isNarrowNResult(encoding)) {
      transposeInPlace(info);
    }
    return info;
  }
};

struct VMVXEncodingResolverMaterializerAttr final
    : EncodingLayoutMaterializerAttrExternalModelBase<
          VMVXEncodingResolverMaterializerAttr, VMVXEncodingResolverAttr> {

  Operation *lowerOp(Attribute attr, OpBuilder &b, Operation *op,
                     TypeRange convertedResTypes,
                     ValueRange convertedOperands) const {
    auto layoutAttr = cast<VMVXEncodingResolverAttr>(attr);
    auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
    if (!linalgOp) {
      return nullptr;
    }
    if (auto fillOp = dyn_cast<linalg::FillOp>(op)) {
      return lowerFillOpWithResolvedLayouts(b, fillOp, convertedResTypes,
                                            convertedOperands);
    }
    if (linalg::isaContractionOpInterface(linalgOp)) {
      return lowerContractionOpWithEncoding(
          b, linalgOp, convertedOperands,
          cast<IREE::Encoding::LayoutMaterializerAttr>(layoutAttr));
    }
    if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {
      return lowerGenericOpWithResolvedLayouts(
          b, genericOp, convertedResTypes, convertedOperands,
          cast<IREE::Encoding::LayoutMaterializerAttr>(attr));
    }
    return nullptr;
  }
};

struct VMVXLayoutResolverAttr final
    : IREE::Encoding::LayoutResolverAttr::ExternalModel<
          VMVXLayoutResolverAttr, VMVXEncodingResolverAttr> {
  Attribute cloneWithSimplifiedConfig(Attribute attr,
                                      DictionaryAttr config) const {
    MLIRContext *ctx = attr.getContext();
    SmallVector<NamedAttribute> configItems;
    storeNamedAttrIfPresent(configItems, config, "ukernels");
    return VMVXEncodingResolverAttr::get(ctx,
                                         DictionaryAttr::get(ctx, configItems));
  }

  Attribute getLayout(Attribute attr, RankedTensorType type) const {
    MLIRContext *ctx = attr.getContext();
    return VMVXEncodingResolverAttr::get(
        ctx, getPackedLayoutImpl(attr, type, /*addEncodingAttr=*/true));
  }
};

struct VMVXSerializableAttr final
    : IREE::Encoding::SerializableAttr::ExternalModel<
          VMVXSerializableAttr, VMVXEncodingResolverAttr> {
  bool isSerialized(Attribute attr) const {
    auto configuration =
        cast<VMVXEncodingResolverAttr>(attr).getConfiguration();
    return configuration && configuration.contains(kEncodingInfoAttrName);
  }

  Value calculateStorageSizeInBytes(Attribute attr, Location loc,
                                    OpBuilder &builder, RankedTensorType type,
                                    ValueRange dynamicDims) const {
    return calculatePackedStorageSizeInBytesImpl(attr, loc, builder, type,
                                                 dynamicDims);
  }
};

struct VMVXEncodingResolverVerifier
    : mlir::VerifiableTensorEncoding::ExternalModel<
          VMVXEncodingResolverVerifier, VMVXEncodingResolverAttr> {
  LogicalResult
  verifyEncoding(Attribute attr, ArrayRef<int64_t> shape, Type elementType,
                 function_ref<InFlightDiagnostic()> emitError) const {
    auto packedLayoutMaterializerAttr =
        cast<Codegen::PackedLayoutMaterializerAttr>(attr);
    return packedLayoutMaterializerAttr.verifyPackedLayoutWithType(
        shape, elementType, emitError);
  }
};

} // namespace

void registerCPUEncodingExternalModels(DialectRegistry &registry) {
  registry.addExtension(
      +[](MLIRContext *ctx, IREE::CPU::IREECPUDialect *dialect) {
        IREE::CPU::CPUEncodingResolverAttr::attachInterface<
            CPUEncodingPackedLayoutMaterializerAttr,
            CPUEncodingResolverMaterializerAttr, CPULayoutResolverAttr,
            CPUSerializableAttr, CPUEncodingResolverVerifier>(*ctx);
        IREE::CPU::VMVXEncodingResolverAttr::attachInterface<
            VMVXEncodingPackedLayoutMaterializerAttr,
            VMVXEncodingResolverMaterializerAttr, VMVXLayoutResolverAttr,
            VMVXSerializableAttr, VMVXEncodingResolverVerifier>(*ctx);
      });
}

} // namespace mlir::iree_compiler::IREE::CPU
