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

#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"

#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenTypes.h"
#include "iree/compiler/Codegen/Utils/CodegenOptions.h"
#include "llvm/ADT/TypeSwitch.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Transform/IR/TransformOps.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/StorageUniquerSupport.h"
#include "mlir/Pass/PassRegistry.h"

#define GET_ATTRDEF_CLASSES
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp.inc"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenEnums.cpp.inc"

static const char kTranslationInfoAttrName[] = "translation_info";
static const char kCompilationInfoAttrName[] = "compilation_info";
static const char kRootOpInfoAttrName[] = "root_op";
static const char kUKernelDescriptorName[] = "iree_codegen.ukernel";

namespace mlir::iree_compiler {

/// Returns an `ArrayAttr` where each element is an `IntegerAttr` of 64-bit
/// integer type whose values is obtained from `values`.
static ArrayAttr getIndexArrayAttr(MLIRContext *context,
                                   ArrayRef<int64_t> values) {
  return ArrayAttr::get(
      context, llvm::map_to_vector(values, [&](int64_t value) -> Attribute {
        return IntegerAttr::get(IndexType::get(context), APInt(64, value));
      }));
}

bool shouldSetTunerAttributes() {
  return CodegenOptions::setTunerAttributes ||
         CodegenOptions::verifyPipelineConstraints ||
         CodegenOptions::emitPipelineConstraints;
}

bool shouldEmitPipelineConstraints() {
  return CodegenOptions::verifyPipelineConstraints ||
         CodegenOptions::emitPipelineConstraints;
}

bool shouldVerifyPipelineConstraints() {
  return CodegenOptions::verifyPipelineConstraints;
}

bool shouldPreservePipelineConstraints() {
  return CodegenOptions::emitPipelineConstraints;
}

} // namespace mlir::iree_compiler

namespace mlir::iree_compiler::IREE::Codegen {

//===----------------------------------------------------------------------===//
// iree_codegen.export_config
//===----------------------------------------------------------------------===//

LogicalResult
ExportConfigAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                         ArrayRef<int64_t> workgroupSize) {
  if (workgroupSize.size() > 3) {
    return emitError() << "expected workgroup size to have atmost 3 entries";
  }
  return success();
}

ArrayAttr ExportConfigAttr::getWorkgroupSizeIndexArray() {
  return getIndexArrayAttr(getContext(), getWorkgroupSize());
}

//===----------------------------------------------------------------------===//
// iree_codegen.vmvx_pipeline
//===----------------------------------------------------------------------===//

static VMVXPipelineBuilder &getVMVXPipelineBuilderStorage() {
  static VMVXPipelineBuilder builder = nullptr;
  return builder;
}

void registerVMVXPipelineBuilder(VMVXPipelineBuilder builder) {
  // Expected to be called exactly once during global init, so thread safety is
  // not a concern.
  [[maybe_unused]] static bool registered = false;
  assert(!registered && "VMVX pipeline builder registered more than once");
  registered = true;
  getVMVXPipelineBuilderStorage() = builder;
}

LogicalResult
VMVXPipelineAttr::buildPipeline(OpPassManager &pm,
                                const CodegenPipelineOptions *options) const {
  VMVXPipelineBuilder builder = getVMVXPipelineBuilderStorage();
  if (!builder) {
    return emitError(UnknownLoc::get(getContext()))
           << "no VMVX pipeline builder registered; ensure "
              "registerCodegenVMVXPasses() was called";
  }
  return builder(*this, pm, options);
}

//===----------------------------------------------------------------------===//
// iree_codegen.transform_dialect_codegen
//===----------------------------------------------------------------------===//

LogicalResult TransformDialectCodegenPipelineAttr::buildPipeline(
    OpPassManager &, const CodegenPipelineOptions *) const {
  return emitError(UnknownLoc::get(getContext()))
         << "transform dialect codegen pipeline should be consumed by "
            "LowerExecutableUsingTransformDialect and replaced with "
            "#iree_codegen.no_pipeline before executable target lowering";
}

//===----------------------------------------------------------------------===//
// iree_codegen.no_pipeline
//===----------------------------------------------------------------------===//

LogicalResult
NoPipelineAttr::buildPipeline(OpPassManager &,
                              const CodegenPipelineOptions *) const {
  return success();
}

//===----------------------------------------------------------------------===//
// iree_codegen.pass_pipeline
//===----------------------------------------------------------------------===//

LogicalResult
PassPipelineAttr::buildPipeline(OpPassManager &pm,
                                const CodegenPipelineOptions *options) const {
  // Textual pipelines ignore options.
  if (failed(parsePassPipeline(getPipeline(), pm))) {
    return failure();
  }
  return success();
}

LogicalResult
PassPipelineAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                         StringRef pipeline) {
  OpPassManager pm("builtin.module");
  if (failed(parsePassPipeline(pipeline, pm))) {
    return emitError() << "invalid pass pipeline specification: '" << pipeline
                       << "'";
  }
  return success();
}

//===----------------------------------------------------------------------===//
// iree_codegen.translation_info
//===----------------------------------------------------------------------===//

LogicalResult TranslationInfoAttr::verify(
    function_ref<InFlightDiagnostic()> emitError, Attribute passPipeline,
    SymbolRefAttr codegenSpec, ArrayRef<int64_t> workgroupSize,
    int64_t subgroupSize, DictionaryAttr configuration) {
  if (!passPipeline) {
    return emitError() << "missing pass pipeline specification";
  }
  if (!passPipeline.hasPromiseOrImplementsInterface<PipelineAttrInterface>()) {
    return emitError() << "pass pipeline must implement PipelineAttrInterface";
  }
  if (codegenSpec && !isa<TransformDialectCodegenPipelineAttr>(passPipeline)) {
    return emitError() << "transform dialect codegen spec requires transform "
                          "dialect codegen pipeline";
  }
  if (workgroupSize.size() > 3) {
    return emitError() << "workgroup size cannot have more than 3 entries";
  }
  if (llvm::any_of(workgroupSize, [](int64_t value) { return value <= 0; })) {
    return emitError() << "workgroup size value has to be greater than zero";
  }
  if (subgroupSize < 0) {
    return emitError() << "subgroup size value cannot be negative";
  }
  return success();
}

//===----------------------------------------------------------------------===//
// iree_codegen.lowering_config_level
//===----------------------------------------------------------------------===//

void LoweringConfigTilingLevelAttr::print(mlir::AsmPrinter &printer) const {
  auto tileInterchange = getInterchange();
  auto printTileSizes = [&] {
    printer << '[';
    if (getScalableFlags().empty()) {
      printer.printStrippedAttrOrType(getSizes());
    } else {
      llvm::interleaveComma(llvm::zip(getSizes(), getScalableFlags()), printer,
                            [&](auto pair) {
                              auto [tileSize, isScalable] = pair;
                              // Wrap scalable sizes in square brackets.
                              if (isScalable) {
                                printer << '[';
                              }
                              printer << tileSize;
                              if (isScalable) {
                                printer << ']';
                              }
                            });
    }
    printer << ']';
  };
  if (tileInterchange.empty()) {
    printTileSizes();
  } else {
    printer << "{sizes = ";
    printTileSizes();
    printer << ", interchange = [";
    printer.printStrippedAttrOrType(tileInterchange);
    printer << "]}";
  }
}

Attribute LoweringConfigTilingLevelAttr::parse(mlir::AsmParser &parser,
                                               mlir::Type) {
  auto loc = parser.getCurrentLocation();
  auto parseListOfSizes = [&](SmallVector<bool> *scalableFlags = nullptr,
                              bool prefixChecked =
                                  false) -> FailureOr<SmallVector<int64_t>> {
    if (!prefixChecked && parser.parseLSquare()) {
      return failure();
    }
    if (parser.parseOptionalRSquare().succeeded()) {
      // Empty list.
      return SmallVector<int64_t>();
    }
    SmallVector<int64_t> sizes;
    bool expectScalableSizes = scalableFlags != nullptr;
    auto listParse =
        parser.parseCommaSeparatedList(AsmParser::Delimiter::None, [&] {
          bool isScalable =
              expectScalableSizes && parser.parseOptionalLSquare().succeeded();
          int64_t size = 0;
          if (parser.parseInteger(size) ||
              (isScalable && parser.parseRSquare())) {
            return failure();
          }
          sizes.push_back(size);
          if (scalableFlags) {
            scalableFlags->push_back(isScalable);
          }
          return success();
        });
    if (failed(listParse) || parser.parseRSquare()) {
      return failure();
    }
    return sizes;
  };
  SmallVector<bool> scalableFlags;
  if (parser.parseOptionalLSquare().succeeded()) {
    // Case 1: Simple list of tile sizes, e.g.:
    // [0, [32], 16]
    auto tileSizes = parseListOfSizes(&scalableFlags, /*prefixChecked=*/true);
    if (failed(tileSizes)) {
      return {};
    }
    return parser.getChecked<LoweringConfigTilingLevelAttr>(
        loc, parser.getContext(), *tileSizes, ArrayRef<int64_t>{},
        scalableFlags);
  }
  // Case 2: sizes and interchange, e.g.:
  // {sizes = [0, [32], 16], interchange = [0, 1, 2]}
  if (parser.parseLBrace() || parser.parseKeyword("sizes") ||
      parser.parseEqual()) {
    return {};
  }
  auto tileSizes = parseListOfSizes(&scalableFlags);
  if (failed(tileSizes) || parser.parseComma() ||
      parser.parseKeyword("interchange") || parser.parseEqual()) {
    return {};
  }
  auto tileInterchange = parseListOfSizes();
  if (failed(tileInterchange) || parser.parseRBrace()) {
    return {};
  }
  return parser.getChecked<LoweringConfigTilingLevelAttr>(
      loc, parser.getContext(), *tileSizes, *tileInterchange, scalableFlags);
}

LogicalResult LoweringConfigTilingLevelAttr::verify(
    function_ref<InFlightDiagnostic()> emitError, ArrayRef<int64_t> tileSizes,
    ArrayRef<int64_t> tileInterchange, ArrayRef<bool> scalableFlags) {
  if (!scalableFlags.empty() && scalableFlags.size() != tileSizes.size()) {
    return emitError() << "scalable flags length does not match tile sizes";
  }
  return success();
}

//===----------------------------------------------------------------------===//
// iree_codegen.lowering_config
//===----------------------------------------------------------------------===//

LoweringConfigAttr
LoweringConfigAttr::get(MLIRContext *context, TileSizesListTypeRef tileSizes,
                        ScalableTileFlagsListTypeRef scalableTileFlags,
                        TileSizesListTypeRef tileInterchange) {
  SmallVector<LoweringConfigTilingLevelAttr> tilinglevels;
  for (auto [level, sizes] : llvm::enumerate(tileSizes)) {
    ArrayRef<int64_t> interchange = level < tileInterchange.size()
                                        ? tileInterchange[level]
                                        : ArrayRef<int64_t>{};
    ArrayRef<bool> scalableFlags = level < scalableTileFlags.size()
                                       ? scalableTileFlags[level]
                                       : ArrayRef<bool>{};
    tilinglevels.push_back(LoweringConfigTilingLevelAttr::get(
        context, sizes, interchange, scalableFlags));
  }
  return get(context,
             LoweringConfigTilingLevelsAttr::get(context, tilinglevels));
}

LoweringConfigAttr
LoweringConfigAttr::get(MLIRContext *context, TileSizesListTypeRef tileSizes,
                        TileSizesListTypeRef tileInterchange) {
  return get(context, tileSizes, {}, tileInterchange);
}

TileSizesListType LoweringConfigAttr::getTileSizeVals() const {
  TileSizesListType tileSizes;
  for (auto &level : getTilingLevels()) {
    tileSizes.push_back(SmallVector<int64_t>(level.getSizes()));
  }
  return tileSizes;
}

SmallVector<int64_t> LoweringConfigAttr::getTileSizeVals(unsigned level) const {
  auto levels = getTilingLevels();
  if (level >= levels.size()) {
    return {};
  }
  return SmallVector<int64_t>(levels[level].getSizes());
}

ScalableTileFlagsListType LoweringConfigAttr::getScalableTileFlagVals() {
  ScalableTileFlagsListType scalableFlags;
  for (auto &level : getTilingLevels()) {
    scalableFlags.push_back(SmallVector<bool>(level.getScalableFlags()));
  }
  return scalableFlags;
}

SmallVector<bool> LoweringConfigAttr::getScalableTileFlagVals(unsigned level) {
  auto levels = getTilingLevels();
  if (level >= levels.size()) {
    return {};
  }
  SmallVector<bool> scalableFlags(levels[level].getScalableFlags());
  // Extend the scalable flags with `false` to match the length of the sizes.
  scalableFlags.resize(levels[level].getSizes().size());
  return scalableFlags;
}

SmallVector<int64_t>
LoweringConfigAttr::getTileInterchangeVals(unsigned level) const {
  auto levels = getTilingLevels();
  if (level >= levels.size()) {
    return {};
  }
  return SmallVector<int64_t>(levels[level].getInterchange());
}

bool LoweringConfigAttr::isInterchangeEmpty() {
  return llvm::none_of(getTilingLevels(), [](auto level) {
    return !level.getInterchange().empty();
  });
}

SmallVector<int64_t> LoweringConfigAttr::getWorkgroupTileSizes() const {
  return getTileSizeVals(0);
}

SmallVector<int64_t> LoweringConfigAttr::getWorkgroupInterchange() const {
  return getTileInterchangeVals(0);
}

std::optional<unsigned> LoweringConfigAttr::getNumTilingLevels() const {
  return getTilingLevels().size();
}

SmallVector<int64_t>
LoweringConfigAttr::getStaticTilingLevelSizes(unsigned level,
                                              Operation *) const {
  return getTileSizeVals(level);
}

Attribute LoweringConfigAttr::getTilingLevelAttr(unsigned level) const {
  return getTilingLevels()[level];
}

SmallVector<OpFoldResult>
LoweringConfigAttr::getTilingLevelSizes(OpBuilder &builder, unsigned level,
                                        Operation *op) const {
  return llvm::map_to_vector(
      getStaticTilingLevelSizes(level, op),
      [&](int64_t t) -> OpFoldResult { return builder.getIndexAttr(t); });
}

bool LoweringConfigAttr::hasTilingLevel(unsigned level) const {
  return !getTileSizeVals(level).empty();
}

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

LogicalResult
LoweringConfigAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                           LoweringConfigTilingLevelsAttr levels) {
  if (!levels) {
    return emitError() << "missing lowering config levels";
  }
  return success();
}

//===----------------------------------------------------------------------===//
// iree.compilation_info
//===----------------------------------------------------------------------===//

LogicalResult
CompilationInfoAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                            LoweringConfigAttrInterface loweringConfig,
                            TranslationInfoAttr translationInfo) {
  if (!loweringConfig) {
    return emitError() << "missing lowering config";
  }
  if (auto defaultConfig = dyn_cast<LoweringConfigAttr>(loweringConfig)) {
    if (failed(LoweringConfigAttr::verify(emitError,
                                          defaultConfig.getTilingLevels()))) {
      return emitError() << "invalid lowering config: " << defaultConfig;
    }
  }
  if (!translationInfo) {
    return emitError() << "missing translation info";
  }
  if (failed(TranslationInfoAttr::verify(
          emitError, translationInfo.getPassPipeline(),
          translationInfo.getCodegenSpec(), translationInfo.getWorkgroupSize(),
          translationInfo.getSubgroupSize(),
          translationInfo.getConfiguration()))) {
    return emitError() << "invalid translation info: " << translationInfo;
  }
  return success();
}

//===----------------------------------------------------------------------===//
// iree_codegen.workgroup_mapping
//===----------------------------------------------------------------------===//

WorkgroupMappingAttr WorkgroupMappingAttr::get(MLIRContext *context,
                                               WorkgroupId id) {
  return WorkgroupMappingAttr::get(context, id, /*delinearizedDim=*/0);
}

LogicalResult
WorkgroupMappingAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                             WorkgroupId id, int64_t delinearizedDim) {
  if (delinearizedDim > 0 && id != WorkgroupId::IdZ) {
    return emitError() << "illegal to use `delinearizationDim` for "
                       << stringifyWorkgroupId(id);
  }
  return success();
}

WorkgroupMappingAttr
WorkgroupMappingAttr::getAttributeFromMappingId(MLIRContext *context,
                                                int64_t mappingId) {
  int64_t linearizedDim = std::max<int64_t>(mappingId - 2, 0);
  WorkgroupId id =
      symbolizeWorkgroupId(std::min<uint64_t>(mappingId, 2)).value();
  return WorkgroupMappingAttr::get(context, id, linearizedDim);
}

bool WorkgroupMappingAttr::operator<(const WorkgroupMappingAttr &rhs) const {
  return getMappingId() < rhs.getMappingId();
}

LogicalResult WorkgroupMappingAttr::verifyAttrList(MLIRContext *context,
                                                   Location loc,
                                                   ArrayRef<Attribute> attrs) {
  if (attrs.empty()) {
    return success();
  }
  SmallVector<IREE::Codegen::WorkgroupMappingAttr> mappingAttrs;
  llvm::SmallDenseSet<IREE::Codegen::WorkgroupMappingAttr, 4> attrSet;
  auto emitError = mlir::detail::getDefaultDiagnosticEmitFn(loc);
  for (auto attr : attrs) {
    auto typedAttr =
        dyn_cast_if_present<IREE::Codegen::WorkgroupMappingAttr>(attr);
    if (!typedAttr) {
      return emitError() << "expected all the mapping attribute to be of "
                            "`WorkgroupMappingAttr` type";
    }
    if (attrSet.contains(typedAttr)) {
      return emitError() << "illegal to repeat mapping specification";
    }
    attrSet.insert(typedAttr);
    mappingAttrs.push_back(typedAttr);
  }

  llvm::sort(mappingAttrs);

  // First element has to be `workgroup_mapping<x>`.
  if (mappingAttrs.front().getId() != IREE::Codegen::WorkgroupId::IdX) {
    return emitError() << "missing `workgroup_mapping<x>`";
  }
  if (mappingAttrs.size() == 1) {
    return success();
  }

  if (mappingAttrs[1].getId() != IREE::Codegen::WorkgroupId::IdY) {
    return emitError() << "missing `workgroup_mapping<y>`";
  }
  if (mappingAttrs.size() == 2) {
    return success();
  }

  auto mappingAttrsRef = ArrayRef<IREE::Codegen::WorkgroupMappingAttr>(
      mappingAttrs.begin(), mappingAttrs.end());
  for (auto [index, attr] : llvm::enumerate(mappingAttrsRef.drop_front(2))) {
    if (attr.getDelinearizedDim() != index) {
      return emitError() << "missing `workgroup_mapping<z:" << index;
    }
  }
  return success();
}

int64_t WorkgroupMappingAttr::getMappingId() const {
  return llvm::to_underlying(getId()) + getDelinearizedDim();
}

bool WorkgroupMappingAttr::isLinearMapping() const { return false; }

int64_t WorkgroupMappingAttr::getRelativeIndex() const {
  return getMappingId();
}

//===----------------------------------------------------------------------===//
// iree_codegen.local_mapping
//===----------------------------------------------------------------------===//

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

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

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

//===----------------------------------------------------------------------===//
// iree_codegen.simple_target
//===----------------------------------------------------------------------===//

LogicalResult
SimpleTargetAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                         ArrayRef<int64_t> maxWorkgroupCount) {
  if (maxWorkgroupCount.size() > 3) {
    return emitError()
           << "expected `max_workgroup_count` to have atmost 3 entries";
  }
  return success();
}

std::array<int64_t, 3> SimpleTargetAttr::getMaximumWorkgroupCount() const {
  std::array<int64_t, 3> maxWorkgroupCount = {
      ShapedType::kDynamic, ShapedType::kDynamic, ShapedType::kDynamic};
  ArrayRef<int64_t> givenMaxWorkgroupCount = getMaxWorkgroupCount();
  assert(givenMaxWorkgroupCount.size() <= 3 &&
         "expected `max_workgroup_count` to have atmost 3 entries");
  for (auto [index, value] : llvm::enumerate(givenMaxWorkgroupCount)) {
    maxWorkgroupCount[index] = value;
  }
  return maxWorkgroupCount;
}

//===---------------------------------------------------------------------===//
// iree_codegen.rotate_rows
//===---------------------------------------------------------------------===//

/// Given `r = |rotationInvariant|`, simplify additions of the following form:
///
/// %offset = arith.addi %0, c
///
/// Where `c` is a constant, to:
///
/// %offset = arith.addi %0, c % r
static OpFoldResult getMinimumConstantOffsetValue(OpBuilder &b, Location loc,
                                                  OpFoldResult offset,
                                                  int64_t rotationInvariant) {
  auto value = dyn_cast_if_present<Value>(offset);
  if (!value) {
    return offset;
  }

  auto add = value.getDefiningOp<arith::AddIOp>();
  if (!add) {
    return offset;
  }

  llvm::APInt constant;
  if (!matchPattern(add.getRhs(), m_ConstantInt(&constant))) {
    return offset;
  }

  int64_t constantOffset = constant.getSExtValue();
  int64_t baseMod = constantOffset % rotationInvariant;

  // Skip constructing the new apply if it's not needed (c < rotationInvariant).
  if (baseMod == constantOffset) {
    return offset;
  }

  Value modOffset = arith::ConstantIndexOp::create(b, loc, baseMod);
  // If the original add is nsw/nuw, then the new add must also be given we're
  // adding a smaller value.
  return arith::AddIOp::create(b, loc, add.getLhs(), modOffset,
                               add.getOverflowFlags())
      .getResult();
}

OpFoldResult RotateRowsAttr::swizzleOffset(OpBuilder &b, Location loc,
                                           OpFoldResult offset,
                                           Value src) const {
  // First normalize the producer of the offset if possible. Note that under
  // this swizzling scheme, every `row_width ^ 2 / access_width` elements yields
  // the same swizzling offset. This means we can modulo the id by this value
  // to get the simplest offset possible in case we are accessing values from
  // successive rows. This allows us to CSE the swizzling computation more
  // effectively.
  OpFoldResult id =
      getMinimumConstantOffsetValue(b, loc, offset, getSwizzlePeriod());

  // Number of elements per row.
  Value rowAlignmentVal = arith::ConstantIndexOp::create(b, loc, getRowWidth());
  // Number of contiguous groups of elements per row (swizzled together).
  Value rowAccessAlignmentVal =
      arith::ConstantIndexOp::create(b, loc, getRowWidth() / getAccessWidth());
  // Number of elements per group.
  Value accessWidthVal =
      arith::ConstantIndexOp::create(b, loc, getAccessWidth());

  // Strip the sub-access_width remainder to align the offset, since the
  // swizzle operates at access_width granularity.
  Value idVal = getValueOrCreateConstantIndexOp(b, loc, id);
  Value remainder = arith::RemUIOp::create(b, loc, idVal, accessWidthVal);
  Value alignedId = arith::SubIOp::create(b, loc, idVal, remainder);

  // i = row # = |offset| floordiv |Num elements per row|
  Value i = arith::DivUIOp::create(b, loc, alignedId, rowAlignmentVal);
  // jByte = element column # = |offset| % |Num elements per row|
  Value jElem = arith::RemUIOp::create(b, loc, alignedId, rowAlignmentVal);
  // j = group column # = jElem / |Num elements per group|
  Value j = arith::DivUIOp::create(b, loc, jElem, accessWidthVal);

  // New j = ((i + j) % |Num groups per row|) * |Num elements per group|
  Value sum = arith::AddIOp::create(b, loc, i, j);
  Value modByte = arith::RemUIOp::create(b, loc, sum, rowAccessAlignmentVal);
  Value mod = arith::MulIOp::create(b, loc, modByte, accessWidthVal);

  // Recompute the swizzled offset `(i * row width) + |new j|`
  Value mul = arith::MulIOp::create(b, loc, i, rowAlignmentVal);
  Value swizzledId = arith::AddIOp::create(b, loc, mod, mul);

  // |swizzledId| is the new offset modulo the swizzle invariant, meaning to
  // get the true swizzled offset take the difference between |swizzledId| and
  // |id| to get change in offset and add it to |offset|.
  //
  // This increases the chance of being able to CSE the offset calculation. When
  // multiple accesses to a memref only differ by a constant value (very common
  // when working with statically shaped memrefs like shared/scratch memory).
  //
  // The diff is computed from |alignedId| but applied to the original |offset|.
  // This implicitly restores the sub-access_width remainder:
  //   offset + (swizzle(alignedId) - alignedId)
  //   = (alignedId + rem) + (swizzle(alignedId) - alignedId)
  //   = swizzle(alignedId) + rem
  Value diff = arith::SubIOp::create(b, loc, swizzledId, alignedId);
  return arith::AddIOp::create(
             b, loc, getValueOrCreateConstantIndexOp(b, loc, offset), diff)
      .getResult();
}

int64_t RotateRowsAttr::getAccessElementCount() const {
  return getAccessWidth();
}

LogicalResult
RotateRowsAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                       int64_t rowWidth, int64_t accessWidth) {
  if (rowWidth % accessWidth != 0) {
    return emitError() << "expected access width to divide row width";
  }
  return success();
}

//===---------------------------------------------------------------------===//
// iree_codegen.symbolic_ukernel_provider
//===---------------------------------------------------------------------===//

FailureOr<Operation *>
SymbolicUKernelProviderAttr::getMLIRUKernel(StringRef name, DictionaryAttr,
                                            Operation *annotationSite) const {
  auto *symbolTableOp = SymbolTable::getNearestSymbolTable(annotationSite);
  SymbolTable symbolTable(symbolTableOp);
  return symbolTable.lookup(name);
}

//===---------------------------------------------------------------------===//
// iree_codegen.xor_shuffle
//===---------------------------------------------------------------------===//

/// Extract column index for XOR swizzling.
/// ((id%rowStride) / accessWidth)
static Value extractCol(OpBuilder &builder, Location loc, OpFoldResult id,
                        OpFoldResult rowAlignment, OpFoldResult accessWidth) {
  AffineExpr d0, s0, s1;
  bindDims(builder.getContext(), d0);
  bindSymbols(builder.getContext(), s0, s1);
  AffineExpr result = (d0 % s0).floorDiv(s1);
  return getValueOrCreateConstantIndexOp(
      builder, loc,
      affine::makeComposedFoldedAffineApply(builder, loc, result,
                                            {id, rowAlignment, accessWidth}));
}

/// Extract row index for XOR swizzling.
/// row = ((id/rowStride) / perPhase ) % rowAccessAlignment
static Value extractRow(OpBuilder &builder, Location loc, OpFoldResult id,
                        OpFoldResult rowStride, OpFoldResult perPhase,
                        OpFoldResult rowAccessAlignment) {
  AffineExpr d0, s0, s1, s2;
  bindDims(builder.getContext(), d0);
  bindSymbols(builder.getContext(), s0, s1, s2);
  AffineExpr result = (d0.floorDiv(s0).floorDiv(s1)) % s2;
  return getValueOrCreateConstantIndexOp(
      builder, loc,
      affine::makeComposedFoldedAffineApply(
          builder, loc, result, {id, rowStride, perPhase, rowAccessAlignment}));
}

/// Swizzle column on id.
/// new_id = id-id%rowAlignmentVal+colSwizzled*accessWidthVal
static Value updateCol(OpBuilder &builder, Location loc, OpFoldResult id,
                       Value colSwizzled, OpFoldResult rowAlignment,
                       OpFoldResult accessWidth) {
  AffineExpr d0, d1, s0, s1;
  bindDims(builder.getContext(), d0, d1);
  bindSymbols(builder.getContext(), s0, s1);
  AffineExpr result = d0 - d0 % s0 + d1 * s1;
  return getValueOrCreateConstantIndexOp(
      builder, loc,
      affine::makeComposedFoldedAffineApply(
          builder, loc, result, {id, colSwizzled, rowAlignment, accessWidth}));
}

OpFoldResult XORShuffleAttr::swizzleOffset(OpBuilder &b, Location loc,
                                           OpFoldResult offset,
                                           Value src) const {
  int64_t rowStride =
      getRowStride() != int64_t() ? getRowStride() : getRowWidth();
  int64_t perPhase = getPerPhase() != int64_t() ? getPerPhase() : 1;

  OpFoldResult id =
      getMinimumConstantOffsetValue(b, loc, offset, getSwizzlePeriod());

  // Strip the sub-access_width remainder to align the offset, since the
  // swizzle operates at access_width granularity.
  Value idVal = getValueOrCreateConstantIndexOp(b, loc, id);
  Value accessWidthVal =
      arith::ConstantIndexOp::create(b, loc, getAccessWidth());
  Value remainder = arith::RemUIOp::create(b, loc, idVal, accessWidthVal);
  Value alignedId = arith::SubIOp::create(b, loc, idVal, remainder);

  // Number of elements per row.
  Value rowAlignmentVal = arith::ConstantIndexOp::create(b, loc, getRowWidth());
  // Number of rows per phase.
  Value perPhaseVal = arith::ConstantIndexOp::create(b, loc, perPhase);
  // Buffer stride.
  Value rowStrideVal = arith::ConstantIndexOp::create(b, loc, rowStride);
  // Number of contiguous groups of elements per row (swizzled together).
  Value rowAccessAlignmentVal =
      arith::ConstantIndexOp::create(b, loc, getRowWidth() / getAccessWidth());

  Value colVal = extractCol(b, loc, alignedId, rowAlignmentVal, accessWidthVal);
  Value rowVal = extractRow(b, loc, alignedId, rowStrideVal, perPhaseVal,
                            rowAccessAlignmentVal);
  auto colSwizzled = arith::XOrIOp::create(b, loc, rowVal, colVal);

  // Update colSwizzled to alignedId
  Value swizzledIdVal = updateCol(b, loc, alignedId, colSwizzled,
                                  rowAlignmentVal, accessWidthVal);
  // The diff is computed from |alignedId| but applied to the original |offset|.
  // This implicitly restores the sub-access_width remainder:
  //   offset + (swizzle(alignedId) - alignedId)
  //   = (alignedId + rem) + (swizzle(alignedId) - alignedId)
  //   = swizzle(alignedId) + rem
  Value diff = arith::SubIOp::create(b, loc, swizzledIdVal, alignedId);
  return arith::AddIOp::create(
             b, loc, getValueOrCreateConstantIndexOp(b, loc, offset), diff)
      .getResult();
}

int64_t XORShuffleAttr::getAccessElementCount() const {
  return getAccessWidth();
}

LogicalResult
XORShuffleAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                       int64_t rowWidth, int64_t accessWidth, int64_t rowStride,
                       int64_t perPhase) {
  if (rowWidth % accessWidth != 0) {
    return emitError() << "expected access width to divide row width";
  }
  int64_t maxPhase = rowWidth / accessWidth;
  if (perPhase > maxPhase) {
    return emitError() << "per_phase must be smaller than max_phase";
  }
  if (rowStride % rowWidth != 0) {
    return emitError() << "expected row width to divide row stride";
  }

  return success();
}

//===----------------------------------------------------------------------===//
// Initialize attributes
//===----------------------------------------------------------------------===//

void IREECodegenDialect::initializeCodegenAttrs() {
  addAttributes<
#define GET_ATTRDEF_LIST
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp.inc" // IWYU pragma: keeep
      >();
}

} // namespace mlir::iree_compiler::IREE::Codegen

namespace mlir::iree_compiler {

//===----------------------------------------------------------------------===//
// Helpers for getting/setting iree_codegen.translation_info attribute on a
// FunctionOpInterface op.
// ===----------------------------------------------------------------------===//

IREE::Codegen::TranslationInfoAttr
getTranslationInfo(FunctionOpInterface funcOp) {
  return funcOp->getAttrOfType<IREE::Codegen::TranslationInfoAttr>(
      kTranslationInfoAttrName);
}

std::optional<SmallVector<int64_t>>
getWorkgroupSize(FunctionOpInterface funcOp) {
  IREE::Codegen::TranslationInfoAttr translationInfo =
      getTranslationInfo(funcOp);
  if (!translationInfo) {
    return std::nullopt;
  }
  return llvm::to_vector(translationInfo.getWorkgroupSize());
}

std::optional<int64_t> getSubgroupSize(FunctionOpInterface funcOp) {
  IREE::Codegen::TranslationInfoAttr translationInfo =
      getTranslationInfo(funcOp);
  if (!translationInfo) {
    return std::nullopt;
  }
  // The underlying storage sets 0 to optional scalar integer value. So if set
  // to 0, return as not set.
  if (translationInfo.getSubgroupSize() == int64_t()) {
    return std::nullopt;
  }
  return translationInfo.getSubgroupSize();
}

LogicalResult
setTranslationInfo(mlir::FunctionOpInterface entryPoint,
                   IREE::Codegen::TranslationInfoAttr translationInfo) {
  entryPoint->setAttr(kTranslationInfoAttrName, translationInfo);
  return success();
}

void eraseTranslationInfo(FunctionOpInterface funcOp) {
  funcOp->removeAttr(kTranslationInfoAttrName);
}

//===----------------------------------------------------------------------===//
// Helpers for getting/setting `iree_codegen.lowering_config` attribute on root
// operations.
// ===----------------------------------------------------------------------===//

SmallVector<int64_t> getTileSizes(Operation *op, unsigned level) {
  IREE::Codegen::LoweringConfigAttrInterface configAttr = getLoweringConfig(op);
  if (!configAttr) {
    return {};
  }
  return configAttr.getStaticTilingLevelSizes(level, op);
}
SmallVector<Value> getTileSizes(OpBuilder &b, Operation *op, unsigned level) {
  IREE::Codegen::LoweringConfigAttrInterface configAttr = getLoweringConfig(op);
  if (!configAttr) {
    return {};
  }
  return llvm::map_to_vector(configAttr.getTilingLevelSizes(b, level, op),
                             [&](OpFoldResult s) -> Value {
                               return getValueOrCreateConstantIndexOp(
                                   b, op->getLoc(), s);
                             });
}

void setLoweringConfig(Operation *op, Attribute config) {
  op->setAttr(kConfigAttrName, config);
}

void eraseLoweringConfig(Operation *op) { op->removeAttr(kConfigAttrName); }

//===----------------------------------------------------------------------===//
// Helpers for getting/setting `iree_codegen.compilation_info` attribute on root
// operations to override IREEs default compilation.
// ===----------------------------------------------------------------------===//

IREE::Codegen::CompilationInfoAttr getCompilationInfo(Operation *op) {
  return op->getAttrOfType<IREE::Codegen::CompilationInfoAttr>(
      kCompilationInfoAttrName);
}

void setCompilationInfo(Operation *op,
                        IREE::Codegen::CompilationInfoAttr config) {
  op->setAttr(kCompilationInfoAttrName, config);
}

void eraseCompilationInfo(Operation *op) {
  op->removeAttr(kCompilationInfoAttrName);
}

//===----------------------------------------------------------------------===//
// Helpers for setting attributes for tuner.
// ===----------------------------------------------------------------------===//

void setRootOpInfo(Operation *op, int64_t set) {
  op->setAttr(kRootOpInfoAttrName,
              IREE::Codegen::RootOpAttr::get(op->getContext(), set));
}

bool hasRootOpInfo(Operation *op) {
  return op->hasAttrOfType<IREE::Codegen::RootOpAttr>(kRootOpInfoAttrName);
}

IREE::Codegen::RootOpAttr getRootOpInfo(Operation *op) {
  return op->getAttrOfType<IREE::Codegen::RootOpAttr>(kRootOpInfoAttrName);
}

//===----------------------------------------------------------------------===//
// Helpers for getting/setting attributes related to ukernels.
//===----------------------------------------------------------------------===//

IREE::Codegen::UKernelProviderInterface
getUKernelProviderFromTarget(DictionaryAttr dict) {
  if (!dict) {
    return {};
  }
  return dict.getAs<IREE::Codegen::UKernelProviderInterface>(
      kUKernelProviderName);
}

IREE::Codegen::UKernelDescriptorAttr getUKernelDescriptor(Operation *op) {
  return op->getAttrOfType<IREE::Codegen::UKernelDescriptorAttr>(
      kUKernelDescriptorName);
}

void setUKernelDescriptor(Operation *op,
                          IREE::Codegen::UKernelDescriptorAttr descriptor) {
  op->setAttr(kUKernelDescriptorName, descriptor);
}

//===----------------------------------------------------------------------===//
// iree_codegen.workgroup_scope
//===----------------------------------------------------------------------===//

// Custom parser/printer to make the <> optional. It will either parse/print as
//   #iree_codegen.workgroup_scope<linearize>
// or without the <>.
//   #iree_codegen.workgroup_scope
Attribute IREE::Codegen::WorkgroupScopeAttr::parse(AsmParser &parser, Type) {
  bool linearize = false;
  if (succeeded(parser.parseOptionalLess())) {
    if (failed(parser.parseKeyword("linearize")) ||
        failed(parser.parseGreater())) {
      return {};
    }
    linearize = true;
  }
  return WorkgroupScopeAttr::get(parser.getContext(), linearize);
}

void IREE::Codegen::WorkgroupScopeAttr::print(AsmPrinter &printer) const {
  if (getLinearize()) {
    printer << "<" << "linearize" << ">";
  }
}

} // namespace mlir::iree_compiler
