Intrinsics.h

Go to the documentation of this file.

//===- Intrinsics.h - MLIR Operations for Declarative Builders ---*- 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 intrinsics for building snippets of MLIR
// declaratively
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_EDSC_INTRINSICS_H_
#define MLIR_EDSC_INTRINSICS_H_

#include "mlir/EDSC/Builders.h"
#include "mlir/Support/LLVM.h"

namespace mlir {

class MemRefType;
class Type;

namespace edsc {

struct IndexHandle : public ValueHandle {
  explicit IndexHandle()
      : ValueHandle(ScopedContext::getBuilder().getIndexType()) {}
  explicit IndexHandle(index_t v) : ValueHandle(v) {}
  explicit IndexHandle(Value v) : ValueHandle(v) {
    assert(v->getType() == ScopedContext::getBuilder().getIndexType() &&
           "Expected index type");
  }
  explicit IndexHandle(ValueHandle v) : ValueHandle(v) {
    assert(v.getType() == ScopedContext::getBuilder().getIndexType() &&
           "Expected index type");
  }
  IndexHandle &operator=(const ValueHandle &v) {
    assert(v.getType() == ScopedContext::getBuilder().getIndexType() &&
           "Expected index type");
    this->t = v.getType();
    this->v = v.getValue();
    return *this;
  }
};

inline SmallVector<IndexHandle, 8> makeIndexHandles(unsigned rank) {
  return SmallVector<IndexHandle, 8>(rank);
}

template <typename T>
inline SmallVector<ValueHandle *, 8>
makeHandlePointers(MutableArrayRef<T> ivs) {
  SmallVector<ValueHandle *, 8> pivs;
  pivs.reserve(ivs.size());
  for (auto &iv : ivs) {
    pivs.push_back(&iv);
  }
  return pivs;
}

inline SmallVector<Value, 8> extractValues(ArrayRef<IndexHandle> ivs) {
  SmallVector<Value, 8> vals;
  vals.reserve(ivs.size());
  for (auto &iv : ivs) {
    vals.push_back(iv.getValue());
  }
  return vals;
}

namespace intrinsics {
namespace detail {
class ValueHandleArray {
public:
  ValueHandleArray(ArrayRef<ValueHandle> vals) {
    values.append(vals.begin(), vals.end());
  }
  ValueHandleArray(ArrayRef<IndexHandle> vals) {
    values.append(vals.begin(), vals.end());
  }
  ValueHandleArray(ArrayRef<index_t> vals) {
    SmallVector<IndexHandle, 8> tmp(vals.begin(), vals.end());
    values.append(tmp.begin(), tmp.end());
  }
  operator ArrayRef<Value>() { return values; }

private:
  ValueHandleArray() = default;
  SmallVector<Value, 8> values;
};

template <typename T> inline T unpack(T value) { return value; }

inline detail::ValueHandleArray unpack(ArrayRef<ValueHandle> values) {
  return detail::ValueHandleArray(values);
}

} // namespace detail

template <typename Op> struct ValueBuilder : public ValueHandle {
  // Builder-based
  template <typename... Args>
  ValueBuilder(Args... args)
      : ValueHandle(ValueHandle::create<Op>(detail::unpack(args)...)) {}
  ValueBuilder(ArrayRef<ValueHandle> vs)
      : ValueBuilder(ValueBuilder::create<Op>(detail::unpack(vs))) {}
  template <typename... Args>
  ValueBuilder(ArrayRef<ValueHandle> vs, Args... args)
      : ValueHandle(ValueHandle::create<Op>(detail::unpack(vs),
                                            detail::unpack(args)...)) {}
  template <typename T, typename... Args>
  ValueBuilder(T t, ArrayRef<ValueHandle> vs, Args... args)
      : ValueHandle(ValueHandle::create<Op>(
            detail::unpack(t), detail::unpack(vs), detail::unpack(args)...)) {}
  template <typename T1, typename T2, typename... Args>
  ValueBuilder(T1 t1, T2 t2, ArrayRef<ValueHandle> vs, Args... args)
      : ValueHandle(ValueHandle::create<Op>(
            detail::unpack(t1), detail::unpack(t2), detail::unpack(vs),
            detail::unpack(args)...)) {}

  template <typename... Args>
  ValueBuilder(OperationFolder *folder, Args... args)
      : ValueHandle(ValueHandle::create<Op>(folder, detail::unpack(args)...)) {}
  ValueBuilder(OperationFolder *folder, ArrayRef<ValueHandle> vs)
      : ValueBuilder(ValueBuilder::create<Op>(folder, detail::unpack(vs))) {}
  template <typename... Args>
  ValueBuilder(OperationFolder *folder, ArrayRef<ValueHandle> vs, Args... args)
      : ValueHandle(ValueHandle::create<Op>(folder, detail::unpack(vs),
                                            detail::unpack(args)...)) {}
  template <typename T, typename... Args>
  ValueBuilder(OperationFolder *folder, T t, ArrayRef<ValueHandle> vs,
               Args... args)
      : ValueHandle(ValueHandle::create<Op>(folder, detail::unpack(t),
                                            detail::unpack(vs),
                                            detail::unpack(args)...)) {}
  template <typename T1, typename T2, typename... Args>
  ValueBuilder(OperationFolder *folder, T1 t1, T2 t2, ArrayRef<ValueHandle> vs,
               Args... args)
      : ValueHandle(ValueHandle::create<Op>(
            folder, detail::unpack(t1), detail::unpack(t2), detail::unpack(vs),
            detail::unpack(args)...)) {}

  ValueBuilder() : ValueHandle(ValueHandle::create<Op>()) {}
};

template <typename Op> struct OperationBuilder : public OperationHandle {
  template <typename... Args>
  OperationBuilder(Args... args)
      : OperationHandle(OperationHandle::create<Op>(detail::unpack(args)...)) {}
  OperationBuilder(ArrayRef<ValueHandle> vs)
      : OperationHandle(OperationHandle::create<Op>(detail::unpack(vs))) {}
  template <typename... Args>
  OperationBuilder(ArrayRef<ValueHandle> vs, Args... args)
      : OperationHandle(OperationHandle::create<Op>(detail::unpack(vs),
                                                    detail::unpack(args)...)) {}
  template <typename T, typename... Args>
  OperationBuilder(T t, ArrayRef<ValueHandle> vs, Args... args)
      : OperationHandle(OperationHandle::create<Op>(
            detail::unpack(t), detail::unpack(vs), detail::unpack(args)...)) {}
  template <typename T1, typename T2, typename... Args>
  OperationBuilder(T1 t1, T2 t2, ArrayRef<ValueHandle> vs, Args... args)
      : OperationHandle(OperationHandle::create<Op>(
            detail::unpack(t1), detail::unpack(t2), detail::unpack(vs),
            detail::unpack(args)...)) {}
  OperationBuilder() : OperationHandle(OperationHandle::create<Op>()) {}
};

using addf = ValueBuilder<AddFOp>;
using affine_apply = ValueBuilder<AffineApplyOp>;
using affine_if = OperationBuilder<AffineIfOp>;
using affine_load = ValueBuilder<AffineLoadOp>;
using affine_store = OperationBuilder<AffineStoreOp>;
using alloc = ValueBuilder<AllocOp>;
using call = OperationBuilder<mlir::CallOp>;
using constant_float = ValueBuilder<ConstantFloatOp>;
using constant_index = ValueBuilder<ConstantIndexOp>;
using constant_int = ValueBuilder<ConstantIntOp>;
using dealloc = OperationBuilder<DeallocOp>;
using dim = ValueBuilder<DimOp>;
using muli = ValueBuilder<MulIOp>;
using mulf = ValueBuilder<MulFOp>;
using memref_cast = ValueBuilder<MemRefCastOp>;
using ret = OperationBuilder<ReturnOp>;
using select = ValueBuilder<SelectOp>;
using std_load = ValueBuilder<LoadOp>;
using std_store = OperationBuilder<StoreOp>;
using subi = ValueBuilder<SubIOp>;
using tanh = ValueBuilder<TanhOp>;
using view = ValueBuilder<ViewOp>;

OperationHandle br(BlockHandle bh, ArrayRef<ValueHandle> operands);

OperationHandle br(BlockHandle *bh, ArrayRef<ValueHandle *> captures,
                   ArrayRef<ValueHandle> operands);

OperationHandle cond_br(ValueHandle cond, BlockHandle trueBranch,
                        ArrayRef<ValueHandle> trueOperands,
                        BlockHandle falseBranch,
                        ArrayRef<ValueHandle> falseOperands);

OperationHandle cond_br(ValueHandle cond, BlockHandle *trueBranch,
                        ArrayRef<ValueHandle *> trueCaptures,
                        ArrayRef<ValueHandle> trueOperands,
                        BlockHandle *falseBranch,
                        ArrayRef<ValueHandle *> falseCaptures,
                        ArrayRef<ValueHandle> falseOperands);
} // namespace intrinsics
} // namespace edsc
} // namespace mlir

#endif // MLIR_EDSC_INTRINSICS_H_