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

//===- AdrenoConfig.h - Adreno CodeGen Configurations ---------------------===//
//
// This file contains CodeGen configurations for Adreno GPUs.
//
//===----------------------------------------------------------------------===//

#include <array>

#include "iree/compiler/Codegen/SPIRV/KernelConfig.h"
#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"

namespace mlir::iree_compiler::detail {

static LogicalResult setAdrenoMatmulConfig(linalg::LinalgOp op,
                                           IREE::GPU::TargetAttr target) {
  const int subgroupSize = target.getPreferredSubgroupSize();
  const std::array<int64_t, 2> workgroupXY = {subgroupSize / 2, 2};
  std::array<int64_t, 3> threadMNK;
  auto inputType = cast<ShapedType>(op.getDpsInputOperand(0)->get().getType());
  if (IREE::Util::getTypeBitWidth(inputType.getElementType()) == 16) {
    threadMNK = {16, 8, 8};
  } else {
    threadMNK = {16, 4, 4};
  }
  return setMatmulOpConfig(target, op, workgroupXY, threadMNK);
}

//===----------------------------------------------------------------------===//
// Entry Point
//===----------------------------------------------------------------------===//

LogicalResult setAdrenoCodeGenConfig(IREE::GPU::TargetAttr target,
                                     Operation *rootOp) {
  int subgroupSize = target.getPreferredSubgroupSize();

  if (!isa<linalg::LinalgOp>(rootOp)) {
    return failure();
  }

  auto linalgOp = cast<linalg::LinalgOp>(rootOp);
  if (isMatmulOrBatchMatmul(linalgOp)) {
    return setAdrenoMatmulConfig(linalgOp, target);
  }

  if (auto convOp = dyn_cast<linalg::ConvolutionOpInterface>(rootOp)) {
    // Use the result type in case of larger bitwidth for accumulators.
    auto type = cast<ShapedType>(convOp->getResult(0).getType());
    const int bitwidth = type.getElementTypeBitWidth();
    if (bitwidth > 32) {
      return failure();
    }
    const int multiplier = 32 / bitwidth;

    auto convDimsOrFailure = linalg::inferConvolutionDims(linalgOp);
    if (failed(convDimsOrFailure)) {
      return failure();
    }
    const int bestTilingFactor =
        (convDimsOrFailure->depth.empty() ? 32 : 16) * multiplier;
    return setConvOpConfig(cast<linalg::LinalgOp>(rootOp), subgroupSize,
                           bestTilingFactor);
  }

  return failure();
}

} // namespace mlir::iree_compiler::detail
