Builders.h

Go to the documentation of this file.

//===- Builders.h - MLIR Declarative Builder Classes ------------*- 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
//
//===----------------------------------------------------------------------===//
//
// Provides intuitive composable interfaces for building structured MLIR
// snippets in a declarative fashion.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_EDSC_BUILDERS_H_
#define MLIR_EDSC_BUILDERS_H_

#include "mlir/Dialect/AffineOps/AffineOps.h"
#include "mlir/Dialect/LoopOps/LoopOps.h"
#include "mlir/Dialect/StandardOps/Ops.h"
#include "mlir/IR/Builders.h"
#include "mlir/Transforms/FoldUtils.h"

namespace mlir {

namespace edsc {

struct index_t {
  explicit index_t(int64_t v) : v(v) {}
  explicit operator int64_t() { return v; }
  int64_t v;
};

class BlockHandle;
class CapturableHandle;
class NestedBuilder;
class ValueHandle;

class ScopedContext {
public:
  ScopedContext(OpBuilder &builder, Location location);

  ScopedContext(OpBuilder &builder, OpBuilder::InsertPoint newInsertPt,
                Location location);
  ~ScopedContext();

  static MLIRContext *getContext();
  static OpBuilder &getBuilder();
  static Location getLocation();

private:
  friend class NestedBuilder;

  ScopedContext() = delete;
  ScopedContext(const ScopedContext &) = delete;
  ScopedContext &operator=(const ScopedContext &) = delete;

  static ScopedContext *&getCurrentScopedContext();

  OpBuilder &builder;
  Optional<OpBuilder::InsertPoint> prevBuilderInsertPoint;
  Location location;
  ScopedContext *enclosingScopedContext;
  NestedBuilder *nestedBuilder;

  // TODO: Implement scoping of ValueHandles. To do this we need a proper data
  // structure to hold ValueHandle objects. We can emulate one but there should
  // already be something available in LLVM for this purpose.
};

class NestedBuilder {
protected:
  NestedBuilder() = default;
  NestedBuilder(const NestedBuilder &) = delete;
  NestedBuilder(NestedBuilder &&other) : bodyScope(other.bodyScope) {
    other.bodyScope = nullptr;
  }

  NestedBuilder &operator=(const NestedBuilder &) = delete;
  NestedBuilder &operator=(NestedBuilder &&other) {
    std::swap(bodyScope, other.bodyScope);
    return *this;
  }

  void enter(mlir::Block *block, int prev = 0) {
    bodyScope = new ScopedContext(
        ScopedContext::getBuilder(),
        OpBuilder::InsertPoint(block, std::prev(block->end(), prev)),
        ScopedContext::getLocation());
    bodyScope->nestedBuilder = this;
  }

  void exit() {
    // Reclaim now to exit the scope.
    bodyScope->nestedBuilder = nullptr;
    delete bodyScope;
    bodyScope = nullptr;
  }

  ~NestedBuilder() {
    assert(!bodyScope &&
           "Illegal use of NestedBuilder; must have called exit()");
  }

private:
  ScopedContext *bodyScope = nullptr;
};

class LoopBuilder : public NestedBuilder {
public:
  static LoopBuilder makeAffine(ValueHandle *iv,
                                ArrayRef<ValueHandle> lbHandles,
                                ArrayRef<ValueHandle> ubHandles, int64_t step);
  static LoopBuilder makeLoop(ValueHandle *iv, ValueHandle lbHandle,
                              ValueHandle ubHandle, ValueHandle stepHandle);
  LoopBuilder(const LoopBuilder &) = delete;
  LoopBuilder(LoopBuilder &&) = default;

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

  void operator()(function_ref<void(void)> fun = nullptr);

private:
  LoopBuilder() = default;
};

class AffineLoopNestBuilder {
public:
  // This entry point accommodates the fact that AffineForOp implicitly uses
  // multiple `lbs` and `ubs` with one single `iv` and `step` to encode `max`
  // and and `min` constraints respectively.
  AffineLoopNestBuilder(ValueHandle *iv, ArrayRef<ValueHandle> lbs,
                        ArrayRef<ValueHandle> ubs, int64_t step);
  AffineLoopNestBuilder(ArrayRef<ValueHandle *> ivs, ArrayRef<ValueHandle> lbs,
                        ArrayRef<ValueHandle> ubs, ArrayRef<int64_t> steps);

  void operator()(function_ref<void(void)> fun = nullptr);

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

class LoopNestBuilder {
public:
  LoopNestBuilder(ArrayRef<edsc::ValueHandle *> ivs, ArrayRef<ValueHandle> lbs,
                  ArrayRef<ValueHandle> ubs, ArrayRef<ValueHandle> steps);
  void operator()(std::function<void(void)> fun = nullptr);

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

// This class exists solely to handle the C++ vexing parse case when
// trying to enter a Block that has already been constructed.
class Append {};

class BlockBuilder : public NestedBuilder {
public:
  BlockBuilder(BlockHandle bh, Append);

  BlockBuilder(BlockHandle *bh, ArrayRef<ValueHandle *> args);

  void operator()(function_ref<void(void)> fun = nullptr);

private:
  BlockBuilder(BlockBuilder &) = delete;
  BlockBuilder &operator=(BlockBuilder &other) = delete;
};

class CapturableHandle {
protected:
  CapturableHandle() = default;
};

class ValueHandle : public CapturableHandle {
public:
  static ValueHandle null() { return ValueHandle(); }

  explicit ValueHandle(Type t) : t(t), v(nullptr) {}

  explicit ValueHandle(Value v) : t(v->getType()), v(v) {}

  ValueHandle(index_t cst);

  ValueHandle(const ValueHandle &other) = default;

  ValueHandle &operator=(const ValueHandle &other);

  void swap(ValueHandle &other) {
    if (this == &other)
      return;
    std::swap(t, other.t);
    std::swap(v, other.v);
  }

  operator Value() const { return getValue(); }
  operator bool() const { return hasValue(); }

  template <typename Op, typename... Args>
  static ValueHandle create(Args... args);

  template <typename Op, typename... Args>
  static ValueHandle create(OperationFolder *folder, Args... args);

  // TODO: createOrFold when available and move inside of the `create` method.
  static ValueHandle createComposedAffineApply(AffineMap map,
                                               ArrayRef<Value> operands);

  static ValueHandle create(StringRef name, ArrayRef<ValueHandle> operands,
                            ArrayRef<Type> resultTypes,
                            ArrayRef<NamedAttribute> attributes = {});

  bool hasValue() const { return v != nullptr; }
  Value getValue() const {
    assert(hasValue() && "Unexpected null value;");
    return v;
  }
  bool hasType() const { return t != Type(); }
  Type getType() const { return t; }

  Operation *getOperation() const {
    if (!v)
      return nullptr;
    return v->getDefiningOp();
  }

protected:
  ValueHandle() : t(), v(nullptr) {}

  Type t;
  Value v;
};

struct OperationHandle : public CapturableHandle {
  OperationHandle() : op(nullptr) {}
  OperationHandle(Operation *op) : op(op) {}

  OperationHandle(const OperationHandle &) = default;
  OperationHandle &operator=(const OperationHandle &) = default;

  template <typename Op, typename... Args>
  static OperationHandle create(Args... args);
  template <typename Op, typename... Args> static Op createOp(Args... args);

  static OperationHandle create(StringRef name, ArrayRef<ValueHandle> operands,
                                ArrayRef<Type> resultTypes,
                                ArrayRef<NamedAttribute> attributes = {});

  operator Operation *() { return op; }
  Operation *getOperation() const { return op; }

private:
  Operation *op;
};

template <typename HandleType> struct CustomOperation {
  CustomOperation(StringRef name) : name(name) {
    static_assert(std::is_same<HandleType, ValueHandle>() ||
                      std::is_same<HandleType, OperationHandle>(),
                  "Only CustomOperation<ValueHandle> or "
                  "CustomOperation<OperationHandle> can be constructed.");
  }
  HandleType operator()(ArrayRef<ValueHandle> operands = {},
                        ArrayRef<Type> resultTypes = {},
                        ArrayRef<NamedAttribute> attributes = {}) {
    return HandleType::create(name, operands, resultTypes, attributes);
  }
  std::string name;
};

class BlockHandle : public CapturableHandle {
public:
  BlockHandle() : block(nullptr) {}

  BlockHandle(mlir::Block *block) : block(block) {}

  BlockHandle(const BlockHandle &) = default;
  BlockHandle &operator=(const BlockHandle &) = default;

  static BlockHandle create(ArrayRef<Type> argTypes);

  operator bool() { return block != nullptr; }
  operator mlir::Block *() { return block; }
  mlir::Block *getBlock() { return block; }

private:
  mlir::Block *block;
};

template <typename Op, typename... Args>
OperationHandle OperationHandle::create(Args... args) {
  return OperationHandle(ScopedContext::getBuilder()
                             .create<Op>(ScopedContext::getLocation(), args...)
                             .getOperation());
}

template <typename Op, typename... Args>
Op OperationHandle::createOp(Args... args) {
  return cast<Op>(
      OperationHandle(ScopedContext::getBuilder()
                          .create<Op>(ScopedContext::getLocation(), args...)
                          .getOperation())
          .getOperation());
}

template <typename Op, typename... Args>
ValueHandle ValueHandle::create(Args... args) {
  Operation *op = ScopedContext::getBuilder()
                      .create<Op>(ScopedContext::getLocation(), args...)
                      .getOperation();
  if (op->getNumResults() == 1) {
    return ValueHandle(op->getResult(0));
  } else if (op->getNumResults() == 0) {
    if (auto f = dyn_cast<AffineForOp>(op)) {
      return ValueHandle(f.getInductionVar());
    }
  }
  llvm_unreachable("unsupported operation, use an OperationHandle instead");
}

template <typename Op, typename... Args>
ValueHandle ValueHandle::create(OperationFolder *folder, Args... args) {
  return folder ? ValueHandle(folder->create<Op>(ScopedContext::getBuilder(),
                                                 ScopedContext::getLocation(),
                                                 args...))
                : ValueHandle(ScopedContext::getBuilder().create<Op>(
                      ScopedContext::getLocation(), args...));
}

namespace op {

ValueHandle operator+(ValueHandle lhs, ValueHandle rhs);
ValueHandle operator-(ValueHandle lhs, ValueHandle rhs);
ValueHandle operator*(ValueHandle lhs, ValueHandle rhs);
ValueHandle operator/(ValueHandle lhs, ValueHandle rhs);
ValueHandle operator%(ValueHandle lhs, ValueHandle rhs);
ValueHandle floorDiv(ValueHandle lhs, ValueHandle rhs);
ValueHandle ceilDiv(ValueHandle lhs, ValueHandle rhs);

ValueHandle operator!(ValueHandle value);
ValueHandle operator&&(ValueHandle lhs, ValueHandle rhs);
ValueHandle operator||(ValueHandle lhs, ValueHandle rhs);
ValueHandle operator^(ValueHandle lhs, ValueHandle rhs);
ValueHandle operator==(ValueHandle lhs, ValueHandle rhs);
ValueHandle operator!=(ValueHandle lhs, ValueHandle rhs);
ValueHandle operator<(ValueHandle lhs, ValueHandle rhs);
ValueHandle operator<=(ValueHandle lhs, ValueHandle rhs);
ValueHandle operator>(ValueHandle lhs, ValueHandle rhs);
ValueHandle operator>=(ValueHandle lhs, ValueHandle rhs);

} // namespace op

template <typename Container>
inline SmallVector<ValueHandle, 8> makeValueHandles(Container values) {
  SmallVector<ValueHandle, 8> res;
  res.reserve(values.size());
  for (auto v : values)
    res.push_back(ValueHandle(v));
  return res;
}

} // namespace edsc
} // namespace mlir

#endif // MLIR_EDSC_BUILDERS_H_