iree/compiler/Codegen/SPIRV/LaunchConfig.h

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

//===- LaunchConfig.h - Configuration used to drive arch specific codegen -===//
//
// This file declares the data structure that is used by the codegeneration to
// lower to target specific IR. The values of the parameters are archtecture
// specific. Once set the same transformations can be used to generate the
// desired code. This allows sharing codegen infra between different backends.
//
//===----------------------------------------------------------------------===//
#ifndef IREE_COMPILER_CODEGEN_COMMON_LAUNCHCONFIG_H_
#define IREE_COMPILER_CODEGEN_COMMON_LAUNCHCONFIG_H_
#include <array>

#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h"

namespace mlir {
namespace iree_compiler {

/// Stores the tile sizes to use at different levels of tiling as a vector of
/// vectors.
/// - First level tiling maps to workgroups.
/// - Second level tiling maps to subgroups.
/// - Third level tiling maps to invocations.
using TileSizesListType = SmallVector<SmallVector<int64_t, 4>, 1>;
using TileSizesListTypeRef = ArrayRef<SmallVector<int64_t, 4>>;

/// Configurations for mapping Linalg ops to CPU/GPU parallel hiearchies.
///
/// Based on the linalg operations in a dispatch region, the number of levels of
/// tiling, the tile sizes needed, the workgroup size, etc. need to be
/// decided. These parameters are called `LaunchConfig`. This class implements
/// one heuristic to compute these for the different linalg operations on
/// buffers. This can be adapted later to support multiple configurations that
/// can be picked based on device information/problem size information. It
/// exposes the information needed by the codegenerators, and hides the
/// implementation from the rest of the pipeline.
class LaunchConfig {
 public:
  LaunchConfig() : workgroupSize({1, 1, 1}), numSubgroups({1, 1, 1}) {}

  /// Removes attributes added to operations for retrieving tile size
  /// information.
  void finalize(FuncOp funcOp);

  /// Gets the tile size computed for an operation at all levels.
  TileSizesListTypeRef getTileSizes(Operation *op) const;

  /// Gets the tile size computed for an operation for an level.
  ArrayRef<int64_t> getTileSizes(Operation *op, size_t level) const;

  /// Returns the workgroup size to use based on the tile sizes.
  ArrayRef<int64_t> getWorkgroupSize() const { return workgroupSize; }

  /// Returns the number of subgroups to use.
  ArrayRef<int64_t> getNumSubgroups() const { return numSubgroups; }

  /// Of the given operations return the operation that has been marked as the
  /// root operation. Within a dispatch region a single root operation (like
  /// matmul, conv, etc.) decides the launch configuration to be used. The rest
  /// of the ops that are fused with it obey this configuration. Returns nullptr
  /// if unable to find an operation that is set as root in the list.
  Operation *getRootOperation(ArrayRef<Operation *> ops);

  /// Returns true if tile sizes have been computed for the operation. If tile
  /// sizes arent set, it implies operation is not to be tiled.
  bool hasTileSizes(Operation *op, size_t level = 0) const {
    return !getTileSizes(op, level).empty();
  }

  /// Use vectorize transformations.
  bool useVectorize() const { return vectorize; }

  /// Sets the tile sizes to use for all levels of tiling of `op`.
  void setTileSizes(Operation *op, TileSizesListType vTileSizes);

  /// Sets the tile sizes to use for a given `level` of tiling of `op`.
  void setTileSizes(Operation *op, ArrayRef<int64_t> vTileSizes, size_t level);

  /// Sets the workgroup size to use for the function.
  void setWorkgroupSize(ArrayRef<int64_t> vWorkgroupSize);

  /// Sets number of subgroups to use.
  void setNumSubgroups(ArrayRef<int64_t> vNumSubgroups);

  /// Sets the root operation. Within a dispatch region a single root operation
  /// (like matmul, conv, etc.) decides the launch configuration to be used. The
  /// rest of the ops that are fused with it obey this configuration.
  void setRootOperation(Operation *root);

  /// Sets the configuration of the `targetOp` to be same as the configuration
  /// of the `sourceOp`.
  void setSameConfig(Operation *sourceOp, Operation *targetOp);

  /// Sets flag to enable vectorization.
  void setVectorize(bool enableVectorize);

 protected:
  /// Current tile size configuration per operation. They key used here to
  /// retrieve the tile size information per operation is the value of a StrAttr
  /// added to operations during `init`. When tiled this attribute is copied
  /// over to the tiled operation, thereby the same key can be used to retrieve
  /// the tile sizes for the next level of tiling. The `finalize` method removes
  /// these attributes.
  llvm::StringMap<TileSizesListType> tileSizes;

  /// Workgroup size to use.
  std::array<int64_t, 3> workgroupSize = {1, 1, 1};

  /// Number of subgroups that are logically distributed along x, y & z.
  std::array<int64_t, 3> numSubgroups = {1, 1, 1};

  /// Use vectorization.
  bool vectorize = false;
};

/// Propogates tile sizes from `rootOperation` to other linalg operations in the
/// dispatch region. This assumes that each dispatch region has a single root
/// operation (like matmul, conv, etc.) that determines the tile sizes to use
/// for tile+fuse+distribute. These are then propogated to the other operations.
/// Note: This is a temporary solution and might be defunct when the codegen
/// becomes more sophisticated.
LogicalResult propogateRootOperationLaunchConfig(
    LaunchConfig &launchConfig, linalg::LinalgOp rootOperation,
    const linalg::LinalgDependenceGraph &dependenceGraph);

}  // namespace iree_compiler
}  // namespace mlir
#endif  // IREE_COMPILER_CODEGEN_COMMON_LAUNCHCONFIG_H_