Utils.h

Go to the documentation of this file.

//===- Utils.h - Utilities to support the Linalg dialect --------*- C++ -*-===//
//
// Part of the MLIR Project, 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
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_DIALECT_LINALG_UTILS_H_
#define MLIR_DIALECT_LINALG_UTILS_H_

#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/LoopOps/LoopOps.h"
#include "mlir/Dialect/StandardOps/Ops.h"
#include "mlir/EDSC/Helpers.h"

#include "llvm/ADT/SetVector.h"

namespace mlir {
class AffineExpr;
class AffineMap;
class OperationFolder;

namespace edsc {

class LoopRangeBuilder : public NestedBuilder {
public:
  LoopRangeBuilder(ValueHandle *iv, ValueHandle range);
  LoopRangeBuilder(ValueHandle *iv, Value range);
  LoopRangeBuilder(ValueHandle *iv, SubViewOp::Range range);

  LoopRangeBuilder(const LoopRangeBuilder &) = delete;
  LoopRangeBuilder(LoopRangeBuilder &&) = default;

  LoopRangeBuilder &operator=(const LoopRangeBuilder &) = delete;
  LoopRangeBuilder &operator=(LoopRangeBuilder &&) = default;

  ValueHandle operator()(std::function<void(void)> fun = nullptr);
};

class LoopNestRangeBuilder {
public:
  LoopNestRangeBuilder(ArrayRef<edsc::ValueHandle *> ivs,
                       ArrayRef<edsc::ValueHandle> ranges);
  LoopNestRangeBuilder(ArrayRef<edsc::ValueHandle *> ivs,
                       ArrayRef<Value> ranges);
  LoopNestRangeBuilder(ArrayRef<edsc::ValueHandle *> ivs,
                       ArrayRef<SubViewOp::Range> ranges);
  edsc::ValueHandle operator()(std::function<void(void)> fun = nullptr);

private:
  SmallVector<LoopRangeBuilder, 4> loops;
};

template <typename LoopTy> class GenericLoopNestRangeBuilder {
public:
  GenericLoopNestRangeBuilder(ArrayRef<edsc::ValueHandle *> ivs,
                              ArrayRef<Value> ranges);
  void operator()(std::function<void(void)> fun = nullptr) { (*builder)(fun); }

private:
  typedef typename std::conditional<std::is_same<LoopTy, AffineForOp>::value,
                                    AffineLoopNestBuilder,
                                    LoopNestRangeBuilder>::type BuilderType;
  std::unique_ptr<BuilderType> builder;
};

} // namespace edsc

namespace linalg {
class LinalgDependenceGraph;

struct FusionInfo {
  LinalgOp originalProducer;
  LinalgOp fusedProducer;
};

bool isProducerLastWriteOfView(const LinalgDependenceGraph &graph,
                               LinalgOp consumer, Value consumedView,
                               LinalgOp producer);

bool isFusableInto(const LinalgDependenceGraph &graph, LinalgOp consumer,
                   Value consumedView, LinalgOp producer);

Optional<FusionInfo> fuseProducerOf(OpBuilder &b, LinalgOp consumer,
                                    unsigned consumerIdx,
                                    const LinalgDependenceGraph &graph,
                                    OperationFolder *folder = nullptr);

template <typename ConcreteOp>
SmallVector<Value, 8> getViewSizes(ConcreteOp linalgOp) {
  SmallVector<Value, 8> res;
  for (auto v : linalgOp.getInputsAndOutputs()) {
    MemRefType t = v->getType().template cast<MemRefType>();
    for (unsigned i = 0; i < t.getRank(); ++i)
      res.push_back(edsc::intrinsics::dim(v, i));
  }
  return res;
}

SmallVector<Value, 4> applyMapToValues(OpBuilder &b, Location loc,
                                       AffineMap map, ArrayRef<Value> values,
                                       OperationFolder *folder = nullptr);

struct TiledLinalgOp {
  LinalgOp op;
  SmallVector<loop::ForOp, 8> loops;
};

Optional<TiledLinalgOp> tileLinalgOp(OpBuilder &b, LinalgOp op,
                                     ArrayRef<Value> tileSizes,
                                     ArrayRef<unsigned> permutation = {},
                                     OperationFolder *folder = nullptr);

Optional<TiledLinalgOp> tileLinalgOp(OpBuilder &b, LinalgOp op,
                                     ArrayRef<int64_t> tileSizes,
                                     ArrayRef<unsigned> permutation = {},
                                     OperationFolder *folder = nullptr);

template <typename... Args>
Optional<TiledLinalgOp> tileLinalgOperation(OpBuilder &b, Operation *op,
                                            Args... args) {
  return tileLinalgOp(b, cast<LinalgOp>(op), args...);
}

struct PromotionInfo {
  Value buffer;
  Value fullLocalView;
  Value partialLocalView;
};

SmallVector<PromotionInfo, 8>
promoteSubViews(OpBuilder &b, Location loc, ArrayRef<Value> subViews,
                bool dynamicBuffers = false, OperationFolder *folder = nullptr);

SmallVector<Value, 4> getAssumedNonViewOperands(LinalgOp linalgOp);

template <typename T, unsigned N>
void applyPermutationToVector(SmallVector<T, N> &inVec,
                              ArrayRef<unsigned> permutation) {
  SmallVector<T, N> auxVec(inVec.size());
  for (unsigned i = 0; i < permutation.size(); ++i)
    auxVec[i] = inVec[permutation[i]];
  inVec = auxVec;
}

LinalgOp promoteSubViewOperands(OpBuilder &b, LinalgOp op,
                                llvm::SetVector<Value> subViews,
                                bool dynamicBuffers = false,
                                OperationFolder *folder = nullptr);

} // namespace linalg
} // namespace mlir

#endif // MLIR_DIALECT_LINALG_UTILS_H_