// Copyright 2020 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/LLVMCPU/KernelDispatch.h"

#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "iree/compiler/Codegen/Transforms/Transforms.h"
#include "iree/compiler/Codegen/Utils/MarkerUtils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/TargetSelect.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

namespace mlir {
namespace iree_compiler {

/// NOTE: None of these flags are supported in any form long term. This are
/// temporary hooks added for development purposes. They could be
/// changed/modified at any time.
/// TODO: Find a way to plumb this through to not rely on these flags.

static llvm::cl::opt<int> clNativeVectorSizeInBytes(
    "iree-codegen-llvm-vector-size-in-bytes",
    llvm::cl::desc("native vector size to use on the hardware"),
    llvm::cl::init(16));

static llvm::cl::opt<int> clNumberOfRuntimeThreads(
    "iree-codegen-llvm-number-of-threads",
    llvm::cl::desc("number of threads that are used at runtime"),
    llvm::cl::init(8));

static llvm::cl::list<int> mmt4dWorkgroupTileSizes(
    "iree-codegen-llvm-mmt4d-workgroup-tile-sizes",
    llvm::cl::desc("linalg.mmt4d workgroup tile size"), llvm::cl::ZeroOrMore);

static llvm::cl::list<int> mmt4dL1TileSizes(
    "iree-codegen-llvm-mmt4d-l1-tile-size",
    llvm::cl::desc("linalg.mmt4d L1 tile size"), llvm::cl::ZeroOrMore);

static llvm::cl::list<int> mmt4dVectorSizes(
    "iree-codegen-llvm-mmt4d-vector-size",
    llvm::cl::desc("linalg.mmt4d vector tile size"), llvm::cl::ZeroOrMore);

static llvm::cl::opt<int> defaultWorkgroupTileSize(
    "iree-codegen-llvm-generic-ops-workgroup-size",
    llvm::cl::desc(
        "linalg.generic and linalg.indexed_generic workgroup tile size"),
    llvm::cl::init(64));

static llvm::cl::opt<bool> useLinalgTransformInterp(
    "iree-codegen-use-linalg-transform-interp",
    llvm::cl::desc(
        "experimental path to use the linalg transform dialect interpreter"),
    llvm::cl::init(false));

using IREE::Codegen::DispatchLoweringPassPipeline;

/// Looks for the `native_vector_size` attribute in the hal.executable.variant
/// op.
static Optional<int64_t> getNativeVectorSizeInBytes(func::FuncOp entryPointFn) {
  auto variantOp =
      entryPointFn->getParentOfType<IREE::HAL::ExecutableVariantOp>();
  if (!variantOp) return llvm::None;
  IREE::HAL::ExecutableTargetAttr targetAttr = variantOp.target();
  if (!targetAttr) return llvm::None;
  auto config = targetAttr.getConfiguration();
  if (!config) return llvm::None;
  auto nativeVectorSizeAttr = config.getAs<IntegerAttr>("native_vector_size");
  if (!nativeVectorSizeAttr) return llvm::None;
  int64_t nativeVectorSizeVal = nativeVectorSizeAttr.getInt();
  if (!nativeVectorSizeVal) return llvm::None;
  return nativeVectorSizeVal;
}

/// For a given `shapedType` or (`byteWidth` of element type) return the number
/// of elements that correspond to the native vector size. Returns 1 as the
/// fallback.
static int64_t getVectorSize(func::FuncOp entryPointFn, unsigned byteWidth) {
  if (Optional<int64_t> nativeVectorSize =
          getNativeVectorSizeInBytes(entryPointFn)) {
    return nativeVectorSize.getValue() / byteWidth;
  }
  return clNativeVectorSizeInBytes / byteWidth;
}
static int64_t getVectorSize(func::FuncOp entryPointFn, ShapedType shapedType) {
  Type elementType = shapedType.getElementType();
  if (!elementType.isIntOrFloat()) return 1;
  unsigned byteWidth = IREE::Util::getRoundedElementByteWidth(elementType);
  return getVectorSize(entryPointFn, byteWidth);
}

/// Returns minimum tiling sizes for each dimension. One dimension is possible
/// to access at different element types. It determines the tiling sizes by
/// looking into all the operands.
static SmallVector<int64_t> getMinTilingSizesForEachDim(
    func::FuncOp entryPointFn, linalg::LinalgOp op) {
  unsigned numLoops = op.getNumLoops();
  SmallVector<int64_t> minTileSizes(numLoops, 1);
  auto inputOutputOpOperands = op.getInputAndOutputOperands();
  for (auto map : llvm::enumerate(op.getIndexingMaps())) {
    // Check the fastest varying dimension of the operand. Set the vector size
    // of the corresponding loop to the vector size.
    if (map.value().getNumResults() == 0) continue;
    auto fastestVaryingDimExpr =
        map.value().getResults().back().dyn_cast<AffineDimExpr>();
    if (!fastestVaryingDimExpr) continue;
    unsigned fastestVaryingDim = fastestVaryingDimExpr.getPosition();

    // If the indexing map has result it has to be a shaped type.
    auto operandType =
        inputOutputOpOperands[map.index()]->get().getType().cast<ShapedType>();
    minTileSizes[fastestVaryingDim] =
        std::max<int64_t>(minTileSizes[fastestVaryingDim],
                          getVectorSize(entryPointFn, operandType));
  }
  return minTileSizes;
}

/// Returns the type length in bytes. Looks through all the interface binding
/// ops to see the ABI types and guess-timates the type size to use. This is
/// used to convert the vector size in bytes to vector size in number of
/// elements.
static unsigned getReferenceTypeLengthInBytes(func::FuncOp entryPointFn) {
  unsigned referenceTypeLengthInBytes = 4;
  entryPointFn.walk([&](IREE::HAL::InterfaceBindingSubspanOp subSpanOp) {
    Type type = subSpanOp.getResult().getType();
    Type elementType = TypeSwitch<Type, Type>(type)
                           .Case<ShapedType, IREE::Flow::DispatchTensorType>(
                               [&](auto shapedType) -> Type {
                                 // Ignore operands that are 0D tensors. These
                                 // are not vector-loadable, so using these to
                                 // get vector length would be a pessimization.
                                 if (!shapedType.getRank()) return nullptr;
                                 return shapedType.getElementType();
                               })
                           .Default([&](Type t) -> Type { return nullptr; });
    if (!elementType || !elementType.isIntOrFloat()) return;
    unsigned typeWidthInBytes =
        IREE::Util::getRoundedElementByteWidth(elementType);
    referenceTypeLengthInBytes =
        std::min<unsigned>(referenceTypeLengthInBytes, typeWidthInBytes);
  });
  return referenceTypeLengthInBytes;
}

/// Returns the default tile sizes to use for the loops that are distributed at
/// Flow level.
static SmallVector<int64_t> getDefaultDistributedLoopTileSizes(
    ArrayRef<int64_t> lbs, ArrayRef<int64_t> ubs,
    ArrayRef<int64_t> minTileSizes, ArrayRef<int64_t> maxTileSizes) {
  assert(lbs.size() == ubs.size() && lbs.size() == minTileSizes.size() &&
         lbs.size() == maxTileSizes.size() &&
         "expected all vectors to be of equal size");
  if (lbs.empty()) {
    return {};
  }
  size_t numDims = lbs.size();
  SmallVector<int64_t> distributedTileSizes(numDims, 1);
  SmallVector<int64_t> numWorkgroupsPerDim(numDims, 1);
  SmallVector<int64_t> workload(numDims, 1);
  auto ceilFn = [](int64_t a, int64_t b) { return (a + b - 1) / b; };

  for (auto i : llvm::seq<size_t>(0, numDims)) {
    if (ShapedType::isDynamic(lbs[i]) || ShapedType::isDynamic(ubs[i])) {
      distributedTileSizes[i] = maxTileSizes[i];
      workload[i] = ShapedType::kDynamicSize;
      continue;
    }
    int64_t candidateTileSize = 1;
    if (ubs[i] > lbs[i]) {
      // Pick a value that evenly distributes the workload.
      candidateTileSize = std::max<int64_t>(
          llvm::PowerOf2Floor(static_cast<uint64_t>(ubs[i] - lbs[i]) / 2),
          minTileSizes[i]);
    }

    // Limit the workload per workgroup to the default being the max to keep the
    // work per invocation reasonable.
    distributedTileSizes[i] =
        std::min<int64_t>(candidateTileSize, maxTileSizes[i]);
    workload[i] = (ubs[i] <= lbs[i] ? 1 : ubs[i] - lbs[i]);
    numWorkgroupsPerDim[i] = ceilFn(workload[i], distributedTileSizes[i]);
  }

  // Reduce the number of workgroups in cases where we are dividing the work too
  // much. Over-provision the number of workgroups to twice the number of
  // threads.
  int64_t numWorkgroupsLimit = 2 * clNumberOfRuntimeThreads;
  int64_t numWorkgroups = 1;
  for (auto ng : numWorkgroupsPerDim) {
    numWorkgroups *= ng;
  }
  unsigned currDim = numDims;
  while (numWorkgroups > numWorkgroupsLimit && currDim > 0) {
    if (distributedTileSizes[currDim - 1] >= maxTileSizes[currDim - 1] ||
        workload[currDim - 1] == ShapedType::kDynamicSize ||
        distributedTileSizes[currDim - 1] >= workload[currDim - 1]) {
      currDim--;
      continue;
    }
    distributedTileSizes[currDim - 1] = std::min<int64_t>(
        distributedTileSizes[currDim - 1] * 2, maxTileSizes[currDim - 1]);
    int64_t nwg =
        ceilFn(workload[currDim - 1], distributedTileSizes[currDim - 1]);
    if (nwg < numWorkgroupsPerDim[currDim - 1]) {
      numWorkgroups /= numWorkgroupsPerDim[currDim - 1];
      numWorkgroups *= nwg;
    } else {
      currDim--;
    }
  }
  return distributedTileSizes;
}

/// Adjusts the workload per workgroup to be a multiple of vector size to ensure
/// that the op vectorizes.
static int64_t getMaxTileSize(int64_t lb, int64_t ub, int64_t maxSize,
                              int64_t vectorSizeVal) {
  if (ub == ShapedType::kDynamicSize || lb == ShapedType::kDynamicSize) {
    return maxSize;
  }
  int64_t dim = ub - lb;
  if (dim < vectorSizeVal) return dim;
  for (int64_t i = std::min(maxSize, dim); i > 0; --i) {
    if (dim % i == 0 && i % vectorSizeVal == 0) {
      return i;
    }
  }
  // If it can't be a multiple of vectorSizeVal, let's choose a factor of dim
  // sizes heuristically.
  int64_t start = std::min(maxSize, dim);
  start = std::min(start, vectorSizeVal * 2);
  for (int64_t i = start; i > 0; --i) {
    if (dim % i == 0) {
      return i;
    }
  }
  return 1;
}

/// Returns the tile size to use for the Flow level of an operation that
/// implements the `PartitionableLoopsInterface`.
static SmallVector<int64_t> getDefaultDistributedLevelTileSizes(
    ArrayRef<Range> iterationDomain,
    IREE::Flow::PartitionableLoopsInterface partitionableLoopInterfaceOp,
    ArrayRef<int64_t> minTileSizes, ArrayRef<int64_t> maxTileSizes) {
  assert(iterationDomain.size() == minTileSizes.size() &&
         "expected as many min tile sizes as number of loops");
  auto getStaticValue = [](Value v) -> int64_t {
    IntegerAttr attr;
    if (!matchPattern(v, m_Constant(&attr))) return ShapedType::kDynamicSize;
    return attr.getInt();
  };
  auto lbs = llvm::to_vector(llvm::map_range(
      iterationDomain, [&](Range r) { return getStaticValue(r.offset); }));
  auto ubs = llvm::to_vector(llvm::map_range(
      iterationDomain, [&](Range r) { return getStaticValue(r.size); }));

  SmallVector<unsigned> partitionableLoops =
      partitionableLoopInterfaceOp.getPartitionableLoops(kNumMaxParallelDims);
  llvm::SmallDenseSet<unsigned, 4> partitionableLoopsSet;
  partitionableLoopsSet.insert(partitionableLoops.begin(),
                               partitionableLoops.end());

  size_t numPartitionedLoops = partitionableLoops.size();
  SmallVector<int64_t> distributedLoopLbs(numPartitionedLoops,
                                          ShapedType::kDynamicSize),
      distributedLoopUbs(numPartitionedLoops, ShapedType::kDynamicSize),
      minDistributedLoopTileSizes(numPartitionedLoops, 1),
      maxDistributedLoopTileSizes(numPartitionedLoops,
                                  defaultWorkgroupTileSize);
  // Find the bounds of the partitionable loops
  unsigned index = 0;
  for (auto range : llvm::enumerate(iterationDomain)) {
    if (!partitionableLoopsSet.count(range.index())) continue;

    minDistributedLoopTileSizes[index] = minTileSizes[range.index()];
    maxDistributedLoopTileSizes[index] = maxTileSizes[range.index()];
    distributedLoopLbs[index] = lbs[range.index()];
    distributedLoopUbs[index] = ubs[range.index()];
    index++;
  }

  SmallVector<int64_t> distributedTileSizes =
      getDefaultDistributedLoopTileSizes(distributedLoopLbs, distributedLoopUbs,
                                         minDistributedLoopTileSizes,
                                         maxDistributedLoopTileSizes);
  SmallVector<int64_t> distributedLevelTileSizes(iterationDomain.size(), 0);
  for (auto loopID : llvm::enumerate(partitionableLoops)) {
    distributedLevelTileSizes[loopID.value()] =
        distributedTileSizes[loopID.index()];
  }
  // Final fix up of the tile sizes to make sure that they divide the problem
  // size to make it vectorizable.
  for (auto i : llvm::seq<unsigned>(0, distributedLevelTileSizes.size())) {
    distributedLevelTileSizes[i] =
        distributedLevelTileSizes[i] != 0
            ? getMaxTileSize(lbs[i], ubs[i], distributedLevelTileSizes[i],
                             minTileSizes[i])
            : 0;
  }
  return distributedLevelTileSizes;
}

/// Splits the tile sizes in `parallelSizes` into `reductionSizes` for the
/// reduction loops.
static void splitParallelAndReductionTiles(
    linalg::LinalgOp op, SmallVectorImpl<int64_t> &parallelSizes,
    SmallVectorImpl<int64_t> &reductionSizes) {
  reductionSizes.assign(parallelSizes.begin(), parallelSizes.end());
  for (auto iteratorType : llvm::enumerate(op.iterator_types())) {
    if (iteratorType.value().cast<StringAttr>().getValue() ==
        getParallelIteratorTypeName()) {
      reductionSizes[iteratorType.index()] = 0;
    } else {
      parallelSizes[iteratorType.index()] = 0;
    }
  }
}

static void setAlwaysVectorizeSizes(linalg::LinalgOp op,
                                    SmallVectorImpl<int64_t> &parallelSizes,
                                    SmallVectorImpl<int64_t> &reductionSizes) {
  Optional<SmallVector<int64_t, 4>> staticLoopRanges = op.getStaticLoopRanges();
  for (auto en :
       llvm::enumerate(llvm::zip(*staticLoopRanges, op.iterator_types()))) {
    auto size = std::get<0>(en.value());
    if (!ShapedType::isDynamic(size)) continue;
    auto iterType = std::get<1>(en.value()).cast<StringAttr>().getValue();
    if (iterType == getParallelIteratorTypeName()) {
      parallelSizes[en.index()] = 1;
    } else {
      reductionSizes[en.index()] = 1;
    }
  }
}

/// Sets the default configuration to use for an operation that implements the
/// `PartitionableLoopsInterface`, given the iteration domain of all the loops.
static LogicalResult setDefaultRootConfig(
    func::FuncOp entryPointFn,
    IREE::Flow::PartitionableLoopsInterface partitionableLoopsInterfaceOp,
    ArrayRef<Range> iterationDomain) {
  if (getLoweringConfig(partitionableLoopsInterfaceOp)) return success();

  SmallVector<unsigned> partitionableLoops =
      partitionableLoopsInterfaceOp.getPartitionableLoops(kNumMaxParallelDims);

  SmallVector<int64_t> minTileSizes(iterationDomain.size(), 1);
  SmallVector<int64_t> maxTileSizes(iterationDomain.size(), 1);
  if (!partitionableLoops.empty()) {
    // TODO: Here the min tile size is just looking at the type of the data in
    // the entry point function, and using a vector size that depends on just
    // that. For `LinalgOp`s we can use the indexing map, find the loops that
    // are fastest varying and set those to have a min tile size of vector
    // length. A version of this is done for generic ops. Generalize that and
    // use it for `LinalgOp`s.
    unsigned typeWidthInBytes = getReferenceTypeLengthInBytes(entryPointFn);
    minTileSizes[partitionableLoops.back()] =
        getVectorSize(entryPointFn, typeWidthInBytes);
    for (auto partitionableLoopId : partitionableLoops) {
      maxTileSizes[partitionableLoopId] = defaultWorkgroupTileSize;
    }
  }

  SmallVector<int64_t> flowTileSizes = getDefaultDistributedLevelTileSizes(
      iterationDomain, partitionableLoopsInterfaceOp, minTileSizes,
      maxTileSizes);
  TileSizesListType tileSizes;
  tileSizes.emplace_back(std::move(flowTileSizes));
  return setOpConfigAndEntryPointFnTranslation(
      entryPointFn, partitionableLoopsInterfaceOp, tileSizes,
      DispatchLoweringPassPipeline::CPUDefault);
}

static LogicalResult setSandboxRootConfig(func::FuncOp entryPointFn,
                                          linalg::ContractionOpInterface op,
                                          ArrayRef<int64_t> flowTileSizes,
                                          ArrayRef<int64_t> target2ndTileSizes,
                                          int vectorSize) {
  assert(target2ndTileSizes.size() == 3 &&
         "the current configuration is driven by matmul which has exactly "
         "three loops");
  // The tiling for parallel dims and reduction dims should be separated.
  SmallVector<int64_t> parallelTileSizes;
  auto linalgOp = cast<linalg::LinalgOp>(op.getOperation());
  int64_t nLoops = linalgOp.getNumLoops();
  if (nLoops >= 3) {
    parallelTileSizes.append(nLoops - 3, 1);
    parallelTileSizes.push_back(getMaxTileSize(
        0, flowTileSizes[nLoops - 3], target2ndTileSizes[0], vectorSize));
  }
  if (nLoops >= 2) {
    parallelTileSizes.push_back(getMaxTileSize(
        0, flowTileSizes[nLoops - 2], target2ndTileSizes[1], vectorSize));
  }
  parallelTileSizes.push_back(0);

  auto lhsShapedType = op.lhs().getType().cast<ShapedType>();
  int64_t K = lhsShapedType.getShape().back();
  SmallVector<int64_t> reductionTileSizes;
  reductionTileSizes.append(nLoops - 1, 0);
  reductionTileSizes.push_back(
      getMaxTileSize(0, K, target2ndTileSizes[2], vectorSize));

  setAlwaysVectorizeSizes(linalgOp, parallelTileSizes, reductionTileSizes);

  TileSizesListType tileSizes;
  tileSizes.emplace_back(flowTileSizes.begin(), flowTileSizes.end());
  tileSizes.push_back(parallelTileSizes);
  tileSizes.push_back(reductionTileSizes);

  return setOpConfigAndEntryPointFnTranslation(
      entryPointFn, op, tileSizes,
      DispatchLoweringPassPipeline::CPUDoubleTilingExpert);
}

static LogicalResult setARMRootConfig(func::FuncOp entryPointFn,
                                      linalg::ContractionOpInterface op,
                                      ArrayRef<int64_t> flowTileSizes,
                                      int vectorSize) {
  // Hardcoded tile sizes, where v is the native vector size.
  // L1 tile sizes are {1, ..., 5v, v, 16v}.
  // Vector tile sizes are {1, ..., v, v, v}
  SmallVector<int64_t> l1TileSizes, vectorTileSizes;
  int64_t nLoops = cast<linalg::LinalgOp>(op.getOperation()).getNumLoops();
  if (nLoops >= 3) {
    l1TileSizes.append(nLoops - 3, 1);
    l1TileSizes.push_back(getMaxTileSize(0, flowTileSizes[nLoops - 3],
                                         5 * vectorSize, vectorSize));
    vectorTileSizes.append(nLoops - 3, 1);
    vectorTileSizes.push_back(vectorSize);
  }
  if (nLoops >= 2) {
    l1TileSizes.push_back(
        getMaxTileSize(0, flowTileSizes[nLoops - 2], vectorSize, vectorSize));
    vectorTileSizes.push_back(vectorSize);
  }

  // L1/vector tile size for k dimensions.
  auto lhsShapedType = op.lhs().getType().cast<ShapedType>();
  int64_t K = lhsShapedType.getShape().back();
  l1TileSizes.push_back(getMaxTileSize(0, K, 16 * vectorSize, vectorSize));
  vectorTileSizes.push_back(vectorSize);
  TileSizesListType tileSizes;
  tileSizes.emplace_back(flowTileSizes.begin(), flowTileSizes.end());
  tileSizes.push_back(l1TileSizes);
  tileSizes.push_back(vectorTileSizes);

  return setOpConfigAndEntryPointFnTranslation(
      entryPointFn, op, tileSizes,
      DispatchLoweringPassPipeline::CPUTileFuseAndVectorize);
}

/// Sets the lowering configuration for dispatch region with root op that
/// implements the contraction operation interface.
static LogicalResult setRootConfig(
    func::FuncOp entryPointFn, linalg::ContractionOpInterface contractionOp,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  auto linalgOp = cast<linalg::LinalgOp>(contractionOp.getOperation());
  unsigned numLoops = linalgOp.getNumLoops();
  {
    SmallVector<unsigned> dims;
    linalgOp.getReductionDims(dims);
    if (dims.size() != 1 || dims[0] != numLoops - 1) {
      return contractionOp.emitOpError(
          "expected to have exactly one reduction dim, and it is the innermost "
          "dim");
    }
  }

  // Consider all element types and use the smallest vector size. The tiling
  // sizes are chosen based on the vector size.
  auto lhsShapedType = contractionOp.lhs().getType().cast<ShapedType>();
  auto rhsShapedType = contractionOp.rhs().getType().cast<ShapedType>();
  auto resShapedType =
      linalgOp.getOutputOperand(0)->get().getType().cast<ShapedType>();
  int64_t vectorSize = getVectorSize(entryPointFn, lhsShapedType);
  vectorSize = std::min(vectorSize, getVectorSize(entryPointFn, rhsShapedType));
  vectorSize = std::min(vectorSize, getVectorSize(entryPointFn, resShapedType));

  // Use the default distribution for the matmul loops.
  SmallVector<int64_t> minTileSizes =
      getMinTilingSizesForEachDim(entryPointFn, linalgOp);
  SmallVector<int64_t> maxTileSizes(numLoops, defaultWorkgroupTileSize);
  if (numLoops > 3) {
    minTileSizes[0] = 1;
    maxTileSizes[0] = 1;
  }

  OpBuilder builder(entryPointFn.getContext());
  builder.setInsertionPoint(contractionOp);
  SmallVector<Range> iterationDomain =
      linalgOp.createLoopRanges(builder, linalgOp->getLoc());
  SmallVector<int64_t> flowTileSizes = getDefaultDistributedLevelTileSizes(
      iterationDomain,
      cast<IREE::Flow::PartitionableLoopsInterface>(
          contractionOp.getOperation()),
      minTileSizes, maxTileSizes);

  // TODO(dcaballe): Find better configurations for RISC-V backends.
  if (isX86(entryPointFn) || isRISCV(entryPointFn)) {
    SmallVector<int64_t> tileSizes = {8, 32, 16};
    return setSandboxRootConfig(entryPointFn, contractionOp, flowTileSizes,
                                tileSizes, vectorSize);
  }

  // Fall back to ARM configurations.
  bool isQuantized =
      lhsShapedType.getElementType() != resShapedType.getElementType();
  if (isQuantized) {
    SmallVector<int64_t> tileSizes = {4, 16, 4};
    return setSandboxRootConfig(entryPointFn, contractionOp, flowTileSizes,
                                tileSizes, vectorSize);
  } else {
    return setARMRootConfig(entryPointFn, contractionOp, flowTileSizes,
                            vectorSize);
  }
}

/// Sets the lowering configuration for dispatch region for linalg.mmt4d root
/// op
static LogicalResult setRootConfig(
    func::FuncOp entryPointFn, linalg::Mmt4DOp mmt4dOp,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  // TODO(ataei): These are hand tuned for some performance benchmarks for
  // now, we want to adapt the same strategy as matmul that dynamically sets
  // tile size.
  auto getWorkgroupTileSizes = [&]() -> SmallVector<int64_t> {
    if (!mmt4dWorkgroupTileSizes.empty()) {
      return SmallVector<int64_t>(mmt4dWorkgroupTileSizes.begin(),
                                  mmt4dWorkgroupTileSizes.end());
    }
    return {48, 32};
  };

  auto getL1TileSizes = [&]() -> SmallVector<int64_t> {
    auto lhsShape = mmt4dOp.inputs()[0].getType().cast<ShapedType>().getShape();
    auto rhsShape = mmt4dOp.inputs()[1].getType().cast<ShapedType>().getShape();
    int M0 = lhsShape[2];
    int N0 = rhsShape[2];
    int K0 = lhsShape[3];
    if (!mmt4dL1TileSizes.empty()) {
      return SmallVector<int64_t>(mmt4dL1TileSizes.begin(),
                                  mmt4dL1TileSizes.end());
    }
    return {1, 1, 1, M0, N0, K0};
  };

  auto getVectorSizes = [&]() -> SmallVector<int64_t> {
    auto lhsShape = mmt4dOp.inputs()[0].getType().cast<ShapedType>().getShape();
    auto rhsShape = mmt4dOp.inputs()[1].getType().cast<ShapedType>().getShape();
    int M0 = lhsShape[2];
    int N0 = rhsShape[2];
    int K0 = lhsShape[3];
    if (!mmt4dVectorSizes.empty()) {
      return SmallVector<int64_t>(mmt4dVectorSizes.begin(),
                                  mmt4dVectorSizes.end());
    }
    return {1, 1, 1, M0, N0, K0};
  };

  SmallVector<int64_t> nativeVectorSize = getVectorSizes();

  TileSizesListType tileSizes = {getWorkgroupTileSizes(), getL1TileSizes(),
                                 nativeVectorSize};

  return setOpConfigAndEntryPointFnTranslation(
      entryPointFn, mmt4dOp, tileSizes,
      DispatchLoweringPassPipeline::CPUTileFuseAndVectorize);
}

/// Sets the lowering configuration for dispatch region for linalg_ext.fft
/// root op.
static LogicalResult setRootConfig(
    func::FuncOp entryPointFn, IREE::LinalgExt::FftOp fftOp,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  unsigned numLoops = fftOp.getLoopIteratorTypes().size();
  auto partitionedLoops = fftOp.getPartitionableLoops(kNumMaxParallelDims);
  SmallVector<int64_t> workgroupTileSizes(numLoops, defaultWorkgroupTileSize);
  llvm::DenseSet<unsigned> partitionedLoopsSet(partitionedLoops.begin(),
                                               partitionedLoops.end());
  for (auto dim : llvm::seq<int64_t>(0, workgroupTileSizes.size())) {
    if (!partitionedLoopsSet.count(dim)) {
      workgroupTileSizes[dim] = 0;
    }
  }

  auto rank = fftOp.getOperandRank();
  if (workgroupTileSizes.size() >= rank && workgroupTileSizes[rank - 1] != 0) {
    APInt value;
    if (matchPattern(fftOp.getStage(), m_ConstantInt(&value))) {
      workgroupTileSizes[rank - 1] = 1ll << value.getSExtValue();
      workgroupTileSizes[rank - 1] =
          std::max(workgroupTileSizes[rank - 1],
                   static_cast<int64_t>(defaultWorkgroupTileSize));
    } else {
      return fftOp.emitOpError("non-constant stage might not work for fft op");
    }
  }
  TileSizesListType tileSizes = {workgroupTileSizes};
  return setOpConfigAndEntryPointFnTranslation(
      entryPointFn, fftOp, tileSizes, DispatchLoweringPassPipeline::CPUDefault);
}

static void setX86WorkgroupTileSizes(
    linalg::GenericOp genericOp, unsigned numLoops,
    ArrayRef<int64_t> flowTileSizes, ArrayRef<int64_t> minTileSizes,
    ArrayRef<int64_t> maxTileSizes,
    SmallVectorImpl<int64_t> &workgroupTileSizes) {
  workgroupTileSizes.append(numLoops, 0);
  Optional<SmallVector<int64_t, 4>> staticLoopRanges =
      genericOp.getStaticLoopRanges();
  for (auto loopNum : llvm::seq<unsigned>(0, numLoops)) {
    if (flowTileSizes[loopNum]) {
      workgroupTileSizes[loopNum] =
          getMaxTileSize(0, flowTileSizes[loopNum], minTileSizes[loopNum],
                         minTileSizes[loopNum]);
    } else {
      // If the flow level tile size is zero, and static loop range is 0 as
      // well, set the tile sizes here to zero as well.
      workgroupTileSizes[loopNum] =
          (staticLoopRanges && staticLoopRanges.getValue()[loopNum] == 1)
              ? 0
              : minTileSizes[loopNum];
    }
  }
}

/// Returns true if the operation is a GenericOp implementing a supported
/// transposition.
static bool isSupportedTransposeOp(linalg::GenericOp genericOp) {
  // Check that the op has at least 2 dimensions.
  if (genericOp.getNumLoops() < 2) {
    return false;
  }

  // Check that the op has only one input and one output.
  // TODO(diegocaballero): Generalize to multiple inputs.
  if ((genericOp.getNumInputs() != 1) || (genericOp.getNumOutputs() != 1)) {
    return false;
  }

  // Check that all the iterators are parallel.
  if (genericOp.getNumParallelLoops() != genericOp.getNumLoops()) {
    return false;
  }

  // Check that the two indexing maps are a permutation of each other.
  auto indexing_maps = genericOp.getIndexingMaps();
  return !indexing_maps[0].isEmpty() && !indexing_maps[1].isEmpty() &&
         ((indexing_maps[0].isIdentity() && !indexing_maps[1].isIdentity() &&
           indexing_maps[1].isPermutation()) ||
          (!indexing_maps[0].isIdentity() && indexing_maps[0].isPermutation() &&
           indexing_maps[1].isIdentity()));
}

/// Sets the default lowering configuration for a generic op to use
/// CPUDoubleTilingExpert pipeline.
static LogicalResult setDefaultGenericOpRootConfig(
    func::FuncOp entryPointFn, linalg::GenericOp genericOp,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  if (getLoweringConfig(genericOp)) {
    return success();
  }

  // If there are no loops, there is nothing to do.
  unsigned numLoops = genericOp.getNumLoops();
  if (numLoops == 0) return success();

  SmallVector<int64_t> minTileSizes =
      getMinTilingSizesForEachDim(entryPointFn, genericOp);
  SmallVector<int64_t> maxTileSizes(numLoops, defaultWorkgroupTileSize);
  if (llvm::all_of(minTileSizes, [](int64_t vs) { return vs == 1; })) {
    // Nothing to vectorize just lower to loops.
    return success();
  }

  // Set the flow level tiling to the default.
  OpBuilder builder(genericOp.getContext());
  builder.setInsertionPoint(genericOp);
  auto linalgOp = cast<linalg::LinalgOp>(genericOp.getOperation());
  SmallVector<Range> iterationDomain =
      linalgOp.createLoopRanges(builder, genericOp.getLoc());
  auto partitionableLoopsInterfaceOp =
      cast<IREE::Flow::PartitionableLoopsInterface>(genericOp.getOperation());
  SmallVector<int64_t> flowTileSizes = getDefaultDistributedLevelTileSizes(
      iterationDomain, partitionableLoopsInterfaceOp, minTileSizes,
      maxTileSizes);

  // Set the next level tile sizes.
  SmallVector<int64_t> parallelTileSizes;
  SmallVector<int64_t> reductionTileSizes;
  setX86WorkgroupTileSizes(genericOp, numLoops, flowTileSizes, minTileSizes,
                           maxTileSizes, parallelTileSizes);
  splitParallelAndReductionTiles(genericOp, parallelTileSizes,
                                 reductionTileSizes);
  setAlwaysVectorizeSizes(linalgOp, parallelTileSizes, reductionTileSizes);

  TileSizesListType tileSizes;
  tileSizes.push_back(flowTileSizes);
  tileSizes.push_back(parallelTileSizes);
  tileSizes.push_back(reductionTileSizes);

  // For non-tensor based ops use the Buffer ops pipeline.
  auto passPipeline =
      genericOp.hasTensorSemantics()
          ? DispatchLoweringPassPipeline::CPUDoubleTilingExpert
          : DispatchLoweringPassPipeline::CPUBufferOpsTileAndVectorize;
  return setOpConfigAndEntryPointFnTranslation(entryPointFn, genericOp,
                                               tileSizes, passPipeline);
}

/// Sets the lowering configuration for a generic op implementing a
/// transposition to use CPUDoubleTilingExpert pipeline.
static LogicalResult setTransposeLikeOpRootConfig(
    func::FuncOp entryPointFn, linalg::GenericOp genericOp,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  if (getLoweringConfig(genericOp)) {
    return success();
  }

  if (!hasAVX2Features(genericOp) || !isSupportedTransposeOp(genericOp)) {
    return success();
  }

  unsigned numLoops = genericOp.getNumLoops();
  SmallVector<int64_t> minTileSizes =
      getMinTilingSizesForEachDim(entryPointFn, genericOp);
  SmallVector<int64_t> maxTileSizes(numLoops, defaultWorkgroupTileSize);
  if (llvm::all_of(minTileSizes, [](int64_t vs) { return vs == 1; })) {
    // Nothing to vectorize just lower to loops.
    return success();
  }

  if (llvm::count_if(minTileSizes,
                     [](int64_t tileSize) { return tileSize > 1; }) != 2) {
    // Transpose patterns are not applicable if vectorizing more or less than
    // two dims.
    return success();
  }

  // Make sure that the original tile sizes are multiple of the tile sizes
  // to be used for the transpose op (i.e., 8x8).
  // TODO(diegocaballero): Enable 4x8 tile sizes if we find it useful.
  if (llvm::any_of(minTileSizes, [](int64_t tileSize) {
        return tileSize > 1 && (tileSize % 8) != 0;
      })) {
    return success();
  }

  // Replace dims to be vectorized with the new 8x8 tile sizes.
  std::replace_if(
      minTileSizes.begin(), minTileSizes.end(),
      [](int64_t tileSize) { return tileSize > 1; }, 8);

  // Set the flow level tiling to the default.
  OpBuilder builder(genericOp.getContext());
  builder.setInsertionPoint(genericOp);
  auto linalgOp = cast<linalg::LinalgOp>(genericOp.getOperation());
  SmallVector<Range> iterationDomain =
      linalgOp.createLoopRanges(builder, genericOp.getLoc());
  auto partitionableLoopsInterfaceOp =
      cast<IREE::Flow::PartitionableLoopsInterface>(genericOp.getOperation());
  SmallVector<int64_t> flowTileSizes = getDefaultDistributedLevelTileSizes(
      iterationDomain, partitionableLoopsInterfaceOp, minTileSizes,
      maxTileSizes);

  // Set the next level tile sizes.
  SmallVector<int64_t> parallelTileSizes;
  setX86WorkgroupTileSizes(genericOp, numLoops, flowTileSizes, minTileSizes,
                           maxTileSizes, parallelTileSizes);

  TileSizesListType tileSizes;
  tileSizes.push_back(flowTileSizes);
  tileSizes.push_back(parallelTileSizes);
  tileSizes.push_back(/*reduction tile sizes*/ {});

  // For non-tensor based ops use the Buffer ops pipeline.
  auto passPipeline =
      genericOp.hasTensorSemantics()
          ? DispatchLoweringPassPipeline::CPUDoubleTilingExpert
          : DispatchLoweringPassPipeline::CPUBufferOpsTileAndVectorize;
  return setOpConfigAndEntryPointFnTranslation(entryPointFn, genericOp,
                                               tileSizes, passPipeline);
}

/// Sets the lowering configuration for a generic op to use
/// CPUDoubleTilingExpert pipeline.
static LogicalResult setRootConfig(
    func::FuncOp entryPointFn, linalg::GenericOp genericOp,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  if (failed(
          setTransposeLikeOpRootConfig(entryPointFn, genericOp, tiledLoops)) ||
      failed(
          setDefaultGenericOpRootConfig(entryPointFn, genericOp, tiledLoops))) {
    return failure();
  }
  return success();
}

/// Sets the lowering configuration for linalg.conv_2d_nhwc_hwcf and
/// linalg.depthwise_conv_2d_nhwc_hwc operations.
static LogicalResult setConvRootConfig(
    func::FuncOp entryPointFn, linalg::LinalgOp convOp,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops,
    ArrayRef<int64_t> targetTileSizes, int64_t vectorSize) {
  if (!isa<linalg::Conv2DNhwcHwcfOp, linalg::DepthwiseConv2DNhwcHwcOp>(
          convOp.getOperation())) {
    return failure();
  }

  // Use the default distribution for the conv loops.
  unsigned numLoops = convOp.getNumLoops();
  SmallVector<int64_t> minTileSizes(numLoops, 1);
  SmallVector<int64_t> maxTileSizes(numLoops, defaultWorkgroupTileSize);

  // Set the flow level tiling to the default.
  OpBuilder builder(convOp.getContext());
  builder.setInsertionPoint(convOp);
  SmallVector<Range> iterationDomain =
      convOp.createLoopRanges(builder, convOp.getLoc());
  auto partitionableLoopsInterfaceOp =
      cast<IREE::Flow::PartitionableLoopsInterface>(convOp.getOperation());
  SmallVector<int64_t> flowTileSizes = getDefaultDistributedLevelTileSizes(
      iterationDomain, partitionableLoopsInterfaceOp, minTileSizes,
      maxTileSizes);

  // Shapes of N, OH, OW, OC, KH, KW, (IC)
  Optional<SmallVector<int64_t, 4>> shapes = convOp.getStaticLoopRanges();
  SmallVector<int64_t> parallelTileSizes(targetTileSizes.begin(),
                                         targetTileSizes.end());
  for (auto i : llvm::seq<unsigned>(0, parallelTileSizes.size())) {
    auto tileSize = flowTileSizes[i] ? flowTileSizes[i] : shapes.getValue()[i];
    // If the tile size is intended to be 1, do not adjust it to `vectorSize`.
    // The ops will be decomposed to lower-rank named ops.
    if (parallelTileSizes[i] != 1) {
      parallelTileSizes[i] =
          getMaxTileSize(0, tileSize, parallelTileSizes[i], vectorSize);
    }
  }
  SmallVector<int64_t> reductionTileSizes;
  splitParallelAndReductionTiles(convOp, parallelTileSizes, reductionTileSizes);
  setAlwaysVectorizeSizes(convOp, parallelTileSizes, reductionTileSizes);

  TileSizesListType tileSizes;
  tileSizes.push_back(flowTileSizes);
  tileSizes.push_back(parallelTileSizes);
  tileSizes.push_back(reductionTileSizes);
  return setOpConfigAndEntryPointFnTranslation(
      entryPointFn, convOp, tileSizes,
      DispatchLoweringPassPipeline::CPUConvTileAndDecomposeExpert);
}

static LogicalResult setRootConfig(
    func::FuncOp entryPointFn, linalg::Conv2DNhwcHwcfOp convOp,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  auto linalgOp = cast<linalg::LinalgOp>(convOp.getOperation());
  int64_t vectorSize =
      getVectorSize(entryPointFn, convOp.getResult(0).getType());
  SmallVector<int64_t> targetTileSizes = {1, 1, 8, vectorSize * 2, 1, 1, 8};
  return setConvRootConfig(entryPointFn, linalgOp, tiledLoops, targetTileSizes,
                           vectorSize);
}

/// Sets the lowering configuration for linalg.depthwise_conv_2d_nhwc_hwc
/// operations.
static LogicalResult setRootConfig(
    func::FuncOp entryPointFn, linalg::DepthwiseConv2DNhwcHwcOp convOp,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  auto linalgOp = cast<linalg::LinalgOp>(convOp.getOperation());
  int64_t vectorSize =
      getVectorSize(entryPointFn, convOp.getResult(0).getType());
  SmallVector<int64_t> targetTileSizes = {1, 1, 8, vectorSize * 2, 1, 3};
  return setConvRootConfig(entryPointFn, linalgOp, tiledLoops, targetTileSizes,
                           vectorSize);
}

/// Set default configuration for Linalg ops.
static LogicalResult setRootConfig(
    func::FuncOp entryPointFn, linalg::LinalgOp linalgOp,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  if (getLoweringConfig(linalgOp)) return success();

  OpBuilder builder(linalgOp.getContext());
  builder.setInsertionPoint(linalgOp);
  SmallVector<Range> iterationDomain =
      linalgOp.createLoopRanges(builder, linalgOp.getLoc());

  auto partitionableLoopOp =
      cast<IREE::Flow::PartitionableLoopsInterface>(linalgOp.getOperation());
  return setDefaultRootConfig(entryPointFn, partitionableLoopOp,
                              iterationDomain);
}

/// Set the default configuration for operations that implement the
/// `TiledOpInterface`.
static LogicalResult setRootConfig(
    func::FuncOp entryPointFn,
    IREE::LinalgExt::TiledOpInterface tiledOpInterfaceOp,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  if (getLoweringConfig(tiledOpInterfaceOp)) return success();

  OpBuilder builder(tiledOpInterfaceOp.getContext());
  builder.setInsertionPoint(tiledOpInterfaceOp);
  SmallVector<Range> iterationDomain =
      tiledOpInterfaceOp.getIterationDomain(builder);
  auto partitionableLoopInterfaceOp =
      cast<IREE::Flow::PartitionableLoopsInterface>(
          tiledOpInterfaceOp.getOperation());
  return setDefaultRootConfig(entryPointFn, partitionableLoopInterfaceOp,
                              iterationDomain);
}

/// Redirects to methods that set the configuration based on operation type.
static LogicalResult setRootConfigImpl(
    func::FuncOp entryPointFn, Operation *op,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  // Do not overwrite default configuration.
  if (getLoweringConfig(op)) return success();

  // Redirect to individual operations.
  auto setRootConfigFn = [&](Operation *op) -> LogicalResult {
    return TypeSwitch<Operation *, LogicalResult>(op)
        .Case<IREE::LinalgExt::FftOp, linalg::GenericOp, linalg::Mmt4DOp,
              linalg::Conv2DNhwcHwcfOp, linalg::DepthwiseConv2DNhwcHwcOp>(
            [&](auto op) {
              return setRootConfig(entryPointFn, op, tiledLoops);
            })
        .Case<linalg::ContractionOpInterface>([&](auto op) {
          return setRootConfig(entryPointFn, op, tiledLoops);
        })
        .Case<linalg::LinalgOp, IREE::LinalgExt::TiledOpInterface>(
            [&](auto op) {
              return setRootConfig(entryPointFn, op, tiledLoops);
            })
        .Default([&](Operation *op) { return success(); });
  };
  return setRootConfigFn(op);
}

/// Redirects to methods that set the configuration based on operation type for
/// VMVX backend.
static LogicalResult setVMVXRootConfigImpl(
    func::FuncOp entryPointFn, Operation *op,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  if (getLoweringConfig(op)) return success();

  // Redirect to individual operations.
  auto setRootConfigFn = [&](Operation *op) -> LogicalResult {
    return TypeSwitch<Operation *, LogicalResult>(op)
        .Case<linalg::LinalgOp, IREE::LinalgExt::TiledOpInterface>(
            [&](auto op) {
              return setRootConfig(entryPointFn, op, tiledLoops);
            })
        .Default([&](Operation *op) { return success(); });
  };
  return setRootConfigFn(op);
}

/// Find the root operation for the dispatch region.
static FailureOr<Operation *> getRootOperation(
    ArrayRef<Operation *> computeOps) {
  Operation *rootOperation = nullptr;
  auto updateRootOperation = [&](Operation *op) -> LogicalResult {
    if (rootOperation) {
      return op->emitOpError(
          "unhandled multiple root operations in dispatch region");
    }
    rootOperation = op;
    return success();
  };
  for (auto op : computeOps) {
    if (auto linalgOp = dyn_cast<linalg::LinalgOp>(op)) {
      // Do not not treat linalg ops that are all parallel as root operations in
      // this sweep.
      if (linalgOp.getNumLoops() == linalgOp.getNumParallelLoops()) continue;

      // All other linalg ops are root ops.
      if (failed(updateRootOperation(op))) return failure();
      continue;
    }

    if (auto tiledOpInterfaceOp =
            dyn_cast<IREE::LinalgExt::TiledOpInterface>(op)) {
      // TODO(ravishankarm): For now
      // `tensor.extract_slice`/`tensor.insert_slice` implement the
      // `tiledInterfaceOp`. With tile + distribute moved out of Flow
      // dialect, this doesnt work anymore. Remove this when the external
      // model implementation of
      // `tensor.extract_slice`/`tensor.insert_slice` are dropped.
      if (isa<tensor::ExtractSliceOp, tensor::InsertSliceOp>(op)) continue;

      // All other operations that implement this interface are root ops.
      if (failed(updateRootOperation(op))) return failure();
      continue;
    }
  }
  if (rootOperation) return rootOperation;

  // If no root operation is found yet. Look for linalg generic ops.
  for (auto op : llvm::reverse(computeOps)) {
    if (isa<linalg::LinalgOp>(op)) {
      if (failed(updateRootOperation(op))) return failure();
    }
  }
  return rootOperation;
}

/// Finds the root operation in the given list of Linalg operations and sets
/// its configuration. Returns error for multiple root operations.
static LogicalResult setRootConfig(
    func::FuncOp entryPointFn, ArrayRef<Operation *> computeOps,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  FailureOr<Operation *> rootOp = getRootOperation(computeOps);
  if (failed(rootOp)) {
    return failure();
  }
  Operation *rootOperation = rootOp.getValue();

  if (rootOperation) {
    if (isVMVXBackend(entryPointFn)) {
      if (failed(
              setVMVXRootConfigImpl(entryPointFn, rootOperation, tiledLoops))) {
        return failure();
      }
    } else {
      if (failed(setRootConfigImpl(entryPointFn, rootOperation, tiledLoops))) {
        return failure();
      }
    }
  }

  if (!getTranslationInfo(entryPointFn)) {
    // Fall back, just set the translation to CPUDefault.
    setTranslationInfo(entryPointFn, DispatchLoweringPassPipeline::CPUDefault,
                       /*workloadPerWorkgroup=*/ArrayRef<int64_t>{},
                       /*workgroupSize=*/ArrayRef<int64_t>{});
  }

  return success();
}

/// Sets the translation information to use for a dispatch region.
static LogicalResult setTranslationInfoAndRootConfig(
    func::FuncOp entryPointFn, ArrayRef<Operation *> computeOps,
    ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) {
  // First check if the operations have a preset pipeline.
  for (auto computeOp : computeOps) {
    if (IREE::Codegen::CompilationInfoAttr compilationInfo =
            getCompilationInfo(computeOp)) {
      // If the function already has a translation, error out.
      if (auto translationInfo = getTranslationInfo(entryPointFn)) {
        return computeOp->emitOpError(
            "multiple ops within dispatch trying to set the translation "
            "info");
      }

      SmallVector<int64_t> workgroupSize =
          compilationInfo.getWorkgroupSizeVals();
      setTranslationInfo(entryPointFn, compilationInfo.getTranslationInfo(),
                         workgroupSize);
      setLoweringConfig(computeOp, compilationInfo.getLoweringConfig());
      eraseCompilationInfo(computeOp);
    }
  }

  // Next set the configuration of the operations.
  return setRootConfig(entryPointFn, computeOps, tiledLoops);
}

LogicalResult initCPULaunchConfig(ModuleOp moduleOp) {
  llvm::StringMap<IREE::HAL::ExecutableEntryPointOp> entryPointOps =
      getAllEntryPoints(moduleOp);
  for (auto funcOp : moduleOp.getOps<func::FuncOp>()) {
    auto entryPointOp = entryPointOps.lookup(funcOp.getName());
    if (!entryPointOp) continue;
    if (getTranslationInfo(entryPointOp)) continue;

    // If using sandbox passes, currently set the workload_per_wg to be
    // empty for single-threaded execution.
    if (useLinalgTransformInterp) {
      auto translationInfo = IREE::Codegen::TranslationInfoAttr::get(
          moduleOp.getContext(), IREE::Codegen::DispatchLoweringPassPipeline::
                                     LinalgTransformInterpCodegen);
      setTranslationInfo(funcOp, translationInfo);
      continue;
    }

    SmallVector<Operation *> computeOps;
    SmallVector<LoopTilingAndDistributionInfo> tiledLoops;

    // If there are no linalg ops, not using Linalg based lowering.
    if (failed(getComputeOps(funcOp, computeOps, tiledLoops))) {
      return failure();
    }

    if (failed(
            setTranslationInfoAndRootConfig(funcOp, computeOps, tiledLoops))) {
      return failure();
    }
  }

  // The root confguration setting introduces `tensor.dim` operations. Resolve
  // those away.
  RewritePatternSet patterns(moduleOp.getContext());
  memref::populateResolveRankedShapeTypeResultDimsPatterns(patterns);
  return applyPatternsAndFoldGreedily(moduleOp, std::move(patterns));
}

}  // namespace iree_compiler
}  // namespace mlir