// 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

//===- KernelDispatchUtils.cpp - Utilities for generating dispatch info ---===//
//
// This file defines utility functions that can be used to get the information
// about tile sizes to use to partition work across workgroups, the workgroup
// sizes and to create information the dispatch on the host side needs to
// execute an entry point function (e.g. total number of workgroups).
//
//===----------------------------------------------------------------------===//

#include "iree/compiler/Codegen/SPIRV/KernelDispatchUtils.h"

#include "iree/compiler/Codegen/Passes.h"
#include "iree/compiler/Codegen/SPIRV/LaunchConfig.h"
#include "iree/compiler/Codegen/SPIRV/Utils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/Shape/IR/ShapeOps.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "llvm/Support/Debug.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/Dialect/Vector/VectorTransforms.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/PatternMatch.h"

#define DEBUG_TYPE "kernel-dispatch-utils"

namespace mlir {
namespace iree_compiler {

/// Given `nprocs` try to distribute it evenly across 2 logical x and y.
static std::tuple<int64_t, int64_t> distributeProcs2D(int64_t nprocs) {
  int64_t nprocs_x = std::max<int64_t>(
      1, static_cast<int64_t>(
             llvm::PowerOf2Ceil(static_cast<uint64_t>(std::sqrt(nprocs)))));
  return std::make_tuple(nprocs_x, nprocs / nprocs_x);
}

/// Returns the minimum of `shape` and `tileSize` if shape is static. If `shape`
/// is dynamic returns `tileSize`.
static int64_t getMinIfShapeStatic(int64_t shape, int64_t tileSize) {
  if (shape == ShapedType::kDynamicSize) return tileSize;
  return std::min(shape, tileSize);
}

/// Sets the `tileSizes` and `workgroupSize` for an Linalg `op` to the default,
/// where at most 3 inner parallel dimensions of `op` are tiled and distributed,
/// and each invocation handles one scalar elements.
// TODO(#5852): revisit the default here: they were chosen to get started and
// not very good.
static LogicalResult setDefaultTilingScheme(
    const spirv::TargetEnv &targetEnv, linalg::LinalgOp op,
    TileSizesListType &tileSizes, std::array<int64_t, 3> &workgroupSize) {
  auto maxWorkgroupSize =
      targetEnv.getResourceLimits().max_compute_workgroup_invocations();

  const int64_t tileSizeX = 32;
  const int64_t tileSizeY = maxWorkgroupSize.getInt() / tileSizeX;

  unsigned numParallelDims = getNumOuterParallelLoops(op);

  SmallVector<int64_t, 4> workgroupLevel(numParallelDims, 0);
  SmallVector<int64_t, 4> invocationLevel(numParallelDims, 0);

  if (numParallelDims >= 1) {
    workgroupLevel.back() = tileSizeX;
    invocationLevel.back() = 1;
  }
  if (numParallelDims >= 2) {
    workgroupLevel[numParallelDims - 2] = tileSizeY;
    invocationLevel[numParallelDims - 2] = 1;
  }
  if (numParallelDims >= 3) {
    workgroupLevel[numParallelDims - 3] = 1;
    invocationLevel[numParallelDims - 3] = 1;
  }

  tileSizes.emplace_back(std::move(workgroupLevel));
  tileSizes.emplace_back();  // Subgroup level
  tileSizes.emplace_back(std::move(invocationLevel));

  workgroupSize = {tileSizeX, tileSizeY, 1};

  return success();
}

/// Fills `inputTypes` and `outputTypes` with the original input/output types
/// for all tiles for `op`.
static std::tuple<SmallVector<ShapedType>, SmallVector<ShapedType>>
getInputOutputTypes(linalg::LinalgOp op) {
  SmallVector<ShapedType> inputTypes(op.getNumInputs()),
      outputTypes(op.getNumOutputs());
  auto inputOperands = op.getInputOperands();
  for (auto operand : enumerate(inputOperands)) {
    assert(!op.isScalar(operand.value()));
    inputTypes[operand.index()] =
        getUntiledType(operand.value()->get()).dyn_cast<ShapedType>();
  }
  auto outputOperands = op.getOutputOperands();
  for (auto operand : enumerate(outputOperands)) {
    outputTypes[operand.index()] =
        getUntiledType(operand.value()->get()).dyn_cast<ShapedType>();
  }
  return std::make_tuple(std::move(inputTypes), std::move(outputTypes));
}

namespace {
struct LaunchConfigInfo {
  std::array<int64_t, 3> workgroupSize = {32, 1, 1};
  std::array<int64_t, 3> numSubgroups = {1, 1, 1};
  bool vectorize = false;
};

struct TileWorkgroupSizePair {
  // How many scalar elements each workgroup should handle along each dimension.
  std::array<int64_t, 3> tileSize;
  std::array<int64_t, 3> workgroupSize;
};
}  // namespace

/// For a given operation `op`, compute the following configurations according
/// to SPIR-V `targetEnv` and `options`:
/// 1) number of tiling levels and tile sizes to use (updates `tileSizes`),
/// 2) workgroup size to use (updates `workgroupSize`),
/// 3) number of subgroups to use if two level tiling is used (updates
///    `numSubgroups`).
template <typename T>
static LogicalResult getOpLaunchConfig(T op, const spirv::TargetEnv &targetEnv,
                                       const SPIRVCodegenOptions &options,
                                       TileSizesListType &tileSizes,
                                       LaunchConfigInfo &config) {
  return setDefaultTilingScheme(targetEnv, op, tileSizes, config.workgroupSize);
}

static void getMaliBestMatMulTileSizes(
    Type elementType, SmallVectorImpl<TileWorkgroupSizePair> &tileSizes,
    int64_t dstSize) {
  if (elementType.isF16()) {
    const int64_t smallMatrixSizeThreshold = 512 * 512;
    // For smaller destination size we cannot fill out the GPU with bigger tile
    // sizes. Instead we pick smaller tiles along M and N to increase the number
    // of workgroups and a larger K tile size since we have lower pressure and
    // need extra instructions to hide latency.
    // TODO: The threshold needs to be fine tuned by doing exploration based on
    // matrix shapes.
    if (dstSize <= smallMatrixSizeThreshold) {
      tileSizes.push_back(TileWorkgroupSizePair({{16, 32, 16}, {8, 2, 1}}));
    } else {
      tileSizes.push_back(TileWorkgroupSizePair({{16, 64, 4}, {8, 2, 1}}));
      tileSizes.push_back(TileWorkgroupSizePair({{8, 128, 4}, {8, 2, 1}}));
      tileSizes.push_back(TileWorkgroupSizePair({{16, 32, 4}, {8, 2, 1}}));
    }
  } else {
    // TODO: Heuristic picked based on MobileNet performance. We need
    // auto-tuning to be able to make a smarter choice.
    const int64_t smallMatrixSizeThreshold = 20000;
    if (dstSize <= smallMatrixSizeThreshold) {
      tileSizes.push_back(TileWorkgroupSizePair({{4, 32, 16}, {8, 2, 1}}));
    }
    tileSizes.push_back(TileWorkgroupSizePair({{12, 32, 4}, {8, 2, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{14, 32, 4}, {8, 2, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{10, 32, 4}, {8, 2, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{7, 64, 4}, {16, 1, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{8, 32, 4}, {8, 2, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{6, 32, 4}, {8, 2, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{24, 16, 4}, {2, 8, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{16, 16, 4}, {2, 8, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{24, 8, 4}, {2, 8, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{40, 8, 4}, {2, 8, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{32, 8, 4}, {2, 8, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{16, 8, 4}, {2, 8, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{1, 32, 16}, {8, 1, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{1, 32, 8}, {8, 1, 1}}));
    tileSizes.push_back(TileWorkgroupSizePair({{1, 32, 4}, {8, 1, 1}}));
  }
}

/// Launch configuration for Mali GPU configuration.
static LogicalResult getMaliSpecificConfig(
    linalg::BatchMatmulOp op, const spirv::TargetEnv &targetEnv,
    const SPIRVCodegenOptions &options, TileSizesListType &tileSizes,
    std::array<int64_t, 3> &workgroupSize,
    std::array<int64_t, 3> &numSubgroups) {
  if (targetEnv.getVendorID() != spirv::Vendor::ARM) return failure();

  SmallVector<ShapedType> inputTypes, outputTypes;
  std::tie(inputTypes, outputTypes) = getInputOutputTypes(op);

  ShapedType lhsType = inputTypes[0], rhsType = inputTypes[1];
  if (!lhsType || !rhsType || !lhsType.hasStaticShape() ||
      !rhsType.hasStaticShape())
    return failure();
  // Get a vector of best tile size ordered from best to worst.
  SmallVector<TileWorkgroupSizePair, 4> workgroupLevelTs;
  int64_t dstSize =
      lhsType.getDimSize(0) * lhsType.getDimSize(1) * rhsType.getDimSize(2);
  getMaliBestMatMulTileSizes(lhsType.getElementType(), workgroupLevelTs,
                             dstSize);
  for (TileWorkgroupSizePair pair : workgroupLevelTs) {
    if (lhsType.getDimSize(1) % pair.tileSize[0] != 0 ||
        rhsType.getDimSize(2) % pair.tileSize[1] != 0 ||
        lhsType.getDimSize(2) % pair.tileSize[2] != 0) {
      continue;
    }

    workgroupSize = pair.workgroupSize;
    SmallVector<int64_t, 4> batchTs;
    batchTs.append({1, pair.tileSize[0], pair.tileSize[1], pair.tileSize[2]});
    tileSizes.emplace_back(batchTs);
    // No tiling at the subgroup level since this target doesn't use subgroup op
    // or shared memory.
    tileSizes.emplace_back();
    SmallVector<int64_t, 4> invocationLevelTs = {
        batchTs[0], batchTs[1] / workgroupSize[1],
        batchTs[2] / workgroupSize[0], batchTs[3]};
    tileSizes.emplace_back(invocationLevelTs);
    return success();
  }
  return failure();
}

/// Launch config for `linalg.batchmatmul`.
template <>
LogicalResult getOpLaunchConfig(linalg::BatchMatmulOp op,
                                const spirv::TargetEnv &targetEnv,
                                const SPIRVCodegenOptions &options,
                                TileSizesListType &tileSizes,
                                LaunchConfigInfo &config) {
  if (succeeded(getMaliSpecificConfig(op, targetEnv, options, tileSizes,
                                      config.workgroupSize,
                                      config.numSubgroups))) {
    config.vectorize = true;
    return success();
  }
  unsigned maxWorkgroupSize = targetEnv.getResourceLimits()
                                  .max_compute_workgroup_invocations()
                                  .getInt();
  std::tie(config.workgroupSize[0], config.workgroupSize[1]) =
      distributeProcs2D(maxWorkgroupSize);
  config.workgroupSize[2] = 1;
  // This is just being hard-wired for now to be minimal viable, but this can be
  // decided better when we have better estimates of device charecteristics.
  const int64_t nRowsPerWorkitem = 1;
  const int64_t nColsPerWorkitem = 1;
  const int64_t nBatchesPerWorkitem = 1;
  int64_t tileSizeK = 0;
  if (options.useWorkgroupMemory) {
    // This number should be decided based on the amount of shared memory
    // available (maybe). For now, just hard-wire it.
    tileSizeK = 32;
  }
  SmallVector<int64_t, 4> workgroupLevel = {
      nBatchesPerWorkitem, nRowsPerWorkitem * config.workgroupSize[1],
      nColsPerWorkitem * config.workgroupSize[0], tileSizeK};
  SmallVector<int64_t, 4> invocationLevel = {
      nBatchesPerWorkitem, nRowsPerWorkitem, nColsPerWorkitem, 0};

  tileSizes.emplace_back(std::move(workgroupLevel));
  tileSizes.emplace_back();  // subgroup level
  tileSizes.emplace_back(std::move(invocationLevel));
  return success();
}

/// Returns the size of the co-operative matrix multiply operations on the
/// device.
static Optional<SmallVector<int64_t, 4>> getCooperativeMatmulSubgroupSize(
    spirv::ResourceLimitsAttr resourceLimits, Type lhsType, Type rhsType,
    Type initType, Type resultType) {
  for (auto coopMatmulProperties :
       resourceLimits.cooperative_matrix_properties_nv()
           .getAsRange<spirv::CooperativeMatrixPropertiesNVAttr>()) {
    if (coopMatmulProperties.a_type().getValue() == lhsType &&
        coopMatmulProperties.b_type().getValue() == rhsType &&
        coopMatmulProperties.c_type().getValue() == initType &&
        coopMatmulProperties.result_type().getValue() == resultType &&
        coopMatmulProperties.scope().getValue() == spirv::Scope::Subgroup) {
      return SmallVector<int64_t, 4>{
          coopMatmulProperties.m_size().getValue().getSExtValue(),
          coopMatmulProperties.n_size().getValue().getSExtValue(),
          coopMatmulProperties.k_size().getValue().getSExtValue()};
    }
  }
  return llvm::None;
}

/// Launch configuration for using spv.CooperativeMatrixMulAddNV
/// operations. Needs two levels of tiling.
static LogicalResult getConfigForCooperativeMatmul(
    linalg::MatmulOp op, const spirv::TargetEnv &targetEnv,
    const SPIRVCodegenOptions &options, TileSizesListType &tileSizes,
    std::array<int64_t, 3> &workgroupSize,
    std::array<int64_t, 3> &numSubgroups) {
  if (!targetEnv.allows(spirv::Capability::CooperativeMatrixNV) ||
      !targetEnv.allows(spirv::Extension::SPV_NV_cooperative_matrix))
    return failure();

  SmallVector<ShapedType> inputTypes, outputTypes;
  std::tie(inputTypes, outputTypes) = getInputOutputTypes(op);

  ShapedType lhsType = inputTypes[0], rhsType = inputTypes[1];
  ShapedType outputType = outputTypes[0];

  auto resourceLimits = targetEnv.getResourceLimits();
  Optional<SmallVector<int64_t, 4>> coopMatmulSize =
      getCooperativeMatmulSubgroupSize(
          resourceLimits, lhsType.getElementType(), rhsType.getElementType(),
          outputType.getElementType(), outputType.getElementType());
  if (!coopMatmulSize) return failure();

  // Check that the matmul sizes are a multiple of the tilesize.
  auto isMultipleOf = [](int64_t s, int64_t ts) {
    return !ShapedType::isDynamic(s) && (s % ts) == 0;
  };

  ArrayRef<int64_t> lhsShape = lhsType.getShape();
  ArrayRef<int64_t> rhsShape = rhsType.getShape();
  if (!isMultipleOf(lhsShape[0], (*coopMatmulSize)[0]) ||
      !isMultipleOf(rhsShape[1], (*coopMatmulSize)[1]) ||
      !isMultipleOf(lhsShape[1], (*coopMatmulSize)[2]) ||
      !isMultipleOf(rhsShape[0], (*coopMatmulSize)[2]))
    return failure();

  if (options.useWorkgroupMemory) {
    numSubgroups[0] = 2;
    numSubgroups[1] = 2;
  } else {
    numSubgroups[0] = 1;
    numSubgroups[1] = 1;
  }
  numSubgroups[2] = 1;

  // For now this is being hard-wired to be {4, 4, 2}. This can actually be set
  // to whatever, but ultimately depends on register pressure.
  const int64_t numVecMatmulPerSubgroupX = 4;
  const int64_t numVecMatmulPerSubgroupY = 4;
  const int64_t numVecMatmulPerSubgroupK = 2;
  SmallVector<int64_t, 4> ts = {
      numVecMatmulPerSubgroupY * (*coopMatmulSize)[0] * numSubgroups[1],
      numVecMatmulPerSubgroupX * (*coopMatmulSize)[1] * numSubgroups[0],
      numVecMatmulPerSubgroupK * (*coopMatmulSize)[2]};
  tileSizes.emplace_back(std::move(ts));

  int64_t subgroupSize =
      resourceLimits.subgroup_size().getValue().getSExtValue();
  workgroupSize[0] = numSubgroups[0] * numSubgroups[1] * subgroupSize;
  workgroupSize[1] = 1;
  workgroupSize[2] = 1;
  // Subgroup tile sizes
  SmallVector<int64_t, 4> subgroupTs = {
      numVecMatmulPerSubgroupY * (*coopMatmulSize)[0],
      numVecMatmulPerSubgroupX * (*coopMatmulSize)[1]};
  tileSizes.emplace_back(std::move(subgroupTs));
  return success();
}

/// Launch config for element-wise linalg.generic.
LogicalResult getGenericOpLaunchConfig(linalg::LinalgOp linalgOp,
                                       const spirv::TargetEnv &targetEnv,
                                       const SPIRVCodegenOptions &options,
                                       TileSizesListType &tileSizes,
                                       LaunchConfigInfo &config) {
  // Skip vectorization for non-minor identity inputs as it generates
  // transfer_read ops with permutation maps that we currently cannot lower.
  // TODO: Remove this restriction once the lowering of the permutation map is
  // supported in core.
  bool vectorize = !linalgOp.hasIndexSemantics() &&
                   llvm::all_of(linalgOp.getIndexingMaps(), [](AffineMap &map) {
                     return map.isMinorIdentity();
                   });
  // TODO(thomasraoux): Lowering of integers other than i32 may require
  // emulation. This is currently not supported for vector operation. Re-enable
  // this when the bug is fixed on SPIR-V lowering side.
  if (llvm::any_of(linalgOp->getOperands(), [](Value operand) {
        Type memrefType = operand.getType().cast<MemRefType>().getElementType();
        return !memrefType.isa<FloatType>() && !memrefType.isInteger(32);
      }))
    vectorize = false;
  int64_t subgroupSize =
      targetEnv.getResourceLimits().subgroup_size().getValue().getSExtValue();
  config.workgroupSize[0] = subgroupSize;
  config.workgroupSize[1] = 1;
  config.workgroupSize[2] = 1;
  SmallVector<ShapedType> inputTypes, outputTypes;
  std::tie(inputTypes, outputTypes) = getInputOutputTypes(linalgOp);
  ShapedType outputShape = outputTypes[0];

  SmallVector<int64_t, 4> candidateTileSizes;
  // When Vectororization is not enabled we skil the second level of tiling and
  // fall back to convertToGPU which will map one element to one thread. To
  // avoid a mismatch in the number of workgroup dispatched, we pick a tile size
  // to have one element per thread.
  // TODO: Remove this once we switch to linalg on tensor path.
  if (vectorize) {
    candidateTileSizes.append({4 * subgroupSize, 2 * subgroupSize});
  }

  candidateTileSizes.push_back(subgroupSize);
  // Use the first tile size that can divide the shape. If the shape is not
  // aligned on any of the tile sizes pick the smallest tile of one element per
  // thread.
  int64_t lowerTs = config.workgroupSize[0];
  for (int64_t size : candidateTileSizes) {
    if (outputShape.getShape().back() % size != 0) continue;
    lowerTs = size;
    break;
  }
  unsigned numLoops = getNumOuterParallelLoops(linalgOp);
  SmallVector<int64_t, 4> ts(numLoops, 1);
  ts.back() = lowerTs;
  tileSizes.emplace_back(ts);  // Workgroup level
  tileSizes.emplace_back();    // Subgroup level

  if (!vectorize || outputShape.getShape().back() % lowerTs != 0) {
    ts.back() = 1;
    tileSizes.emplace_back(ts);  // Thread level
    config.vectorize = false;
  } else {
    ts.back() = lowerTs / subgroupSize;
    tileSizes.emplace_back(ts);  // Thread level
    // Vectorize only if we are processing more than one element per thread.
    config.vectorize = vectorize && (ts.back() > 1);
  }
  return success();
}

#define GET_GENERIC_OP_LAUNCH_CONFIG(opType)                            \
  template <>                                                           \
  LogicalResult getOpLaunchConfig(                                      \
      opType op, const spirv::TargetEnv &targetEnv,                     \
      const SPIRVCodegenOptions &options, TileSizesListType &tileSizes, \
      LaunchConfigInfo &config) {                                       \
    return getGenericOpLaunchConfig(op, targetEnv, options, tileSizes,  \
                                    config);                            \
  }

GET_GENERIC_OP_LAUNCH_CONFIG(linalg::GenericOp)

#undef GET_GENERIC_OP_LAUNCH_CONFIG

/// Launch configuration for different known GPU configuration.
static LogicalResult getTargetSpecificConfig(
    linalg::MatmulOp op, const spirv::TargetEnv &targetEnv,
    const SPIRVCodegenOptions &options, TileSizesListType &tileSizes,
    std::array<int64_t, 3> &workgroupSize,
    std::array<int64_t, 3> &numSubgroups) {
  if (targetEnv.getVendorID() != spirv::Vendor::ARM) return failure();

  SmallVector<ShapedType> inputTypes, outputTypes;
  std::tie(inputTypes, outputTypes) = getInputOutputTypes(op);

  ShapedType lhsType = inputTypes[0], rhsType = inputTypes[1];
  // If the shape size is unknonw fall back to none vectorized path.
  if (!lhsType || !rhsType || !lhsType.hasStaticShape() ||
      !rhsType.hasStaticShape())
    return failure();

  // Pick ideal tile size based on the type.
  SmallVector<TileWorkgroupSizePair, 4> workgroupLevelTs;
  int64_t dstSize = lhsType.getDimSize(0) * rhsType.getDimSize(1);
  getMaliBestMatMulTileSizes(lhsType.getElementType(), workgroupLevelTs,
                             dstSize);
  for (TileWorkgroupSizePair pair : workgroupLevelTs) {
    if (lhsType.getDimSize(0) % pair.tileSize[0] != 0 ||
        rhsType.getDimSize(1) % pair.tileSize[1] != 0 ||
        lhsType.getDimSize(1) % pair.tileSize[2] != 0) {
      continue;
    }

    workgroupSize = pair.workgroupSize;
    SmallVector<int64_t, 4> matmulTS(pair.tileSize.begin(),
                                     pair.tileSize.end());
    tileSizes.emplace_back(matmulTS);
    // No tiling at the subgroup level since this target doesn't use subgroup op
    // or shared memory.
    tileSizes.emplace_back();
    SmallVector<int64_t, 4> invocationLevelTs = {matmulTS[0] / workgroupSize[1],
                                                 matmulTS[1] / workgroupSize[0],
                                                 matmulTS[2]};
    tileSizes.emplace_back(invocationLevelTs);
    return success();
  }
  return failure();
}

template <>
LogicalResult getOpLaunchConfig(linalg::MatmulOp op,
                                const spirv::TargetEnv &targetEnv,
                                const SPIRVCodegenOptions &options,
                                TileSizesListType &tileSizes,
                                LaunchConfigInfo &config) {
  if (succeeded(getConfigForCooperativeMatmul(op, targetEnv, options, tileSizes,
                                              config.workgroupSize,
                                              config.numSubgroups))) {
    config.vectorize = true;
    return success();
  }
  if (succeeded(getTargetSpecificConfig(op, targetEnv, options, tileSizes,
                                        config.workgroupSize,
                                        config.numSubgroups))) {
    config.vectorize = true;
    return success();
  }

  unsigned maxWorkgroupSize = targetEnv.getResourceLimits()
                                  .max_compute_workgroup_invocations()
                                  .getInt();
  std::tie(config.workgroupSize[0], config.workgroupSize[1]) =
      distributeProcs2D(maxWorkgroupSize);
  config.workgroupSize[2] = 1;
  const int nRowsPerWorkitem = 1;
  const int nColsPerWorkitem = 1;
  int64_t tileSizeK = 0;
  if (options.useWorkgroupMemory) {
    // TODO(#3131): This number should be decided based on the amount of shared
    // memory available (maybe). For now, just hard-wire it.
    tileSizeK = 32;
  }

  SmallVector<ShapedType> inputTypes;
  std::tie(inputTypes, std::ignore) = getInputOutputTypes(op);
  int64_t M = inputTypes[0].getShape()[0];
  int64_t N = inputTypes[1].getShape()[1];
  int64_t K = inputTypes[0].getShape()[1];

  SmallVector<int64_t, 4> workgroupLevel = {
      getMinIfShapeStatic(M, nRowsPerWorkitem * config.workgroupSize[1]),
      getMinIfShapeStatic(N, nColsPerWorkitem * config.workgroupSize[0]),
      getMinIfShapeStatic(K, tileSizeK)};
  SmallVector<int64_t, 4> invocationLevel = {1, 1, 0};

  tileSizes.emplace_back(std::move(workgroupLevel));
  tileSizes.emplace_back();  // subgroup level
  tileSizes.emplace_back(std::move(invocationLevel));
  return success();
}

static LogicalResult getMaliSpecificConfig(linalg::ConvInputNHWCFilterHWCFOp op,
                                           TileSizesListType &tileSizes,
                                           LaunchConfigInfo &config) {
  SmallVector<ShapedType> inputTypes, outputTypes;
  std::tie(inputTypes, outputTypes) = getInputOutputTypes(op);

  ShapedType inputType = inputTypes[0], outputType = outputTypes[0];
  if (!inputType || !outputType || !inputType.hasStaticShape() ||
      !outputType.hasStaticShape())
    return failure();

  bool isInputTilable =
      inputType.getDimSize(3) % 4 == 0 || inputType.getDimSize(3) < 4;
  if (!isInputTilable) return failure();

  // A list of preferred tile sizes and workgroup sizes. This is for Mali
  // G77 now and it's fairly ad-hoc. We need to have a better story for
  // incorporating such information.
  static const TileWorkgroupSizePair tileWorkgroupSizePairs[] = {
      {{4, 4, 16}, {4, 4, 1}},
      {{2, 2, 64}, {16, 1, 1}},
      {{4, 8, 8}, {2, 4, 2}},
      {{2, 2, 32}, {8, 2, 1}},
      {{1, 1, 32}, {8, 1, 1}}};

  for (const auto &pair : tileWorkgroupSizePairs) {
    const std::array<int64_t, 3> &tileSize = pair.tileSize;
    const std::array<int64_t, 3> &workgroupSize = pair.workgroupSize;

    auto outputShape = outputType.getShape();
    bool isOutputTilable = (outputShape[0] == 1) &&
                           (outputShape[1] % tileSize[0] == 0) &&
                           (outputShape[2] % tileSize[1] == 0) &&
                           (outputShape[3] % tileSize[2] == 0);
    if (!isOutputTilable) continue;

    SmallVector<int64_t, 4> workgroupLevel = {
        /*batch=*/0, /*output_height=*/tileSize[0],
        /*output_width=*/tileSize[1], /*output_channel=*/tileSize[2]};
    tileSizes.emplace_back(std::move(workgroupLevel));

    // No tiling at the subgroup level given that we don't use subgroup
    // level syncrhonization  or shared memory.
    tileSizes.emplace_back();

    SmallVector<int64_t, 4> invocationLevel = {
        /*batch=*/0, /*output_height=*/tileSize[0] / workgroupSize[2],
        /*output_width=*/tileSize[1] / workgroupSize[1],
        /*output_channel=*/tileSize[2] / workgroupSize[0]};
    tileSizes.emplace_back(invocationLevel);

    // Finally, for each invocation, we use tiling to generate loops to loop
    // over the filter's height (step 1), width (step 1), and input channel
    // (step 4) dimensions.
    SmallVector<int64_t, 4> fourthLevel = {0, 0, 0, 0, 1, 1, 4};
    tileSizes.emplace_back(fourthLevel);

    config.workgroupSize = workgroupSize;
    config.vectorize = true;

    return success();
  }

  return failure();
}

template <>
LogicalResult getOpLaunchConfig(linalg::ConvInputNHWCFilterHWCFOp op,
                                const spirv::TargetEnv &targetEnv,
                                const SPIRVCodegenOptions &options,
                                TileSizesListType &tileSizes,
                                LaunchConfigInfo &config) {
  if (targetEnv.getVendorID() == spirv::Vendor::ARM &&
      succeeded(getMaliSpecificConfig(op, tileSizes, config))) {
    return success();
  }

  return setDefaultTilingScheme(targetEnv, op, tileSizes, config.workgroupSize);
}

static LogicalResult getMaliSpecificConfig(
    linalg::DepthwiseConvInputNHWCFilterHWCOp op, TileSizesListType &tileSizes,
    LaunchConfigInfo &config) {
  SmallVector<ShapedType> inputTypes, outputTypes;
  std::tie(inputTypes, outputTypes) = getInputOutputTypes(op);

  ShapedType inputType = inputTypes[0], outputType = outputTypes[0];
  if (!inputType || !outputType || !inputType.hasStaticShape() ||
      !outputType.hasStaticShape())
    return failure();

  // A list of preferred tile sizes and workgroup sizes. This is for Mali
  // G77 now and it's fairly ad-hoc. We need to have a better story for
  // incorporating such information.
  static const TileWorkgroupSizePair tileWorkgroupSizePairs[] = {
      {{2, 2, 32}, {8, 2, 2}},
      {{1, 4, 16}, {4, 4, 1}},
      {{1, 1, 64}, {16, 1, 1}},
  };

  for (const auto &pair : tileWorkgroupSizePairs) {
    const std::array<int64_t, 3> &tileSize = pair.tileSize;
    const std::array<int64_t, 3> &workgroupSize = pair.workgroupSize;

    auto outputShape = outputType.getShape();
    bool isOutputTilable = outputShape[0] == 1 &&
                           (outputShape[1] % tileSize[0] == 0) &&
                           (outputShape[2] % tileSize[1] == 0) &&
                           (outputShape[3] % tileSize[2] == 0);
    if (!isOutputTilable) continue;

    SmallVector<int64_t, 4> workgroupLevel = {/*batch=*/0,
                                              /*output_height=*/tileSize[0],
                                              /*output_width=*/tileSize[1],
                                              /*output_channel=*/tileSize[2]};
    tileSizes.emplace_back(std::move(workgroupLevel));

    // No tiling at the subgroup level given that we don't use subgroup
    // level syncrhonization  or shared memory.
    tileSizes.emplace_back();

    SmallVector<int64_t, 4> invocationLevel = {
        /*batch=*/0, /*output_height=*/tileSize[0] / workgroupSize[2],
        /*output_width=*/tileSize[1] / workgroupSize[1],
        /*output_channel=*/tileSize[2] / workgroupSize[0]};
    tileSizes.emplace_back(invocationLevel);

    // Finally, for each invocation, we use tiling to generate loops to loop
    // over the filter's height (step 1) and width (step 1) dimensions.
    SmallVector<int64_t, 4> fourthLevel = {0, 0, 0, 0, 1, 1};
    tileSizes.emplace_back(fourthLevel);

    config.workgroupSize = workgroupSize;
    config.vectorize = true;

    return success();
  }
  return failure();
}

template <>
LogicalResult getOpLaunchConfig(linalg::DepthwiseConvInputNHWCFilterHWCOp op,
                                const spirv::TargetEnv &targetEnv,
                                const SPIRVCodegenOptions &options,
                                TileSizesListType &tileSizes,
                                LaunchConfigInfo &config) {
  if (targetEnv.getVendorID() == spirv::Vendor::ARM &&
      succeeded(getMaliSpecificConfig(op, tileSizes, config))) {
    return success();
  }

  return setDefaultTilingScheme(targetEnv, op, tileSizes, config.workgroupSize);
}

Optional<LaunchConfig> initGPULaunchConfig(
    MLIRContext *context, const linalg::LinalgDependenceGraph &dependenceGraph,
    const SPIRVCodegenOptions &options, ArrayRef<linalg::LinalgOp> linalgOps) {
  LaunchConfig launchConfig;
  if (!options.workgroupSize.empty()) {
    SmallVector<int64_t, 3> workgroupTileSizes(
        options.workgroupTileSizes.begin(), options.workgroupTileSizes.end());
    SmallVector<int64_t, 3> invocationTileSizes(
        options.invocationTileSizes.begin(), options.invocationTileSizes.end());
    for (linalg::LinalgOp linalgOp : linalgOps) {
      launchConfig.setTileSizes(linalgOp.getOperation(), workgroupTileSizes, 0);
      // Subgroup level.
      launchConfig.setTileSizes(linalgOp.getOperation(), {}, 1);
      // Invocation level.
      launchConfig.setTileSizes(linalgOp.getOperation(), invocationTileSizes,
                                2);
      launchConfig.setVectorize(true);
    }
    SmallVector<int64_t, 3> workgroupSize(options.workgroupSize.begin(),
                                          options.workgroupSize.end());
    launchConfig.setWorkgroupSize(workgroupSize);
  }

  if (linalgOps.empty()) return launchConfig;

  spirv::TargetEnv targetEnv(spirv::lookupTargetEnv(*linalgOps.begin()));

  Optional<linalg::LinalgOp> rootOperation = {};
  LaunchConfigInfo config;
#define DISPATCH(opName)                                                \
  if (auto op = dyn_cast<opName>(linalgOp.getOperation())) {            \
    rootOperation = linalgOp;                                           \
    if (launchConfig.hasTileSizes(linalgOp.getOperation())) break;      \
    TileSizesListType tileSizesInfo;                                    \
    if (failed(getOpLaunchConfig(op, targetEnv, options, tileSizesInfo, \
                                 config))) {                            \
      return llvm::None;                                                \
    }                                                                   \
    launchConfig.setTileSizes(op, tileSizesInfo);                       \
    break;                                                              \
  }

  for (linalg::LinalgOp linalgOp : linalgOps) {
    DISPATCH(linalg::BatchMatmulOp)
    DISPATCH(linalg::DepthwiseConvInputNHWCFilterHWCOp)
    DISPATCH(linalg::DepthwiseConvInputNHWCFilterHWCFOp)
    DISPATCH(linalg::ConvInputNWCFilterWCFOp)
    DISPATCH(linalg::ConvInputNHWCFilterHWCFOp)
    DISPATCH(linalg::ConvInputNDHWCFilterDHWCFOp)
    DISPATCH(linalg::MatmulOp)
    DISPATCH(linalg::PoolingNhwcMaxOp)
    DISPATCH(linalg::PoolingNhwcMinOp)
    DISPATCH(linalg::PoolingNhwcSumOp)
  }

  // Any generic operations found are made the root if no other op is the root
  if (!rootOperation) {
    for (linalg::LinalgOp linalgOp : reverse(linalgOps)) {
      size_t numLoops = getNumOuterParallelLoops(linalgOp);
      if (numLoops == 0 ||
          llvm::any_of(linalgOp.getIndexingMaps(), [](AffineMap &map) {
            return !map.isProjectedPermutation();
          })) {
        return llvm::None;
      }

      DISPATCH(linalg::GenericOp)
    }
  }

#undef DISPATCH

  if (!rootOperation) {
    return llvm::None;
  }

  launchConfig.setRootOperation(*rootOperation);
  if (options.workgroupSize.empty()) {
    launchConfig.setWorkgroupSize(config.workgroupSize);
    launchConfig.setVectorize(config.vectorize);
  }
  launchConfig.setNumSubgroups(config.numSubgroups);

  if (failed(propogateRootOperationLaunchConfig(launchConfig, *rootOperation,
                                                dependenceGraph)))
    return llvm::None;

  // TODO(ravishankarm): Verify that the set configurations is within the device
  // limits.
  return launchConfig;
}

template <typename OpTy>
static Optional<SmallVector<int64_t, 4>> getOpNativeVectorSize(OpTy op) {
  return llvm::None;
}

template <>
Optional<SmallVector<int64_t, 4>> getOpNativeVectorSize<vector::ContractionOp>(
    vector::ContractionOp op) {
  auto targetEnvAttr = spirv::lookupTargetEnv(op);
  auto targetEnv = spirv::TargetEnv(targetEnvAttr);
  if (targetEnv.allows(spirv::Capability::CooperativeMatrixNV) &&
      targetEnv.allows(spirv::Extension::SPV_NV_cooperative_matrix)) {
    return getCooperativeMatmulSubgroupSize(
        targetEnv.getResourceLimits(), op.getLhsType().getElementType(),
        op.getRhsType().getElementType(),
        op.getAccType().cast<VectorType>().getElementType(),
        op.getResultType().cast<VectorType>().getElementType());
  } else {
    unsigned lastParalleldim = 0;
    for (auto it : llvm::enumerate(op.iterator_types())) {
      if (isParallelIterator(it.value())) lastParalleldim = it.index();
    }
    SmallVector<int64_t, 4> nativeSize(op.iterator_types().size(), 1);
    nativeSize[lastParalleldim] = 4;
    // Map to vec4 fma operations.
    return nativeSize;
  }
}

template <>
Optional<SmallVector<int64_t, 4>> getOpNativeVectorSize<vector::FMAOp>(
    vector::FMAOp op) {
  SmallVector<int64_t, 4> size(op.getType().getRank(), 1);
  size.back() = 4;
  return size;
}

template <>
Optional<SmallVector<int64_t, 4>> getOpNativeVectorSize<vector::TransferReadOp>(
    vector::TransferReadOp op) {
  auto targetEnv = spirv::TargetEnv(spirv::lookupTargetEnv(op));
  if (targetEnv.allows(spirv::Capability::CooperativeMatrixNV) &&
      targetEnv.allows(spirv::Extension::SPV_NV_cooperative_matrix)) {
    // Unroll cooperative martrix load based on the size of the contract.
    VectorType dstVec;
    for (Operation *users : op->getUsers()) {
      auto extract = dyn_cast<vector::ExtractStridedSliceOp>(users);
      if (!extract) return llvm::None;
      auto vecType = extract.getResult().getType().cast<VectorType>();
      if (dstVec && dstVec != vecType) return llvm::None;
      dstVec = vecType;
    }
    return SmallVector<int64_t, 4>(dstVec.getShape().begin(),
                                   dstVec.getShape().end());
  }

  // Map to load4.
  auto rank = op.getVectorType().getRank();
  SmallVector<int64_t, 4> nativeSize(rank, 1);
  // Load 4 elements on the most inner dimension.
  for (auto dim : llvm::enumerate(op.permutation_map().getResults())) {
    if (auto dimExpr = dim.value().dyn_cast<AffineDimExpr>()) {
      if (dimExpr.getPosition() == op.permutation_map().getNumDims() - 1)
        nativeSize[dim.index()] = 4;
    }
  }
  return nativeSize;
}

template <>
Optional<SmallVector<int64_t, 4>>
getOpNativeVectorSize<vector::TransferWriteOp>(vector::TransferWriteOp op) {
  auto targetEnv = spirv::TargetEnv(spirv::lookupTargetEnv(op));
  if (targetEnv.allows(spirv::Capability::CooperativeMatrixNV) &&
      targetEnv.allows(spirv::Extension::SPV_NV_cooperative_matrix)) {
    // Unroll cooperative martrix store based on the size of the contract.
    auto insert = op.vector().getDefiningOp<vector::InsertStridedSliceOp>();
    if (!insert) return llvm::None;
    ArrayRef<int64_t> shape = insert.getSourceVectorType().getShape();
    return SmallVector<int64_t, 4>(shape.begin(), shape.end());
  }

  // Map to store4.
  auto rank = op.getVectorType().getRank();
  SmallVector<int64_t, 4> nativeSize(rank, 1);
  // Store 4 elements on the most inner dimension.
  for (auto dim : llvm::enumerate(op.permutation_map().getResults())) {
    if (auto dimExpr = dim.value().dyn_cast<AffineDimExpr>()) {
      if (dimExpr.getPosition() == op.permutation_map().getNumDims() - 1)
        nativeSize[dim.index()] = 4;
    }
  }
  return nativeSize;
}

Optional<SmallVector<int64_t, 4>> getSPIRVNativeVectorSize(Operation *op) {
#define DISPATCH(opname)                            \
  if (isa<opname>(op)) {                            \
    return getOpNativeVectorSize(cast<opname>(op)); \
  }

  DISPATCH(vector::ContractionOp)
  DISPATCH(vector::FMAOp)
  DISPATCH(vector::TransferReadOp)
  DISPATCH(vector::TransferWriteOp)

#undef DISPATCH

  if (OpTrait::hasElementwiseMappableTraits(op) && op->getNumResults() == 1) {
    if (auto vecType = op->getResultTypes()[0].dyn_cast<VectorType>()) {
      // Map elementwise ops to vec4.
      SmallVector<int64_t, 4> nativeSize(vecType.getRank() - 1, 1);
      nativeSize.push_back(4);
      return nativeSize;
    }
  }
  return llvm::None;
}

}  // namespace iree_compiler
}  // namespace mlir