<commit_before>// 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/LLVMGPU/KernelConfig.h"

#include "iree/compiler/Codegen/Utils/Utils.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"

using namespace mlir;
using namespace mlir::iree_compiler;

static constexpr unsigned cudaWarpSize = 32;

static void setConfig(TileSizesListType tileSizes, Operation* op) {
  IREE::HAL::LoweringConfig config =
      buildConfigAttr(tileSizes, ArrayRef<int64_t>{}, op->getContext());
  setLoweringConfig(op, config);
}

/// Sets the translation info on the `hal.executable.entry_point` op
/// corresponding to the `entryPointFn`. Returns failure if a translation info
/// is already set on the entry point op and is incompatible with what is being
/// set.
static LogicalResult setTranslationInfo(
    FuncOp entryPointFn, IREE::HAL::DispatchLoweringPassPipeline passPipeline,
    ArrayRef<int64_t> workloadPerWorkgroup) {
  auto entryPointOp = getEntryPoint(entryPointFn);
  auto translationInfo = buildTranslationInfo(
      passPipeline, workloadPerWorkgroup, entryPointFn.getContext());
  return setTranslationInfo(entryPointOp, translationInfo);
}

static LogicalResult setRootConfig(FuncOp entryPoint, linalg::GenericOp op) {
  IREE::HAL::DispatchLoweringPassPipeline passPipeline =
      IREE::HAL::DispatchLoweringPassPipeline::LLVMGPUVectorize;
  std::array<int64_t, 3> workgroupSize = {1, 1, 1};
  TileSizesListType tileSizes;
  size_t numLoops = getNumOuterParallelLoops(op);
  if (numLoops == 0) {
    // Pure reduction, we serialize the operation on a single thread.
    // TODO: Use atomic to allow distributing reduction loops.
    tileSizes.push_back({});
    setConfig(tileSizes, op);
    return setTranslationInfo(entryPoint, passPipeline, workgroupSize);
  }

  workgroupSize = {cudaWarpSize, 1, 1};
  // Pick a fixed tile size independent of the original shape.
  // TODO(thomasraoux): Currently the original shape information is lost during
  // tiling at the flow level. We need way to access it to be able to make a
  // better choice of tile size.
  int64_t lowerTs = 4 * cudaWarpSize;
  SmallVector<int64_t, 4> ts;
  ts.resize(numLoops, 1);
  ts.back() = lowerTs;
  tileSizes.push_back(ts);  // Workgroup level.
  tileSizes.push_back({});  // Subgroup level.
  ts.back() = lowerTs / cudaWarpSize;
  tileSizes.push_back(ts);  // Thread level.
  setConfig(tileSizes, op);
  return setTranslationInfo(entryPoint, passPipeline, workgroupSize);
}

static LogicalResult setContractConfig(FuncOp entryPoint, linalg::LinalgOp op) {
  IREE::HAL::DispatchLoweringPassPipeline passPipeline =
      IREE::HAL::DispatchLoweringPassPipeline::LLVMGPUVectorize;
  TileSizesListType tileSizes;
  const int64_t numWarp = 2;
  SmallVector<int64_t, 3> workgroupSize = {numWarp * cudaWarpSize, 1, 1};
  // Currently just a basic tile size to enable tiling and vectorization.
  // TODO: pick a more efficient tile size and tile at subgroup level.
  SmallVector<int64_t, 4> ts;
  // Tile all the higher parallel dimension with a size of 1 and the 2 most
  // inner dimension with the tileX/tileY size.
  ts.append(op.getNumParallelLoops() - 2, 1);
  int64_t tileX = 2;
  int64_t tileY = 256;
  ts.append({tileX, tileY});
  // Tile all the reduction dimension with a size of 4.
  ts.append(op.getNumReductionLoops(), 4);
  tileSizes.push_back(ts);  // Workgroup level.
  tileSizes.push_back({});  // Subgroup level.
  // At the thread level only tile parallel loops.
  SmallVector<int64_t, 4> invocationLevelTs(op.getNumParallelLoops() - 2, 1);
  invocationLevelTs.append(<commit_msg>[CUDA] Add basic logic to pick the tile size based on contractioon shape (#6523)<commit_after>// 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/LLVMGPU/KernelConfig.h"

#include "iree/compiler/Codegen/Utils/Utils.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"

using namespace mlir;
using namespace mlir::iree_compiler;

static constexpr unsigned cudaWarpSize = 32;

static void setConfig(TileSizesListType tileSizes, Operation* op) {
  IREE::HAL::LoweringConfig config =
      buildConfigAttr(tileSizes, ArrayRef<int64_t>{}, op->getContext());
  setLoweringConfig(op, config);
}

/// Sets the translation info on the `hal.executable.entry_point` op
/// corresponding to the `entryPointFn`. Returns failure if a translation info
/// is already set on the entry point op and is incompatible with what is being
/// set.
static LogicalResult setTranslationInfo(
    FuncOp entryPointFn, IREE::HAL::DispatchLoweringPassPipeline passPipeline,
    ArrayRef<int64_t> workloadPerWorkgroup) {
  auto entryPointOp = getEntryPoint(entryPointFn);
  auto translationInfo = buildTranslationInfo(
      passPipeline, workloadPerWorkgroup, entryPointFn.getContext());
  return setTranslationInfo(entryPointOp, translationInfo);
}

static LogicalResult setRootConfig(FuncOp entryPoint, linalg::GenericOp op) {
  IREE::HAL::DispatchLoweringPassPipeline passPipeline =
      IREE::HAL::DispatchLoweringPassPipeline::LLVMGPUVectorize;
  std::array<int64_t, 3> workgroupSize = {1, 1, 1};
  TileSizesListType tileSizes;
  size_t numLoops = getNumOuterParallelLoops(op);
  if (numLoops == 0) {
    // Pure reduction, we serialize the operation on a single thread.
    // TODO: Use atomic to allow distributing reduction loops.
    tileSizes.push_back({});
    setConfig(tileSizes, op);
    return setTranslationInfo(entryPoint, passPipeline, workgroupSize);
  }

  workgroupSize = {cudaWarpSize, 1, 1};
  // Pick a fixed tile size independent of the original shape.
  // TODO(thomasraoux): Currently the original shape information is lost during
  // tiling at the flow level. We need way to access it to be able to make a
  // better choice of tile size.
  int64_t lowerTs = 4 * cudaWarpSize;
  SmallVector<int64_t, 4> ts;
  ts.resize(numLoops, 1);
  ts.back() = lowerTs;
  tileSizes.push_back(ts);  // Workgroup level.
  tileSizes.push_back({});  // Subgroup level.
  ts.back() = lowerTs / cudaWarpSize;
  tileSizes.push_back(ts);  // Thread level.
  setConfig(tileSizes, op);
  return setTranslationInfo(entryPoint, passPipeline, workgroupSize);
}

namespace {
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

/// Return the best combination of tile size and wg size. It will then used to
/// pick the best size aligned with the shape dimension.
static void getMatmulConfig(SmallVectorImpl<TileWorkgroupSizePair>& tileSizes) {
  tileSizes.push_back(TileWorkgroupSizePair({{16, 128, 4}, {32, 1, 1}}));
  tileSizes.push_back(TileWorkgroupSizePair({{8, 128, 4}, {32, 1, 1}}));
  tileSizes.push_back(TileWorkgroupSizePair({{16, 64, 4}, {16, 2, 1}}));
}

static LogicalResult setContractConfig(FuncOp entryPoint, linalg::LinalgOp op) {
  IREE::HAL::DispatchLoweringPassPipeline passPipeline =
      IREE::HAL::DispatchLoweringPassPipeline::LLVMGPUVectorize;
  TileSizesListType tileSizes;
  // Infer the MxN size of the matmul based on operands and indexing maps.
  auto lhsShape = getUntiledShape(op.getInputOperand(0)->get());
  auto rhsShape = getUntiledShape(op.getInputOperand(1)->get());
  int64_t sizeM = -1;
  int64_t sizeN = -1;
  auto outputMap = op.getTiedIndexingMap(op.getOutputOperand(0));
  for (unsigned i = 0; i < lhsShape.size(); i++) {
    if (op.getTiedIndexingMap(op.getInputOperand(0)).getDimPosition(i) ==
        outputMap.getDimPosition(outputMap.getNumResults() - 2)) {
      sizeM = lhsShape[i];
      break;
    }
  }
  for (unsigned i = 0; i < rhsShape.size(); i++) {
    if (op.getTiedIndexingMap(op.getInputOperand(1)).getDimPosition(i) ==
        outputMap.getDimPosition(outputMap.getNumResults() - 1))
      sizeN = rhsShape[i];
  }
  // Default tile size and workgroup size.
  int64_t tileX = 2;
  int64_t tileY = 256;
  SmallVector<int64_t, 3> workgroupSize = {2 * cudaWarpSize, 1, 1};
  SmallVector<TileWorkgroupSizePair> tileSizeConfig;
  // Query the best configuration.
  getMatmulConfig(tileSizeConfig);
  // Pick the best configuration where the original shape is aligned on the tile
  // size.
  for (TileWorkgroupSizePair& config : tileSizeConfig) {
    if (sizeN % config.tileSize[1] == 0 && sizeM % config.tileSize[0] == 0) {
      tileX = config.tileSize[0];
      tileY = config.tileSize[1];
      workgroupSize.assign(config.workgroupSize.begin(),
                           config.workgroupSize.end());
      break;
    }
  }
  // Currently just a basic tile size to enable tiling and vectorization.
  // TODO: pick a more efficient tile size and tile at subgroup level.
  SmallVector<int64_t, 4> ts;
  // Tile all the higher parallel dimension with a size of 1 and the 2 most
  // inner dimension with the tileX/tileY size.
  ts.append(op.getNumParallelLoops() - 2, 1);
  ts.append({tileX, tileY});
  // Tile all the reduction dimension with a size of 4.
  ts.append(op.getNumReductionLoops(), 4);
  tileSizes.push_back(ts);  // Workgroup level.
  tileSizes.push_back({});  // Subgroup level.
  // At the thread level only tile parallel loops.
  SmallVector<int64_t, 4> invocationLevelTs(op.getNumParallelLoops() - 2, 1);
  invocationLevelTs.append(
{'fraction_non_alphanumeric': 0.10903426791277258, 'fraction_numerical': 0.01838006230529595, 'mean_word_length': 4.329828444936359, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 25, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6978464', 'n_tokens_mistral': 3181, 'n_tokens_neox': 3066, 'n_words': 865}