forked from OSchip/llvm-project
Start a new implementation for edsc::Builder
This CL reworks the design of EDSCs from first principles. It introduces a ValueHandle which can hold either: 1. an eagerly typed, delayed Value* 2. a precomputed Value* A ValueHandle can be manipulated with intrinsic operations a nested within a NestedBuilder. These NestedBuilder are a more idiomatic nested builder abstraction that should feel intuitive to program in C++. Notably, this abstraction does not require an AST to stage the construction of MLIR snippets from C++. Instead, the abstraction makes use of orderings between definition and declaration of ValueHandles and provides a NestedBuilder and a LoopBuilder helper classes to handle those orderings. All instruction creations are meant to use the templated ValueHandle::create<> which directly calls mlir::Builder.create<>. For now the EDSC AST and the builders live side-by-side until the C API is ported. PiperOrigin-RevId: 237030945
This commit is contained in:
Nicolas Vasilache
jpienaar
parent
95949a0d09
commit
af6c3f7a63
|
|
@ -0,0 +1,251 @@
|
|||
//===- Builders.h - MLIR Declarative Builder Classes ------------*- C++ -*-===//
|
||||
//
|
||||
// Copyright 2019 The MLIR Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
// =============================================================================
|
||||
//
|
||||
// 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/AffineOps/AffineOps.h"
|
||||
#include "mlir/IR/Builders.h"
|
||||
#include "mlir/StandardOps/Ops.h"
|
||||
|
||||
namespace mlir {
|
||||
|
||||
namespace edsc {
|
||||
|
||||
struct index_t {
|
||||
explicit index_t(int64_t v) : v(v) {}
|
||||
int64_t v;
|
||||
};
|
||||
|
||||
class NestedBuilder;
|
||||
class ValueHandle;
|
||||
|
||||
/// Helper class to transparently handle builder insertion points by RAII.
|
||||
/// As its name indicates, a ScopedContext is means to be used locally in a
|
||||
/// scoped fashion. This abstracts away all the boilerplate related to
|
||||
/// checking proper usage of captures, NestedBuilders as well as handling the
|
||||
/// setting and restoring of insertion points.
|
||||
class ScopedContext {
|
||||
public:
|
||||
/// Sets location to fun->getLoc() in case the provided Loction* is null.
|
||||
ScopedContext(Function *fun, Location *loc = nullptr);
|
||||
ScopedContext(FuncBuilder builder, Location location);
|
||||
~ScopedContext();
|
||||
|
||||
static MLIRContext *getContext();
|
||||
static FuncBuilder *getBuilder();
|
||||
static Location getLocation();
|
||||
|
||||
private:
|
||||
/// Only NestedBuilder (which is used to create an instruction with a body)
|
||||
/// may access private members in order to implement scoping.
|
||||
friend class NestedBuilder;
|
||||
|
||||
ScopedContext() = delete;
|
||||
ScopedContext(const ScopedContext &) = delete;
|
||||
ScopedContext &operator=(const ScopedContext &) = delete;
|
||||
|
||||
static ScopedContext *&getCurrentScopedContext();
|
||||
|
||||
/// Current FuncBuilder.
|
||||
FuncBuilder builder;
|
||||
/// Current location.
|
||||
Location location;
|
||||
/// Parent context we return into.
|
||||
ScopedContext *enclosingScopedContext;
|
||||
/// Defensively keeps track of the current NestedBuilder to ensure proper
|
||||
/// scoping usage.
|
||||
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.
|
||||
};
|
||||
|
||||
/// A NestedBuilder is a scoping abstraction to create an idiomatic syntax
|
||||
/// embedded in C++ that serves the purpose of building nested MLIR.
|
||||
/// Nesting and compositionality is obtained by using the strict ordering that
|
||||
/// exists between object construction and method invocation on said object (in
|
||||
/// our case, the call to `operator()`).
|
||||
/// This ordering allows implementing an abstraction that decouples definition
|
||||
/// from declaration (in a PL sense) on placeholders of type ValueHandle and
|
||||
/// BlockHandle.
|
||||
class NestedBuilder {
|
||||
protected:
|
||||
/// Enter an mlir::Block and setup a ScopedContext to insert instructions at
|
||||
/// the end of it. Since we cannot use c++ language-level scoping to implement
|
||||
/// scoping itself, we use enter/exit pairs of instructions.
|
||||
/// As a consequence we must allocate a new FuncBuilder + ScopedContext and
|
||||
/// let the escape.
|
||||
void enter(mlir::Block *block) {
|
||||
bodyScope = new ScopedContext(FuncBuilder(block, block->end()),
|
||||
ScopedContext::getLocation());
|
||||
bodyScope->nestedBuilder = this;
|
||||
}
|
||||
|
||||
/// Exit the current mlir::Block by explicitly deleting the dynamically
|
||||
/// allocated FuncBuilder and ScopedContext.
|
||||
void exit() {
|
||||
// Reclaim now to exit the scope.
|
||||
bodyScope->nestedBuilder = nullptr;
|
||||
delete bodyScope;
|
||||
bodyScope = nullptr;
|
||||
}
|
||||
|
||||
/// Custom destructor does nothing because we already destroyed bodyScope
|
||||
/// manually in `exit`. Insert an assertion to defensively guard against
|
||||
/// improper usage of scoping.
|
||||
~NestedBuilder() {
|
||||
assert(!bodyScope &&
|
||||
"Illegal use of NestedBuilder; must have called exit()");
|
||||
}
|
||||
|
||||
private:
|
||||
ScopedContext *bodyScope = nullptr;
|
||||
};
|
||||
|
||||
/// A LoopBuilder is a generic NestedBuilder for loop-like MLIR instructions.
|
||||
/// More specifically it is meant to be used as a temporary object for
|
||||
/// representing any nested MLIR construct that is "related to" an mlir::Value*
|
||||
/// (for now an induction variable).
|
||||
/// This is extensible and will evolve in the future as MLIR evolves, hence
|
||||
/// the name LoopBuilder (as opposed to say ForBuilder or AffineForBuilder).
|
||||
class LoopBuilder : public NestedBuilder {
|
||||
public:
|
||||
/// Constructs a new AffineForOp and captures the associated induction
|
||||
/// variable. A ValueHandle pointer is passed as the first argument and is the
|
||||
/// *only* way to capture the loop induction variable.
|
||||
LoopBuilder(ValueHandle *iv, ArrayRef<ValueHandle> lbHandles,
|
||||
ArrayRef<ValueHandle> ubHandles, int64_t step);
|
||||
|
||||
/// The only purpose of this operator is to serve as a sequence point so that
|
||||
/// the evaluation of `stmts` (which build IR snippets in a scoped fashion) is
|
||||
/// sequenced strictly after the constructor of LoopBuilder.
|
||||
/// In order to be admissible in a nested ArrayRef<ValueHandle>, operator()
|
||||
/// returns a ValueHandle::null() that cannot be captured.
|
||||
// TODO(ntv): when loops return escaping ssa-values, this should be adapted.
|
||||
ValueHandle operator()(ArrayRef<ValueHandle> stmts);
|
||||
};
|
||||
|
||||
/// ValueHandle implements a (potentially "delayed") typed Value abstraction.
|
||||
/// ValueHandle should be captured by pointer but otherwise passed by Value
|
||||
/// everywhere.
|
||||
/// A ValueHandle can have 3 states:
|
||||
/// 1. null state (empty type and empty value), in which case it does not hold
|
||||
/// a value and may never hold a Value (not now of in the future). This is
|
||||
/// used for MLIR operations with zero returns as well as the result of
|
||||
/// calling a NestedBuilder::operator(). In both cases the objective is to
|
||||
/// have an object that can be inserted in an ArrayRef<ValueHandle> to
|
||||
/// implement nesting;
|
||||
/// 2. delayed state (empty value), in which case it represents an eagerly
|
||||
/// typed "delayed" value that can be hold a Value in the future;
|
||||
/// 3. constructed state,in which case it holds a Value.
|
||||
class ValueHandle {
|
||||
public:
|
||||
/// A ValueHandle in a null state can never be captured;
|
||||
static ValueHandle null() { return ValueHandle(); }
|
||||
|
||||
/// A ValueHandle that is constructed from a Type represents a typed "delayed"
|
||||
/// Value. A delayed Value can only capture Values of the specified type.
|
||||
/// Such a delayed value represents the declaration (in the PL sense) of a
|
||||
/// placeholder for an mlir::Value* that will be constructed and captured at
|
||||
/// some later point in the program.
|
||||
explicit ValueHandle(Type t) : t(t), v(nullptr) {}
|
||||
|
||||
/// A ValueHandle that is constructed from an mlir::Value* is an "eager"
|
||||
/// Value. An eager Value represents both the declaration and the definition
|
||||
/// (in the PL sense) of a placeholder for an mlir::Value* that has already
|
||||
/// been constructed in the past and that is captured "now" in the program.
|
||||
explicit ValueHandle(Value *v) : t(v->getType()), v(v) {}
|
||||
|
||||
/// Builds a ConstantIndexOp of value `cst`. The constant is created at the
|
||||
/// current insertion point.
|
||||
/// This implicit constructor is provided to each build an eager Value for a
|
||||
/// constant at the current insertion point in the IR. An implicit constructor
|
||||
/// allows idiomatic expressions mixing ValueHandle and literals.
|
||||
ValueHandle(index_t cst);
|
||||
|
||||
/// ValueHandle is a value type, use the default copy constructor.
|
||||
ValueHandle(const ValueHandle &other) = default;
|
||||
|
||||
/// ValueHandle is a value type, the assignment operator typechecks before
|
||||
/// assigning.
|
||||
/// ```
|
||||
ValueHandle &operator=(const ValueHandle &other);
|
||||
|
||||
/// Implicit conversion useful for automatic conversion to Container<Value*>.
|
||||
operator Value *() const { return getValue(); }
|
||||
|
||||
/// Generic mlir::Op create. This is the key to being extensible to the whole
|
||||
/// of MLIR without duplicating the type system or the AST.
|
||||
template <typename Op, typename... Args>
|
||||
static ValueHandle create(Args... args);
|
||||
|
||||
/// Special case to build composed AffineApply operations.
|
||||
// TODO: createOrFold when available and move inside of the `create` method.
|
||||
static ValueHandle createComposedAffineApply(AffineMap map,
|
||||
ArrayRef<Value *> operands);
|
||||
|
||||
bool hasValue() const { return v != nullptr; }
|
||||
Value *getValue() const { return v; }
|
||||
bool hasType() const { return t != Type(); }
|
||||
Type getType() const { return t; }
|
||||
|
||||
private:
|
||||
ValueHandle() : t(), v(nullptr) {}
|
||||
|
||||
Type t;
|
||||
Value *v;
|
||||
};
|
||||
|
||||
template <typename Op, typename... Args>
|
||||
ValueHandle ValueHandle::create(Args... args) {
|
||||
Instruction *inst = ScopedContext::getBuilder()
|
||||
->create<Op>(ScopedContext::getLocation(), args...)
|
||||
->getInstruction();
|
||||
if (inst->getNumResults() == 1) {
|
||||
return ValueHandle(inst->getResult(0));
|
||||
} else if (inst->getNumResults() == 0) {
|
||||
if (auto f = inst->dyn_cast<AffineForOp>()) {
|
||||
f->createBody();
|
||||
return ValueHandle(f->getInductionVar());
|
||||
}
|
||||
return ValueHandle();
|
||||
}
|
||||
llvm_unreachable("unsupported inst with > 1 results");
|
||||
}
|
||||
|
||||
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);
|
||||
|
||||
} // namespace op
|
||||
} // namespace edsc
|
||||
} // namespace mlir
|
||||
|
||||
#endif // MLIR_EDSC_BUILDERS_H_
|
||||
|
|
@ -0,0 +1,51 @@
|
|||
//===- Intrinsics.h - MLIR Operations for Declarative Builders ---*- C++-*-===//
|
||||
//
|
||||
// Copyright 2019 The MLIR Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
// =============================================================================
|
||||
//
|
||||
// Provides intuitive composable intrinsics for building snippets of MLIR
|
||||
// declaratively
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef MLIR_EDSC_INTRINSICS_H_
|
||||
#define MLIR_EDSC_INTRINSICS_H_
|
||||
|
||||
#include "mlir/Support/LLVM.h"
|
||||
|
||||
namespace mlir {
|
||||
|
||||
namespace edsc {
|
||||
|
||||
class ValueHandle;
|
||||
|
||||
/// Provides a set of first class intrinsics.
|
||||
/// In the future, most of intrinsics reated to Instruction that don't contain
|
||||
/// other instructions should be Tablegen'd.
|
||||
namespace intrinsics {
|
||||
|
||||
// TODO(ntv): Intrinsics below this line should be TableGen'd.
|
||||
/// Builds an mlir::ReturnOp with the proper `operands` that each must have
|
||||
/// captured an mlir::Value*.
|
||||
/// Returns an empty ValueHandle
|
||||
ValueHandle RETURN(llvm::ArrayRef<ValueHandle> operands);
|
||||
|
||||
} // namespace intrinsics
|
||||
|
||||
} // namespace edsc
|
||||
|
||||
} // namespace mlir
|
||||
|
||||
#endif // MLIR_EDSC_INTRINSICS_H_
|
||||
|
|
@ -0,0 +1,264 @@
|
|||
//===- Builders.cpp - MLIR Declarative Builder Classes ----------*- C++ -*-===//
|
||||
//
|
||||
// Copyright 2019 The MLIR Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
// =============================================================================
|
||||
|
||||
#include "mlir/EDSC/Builders.h"
|
||||
#include "mlir/IR/AffineExpr.h"
|
||||
#include "mlir/StandardOps/Ops.h"
|
||||
|
||||
#include "llvm/ADT/Optional.h"
|
||||
|
||||
using namespace mlir;
|
||||
using namespace mlir::edsc;
|
||||
|
||||
mlir::edsc::ScopedContext::ScopedContext(Function *fun, Location *loc)
|
||||
: builder(fun), location(loc ? *loc : fun->getLoc()),
|
||||
enclosingScopedContext(ScopedContext::getCurrentScopedContext()),
|
||||
nestedBuilder(nullptr) {
|
||||
getCurrentScopedContext() = this;
|
||||
}
|
||||
|
||||
mlir::edsc::ScopedContext::ScopedContext(FuncBuilder builder, Location location)
|
||||
: builder(builder), location(location),
|
||||
enclosingScopedContext(ScopedContext::getCurrentScopedContext()),
|
||||
nestedBuilder(nullptr) {
|
||||
getCurrentScopedContext() = this;
|
||||
}
|
||||
|
||||
mlir::edsc::ScopedContext::~ScopedContext() {
|
||||
assert(!nestedBuilder &&
|
||||
"Active NestedBuilder must have been exited at this point!");
|
||||
getCurrentScopedContext() = enclosingScopedContext;
|
||||
}
|
||||
|
||||
ScopedContext *&mlir::edsc::ScopedContext::getCurrentScopedContext() {
|
||||
thread_local ScopedContext *context = nullptr;
|
||||
return context;
|
||||
}
|
||||
|
||||
FuncBuilder *mlir::edsc::ScopedContext::getBuilder() {
|
||||
assert(ScopedContext::getCurrentScopedContext() &&
|
||||
"Unexpected Null ScopedContext");
|
||||
return &ScopedContext::getCurrentScopedContext()->builder;
|
||||
}
|
||||
|
||||
Location mlir::edsc::ScopedContext::getLocation() {
|
||||
assert(ScopedContext::getCurrentScopedContext() &&
|
||||
"Unexpected Null ScopedContext");
|
||||
return ScopedContext::getCurrentScopedContext()->location;
|
||||
}
|
||||
|
||||
MLIRContext *mlir::edsc::ScopedContext::getContext() {
|
||||
assert(getBuilder() && "Unexpected null builder");
|
||||
return getBuilder()->getContext();
|
||||
}
|
||||
|
||||
mlir::edsc::ValueHandle::ValueHandle(index_t cst) {
|
||||
auto *b = ScopedContext::getBuilder();
|
||||
auto loc = ScopedContext::getLocation();
|
||||
v = b->create<ConstantIndexOp>(loc, cst.v)->getResult();
|
||||
t = v->getType();
|
||||
}
|
||||
|
||||
ValueHandle &mlir::edsc::ValueHandle::operator=(const ValueHandle &other) {
|
||||
assert(t == other.t && "Wrong type capture");
|
||||
assert(!v && "ValueHandle has already been captured, use a new name!");
|
||||
v = other.v;
|
||||
return *this;
|
||||
}
|
||||
|
||||
ValueHandle
|
||||
mlir::edsc::ValueHandle::createComposedAffineApply(AffineMap map,
|
||||
ArrayRef<Value *> operands) {
|
||||
assert(ScopedContext::getBuilder() && "Unexpected null builder");
|
||||
Instruction *inst =
|
||||
makeComposedAffineApply(ScopedContext::getBuilder(),
|
||||
ScopedContext::getLocation(), map, operands)
|
||||
->getInstruction();
|
||||
assert(inst->getNumResults() == 1 && "Not a single result AffineApply");
|
||||
return ValueHandle(inst->getResult(0));
|
||||
}
|
||||
|
||||
static llvm::Optional<ValueHandle> emitStaticFor(ArrayRef<ValueHandle> lbs,
|
||||
ArrayRef<ValueHandle> ubs,
|
||||
int64_t step) {
|
||||
if (lbs.size() != 1 || ubs.size() != 1)
|
||||
return llvm::Optional<ValueHandle>();
|
||||
|
||||
auto *lbDef = lbs.front().getValue()->getDefiningInst();
|
||||
auto *ubDef = ubs.front().getValue()->getDefiningInst();
|
||||
if (!lbDef || !ubDef)
|
||||
return llvm::Optional<ValueHandle>();
|
||||
|
||||
auto lbConst = lbDef->dyn_cast<ConstantIndexOp>();
|
||||
auto ubConst = ubDef->dyn_cast<ConstantIndexOp>();
|
||||
if (!lbConst || !ubConst)
|
||||
return llvm::Optional<ValueHandle>();
|
||||
|
||||
return ValueHandle::create<AffineForOp>(lbConst->getValue(),
|
||||
ubConst->getValue(), step);
|
||||
}
|
||||
|
||||
mlir::edsc::LoopBuilder::LoopBuilder(ValueHandle *iv,
|
||||
ArrayRef<ValueHandle> lbHandles,
|
||||
ArrayRef<ValueHandle> ubHandles,
|
||||
int64_t step) {
|
||||
if (auto res = emitStaticFor(lbHandles, ubHandles, step)) {
|
||||
*iv = res.getValue();
|
||||
} else {
|
||||
SmallVector<Value *, 4> lbs(lbHandles.begin(), lbHandles.end());
|
||||
SmallVector<Value *, 4> ubs(ubHandles.begin(), ubHandles.end());
|
||||
*iv = ValueHandle::create<AffineForOp>(
|
||||
lbs, ScopedContext::getBuilder()->getMultiDimIdentityMap(lbs.size()),
|
||||
ubs, ScopedContext::getBuilder()->getMultiDimIdentityMap(ubs.size()),
|
||||
step);
|
||||
}
|
||||
auto *body = getForInductionVarOwner(iv->getValue())->getBody();
|
||||
enter(body);
|
||||
}
|
||||
|
||||
ValueHandle mlir::edsc::LoopBuilder::operator()(ArrayRef<ValueHandle> stmts) {
|
||||
// Call to `exit` must be explicit and asymmetric (cannot happen in the
|
||||
// destructor) because of ordering wrt comma operator.
|
||||
/// The particular use case concerns nested blocks:
|
||||
///
|
||||
/// ```c++
|
||||
/// For (&i, lb, ub, 1)({
|
||||
/// /--- destructor for this `For` is not always called before ...
|
||||
/// V
|
||||
/// For (&j1, lb, ub, 1)({
|
||||
/// some_op_1,
|
||||
/// }),
|
||||
/// /--- ... this scope is entered, resulting in improperly nested IR.
|
||||
/// V
|
||||
/// For (&j2, lb, ub, 1)({
|
||||
/// some_op_2,
|
||||
/// }),
|
||||
/// });
|
||||
/// ```
|
||||
exit();
|
||||
return ValueHandle::null();
|
||||
}
|
||||
|
||||
template <typename Op>
|
||||
static ValueHandle createBinaryHandle(ValueHandle lhs, ValueHandle rhs) {
|
||||
return ValueHandle::create<Op>(lhs.getValue(), rhs.getValue());
|
||||
}
|
||||
|
||||
static std::pair<AffineExpr, Value *>
|
||||
categorizeValueByAffineType(MLIRContext *context, Value *val, unsigned &numDims,
|
||||
unsigned &numSymbols) {
|
||||
AffineExpr d;
|
||||
Value *resultVal = nullptr;
|
||||
auto *constant = val->getDefiningInst()
|
||||
? val->getDefiningInst()->dyn_cast<ConstantIndexOp>()
|
||||
: nullptr;
|
||||
if (constant) {
|
||||
d = getAffineConstantExpr(constant->getValue(), context);
|
||||
} else if (isValidSymbol(val) && !isValidDim(val)) {
|
||||
d = getAffineSymbolExpr(numSymbols++, context);
|
||||
resultVal = val;
|
||||
} else {
|
||||
assert(isValidDim(val) && "Must be a valid Dim");
|
||||
d = getAffineDimExpr(numDims++, context);
|
||||
resultVal = val;
|
||||
}
|
||||
return std::make_pair(d, resultVal);
|
||||
}
|
||||
|
||||
static ValueHandle createBinaryIndexHandle(
|
||||
ValueHandle lhs, ValueHandle rhs,
|
||||
llvm::function_ref<AffineExpr(AffineExpr, AffineExpr)> affCombiner) {
|
||||
MLIRContext *context = ScopedContext::getContext();
|
||||
unsigned numDims = 0, numSymbols = 0;
|
||||
AffineExpr d0, d1;
|
||||
Value *v0, *v1;
|
||||
std::tie(d0, v0) =
|
||||
categorizeValueByAffineType(context, lhs.getValue(), numDims, numSymbols);
|
||||
std::tie(d1, v1) =
|
||||
categorizeValueByAffineType(context, rhs.getValue(), numDims, numSymbols);
|
||||
SmallVector<Value *, 2> operands;
|
||||
if (v0) {
|
||||
operands.push_back(v0);
|
||||
}
|
||||
if (v1) {
|
||||
operands.push_back(v1);
|
||||
}
|
||||
auto map = AffineMap::get(numDims, numSymbols, {affCombiner(d0, d1)}, {});
|
||||
// TODO: createOrFold when available.
|
||||
return ValueHandle::createComposedAffineApply(map, operands);
|
||||
}
|
||||
|
||||
template <typename IOp, typename FOp>
|
||||
static ValueHandle createBinaryHandle(
|
||||
ValueHandle lhs, ValueHandle rhs,
|
||||
llvm::function_ref<AffineExpr(AffineExpr, AffineExpr)> affCombiner) {
|
||||
auto thisType = lhs.getValue()->getType();
|
||||
auto thatType = rhs.getValue()->getType();
|
||||
assert(thisType == thatType && "cannot mix types in operators");
|
||||
(void)thisType;
|
||||
(void)thatType;
|
||||
if (thisType.isIndex()) {
|
||||
return createBinaryIndexHandle(lhs, rhs, affCombiner);
|
||||
} else if (thisType.isa<IntegerType>()) {
|
||||
return createBinaryHandle<IOp>(lhs, rhs);
|
||||
} else if (thisType.isa<FloatType>()) {
|
||||
return createBinaryHandle<FOp>(lhs, rhs);
|
||||
} else if (auto aggregateType = thisType.dyn_cast<VectorOrTensorType>()) {
|
||||
if (aggregateType.getElementType().isa<IntegerType>())
|
||||
return createBinaryHandle<IOp>(lhs, rhs);
|
||||
else if (aggregateType.getElementType().isa<FloatType>())
|
||||
return createBinaryHandle<FOp>(lhs, rhs);
|
||||
}
|
||||
llvm_unreachable("failed to create a ValueHandle");
|
||||
}
|
||||
|
||||
ValueHandle mlir::edsc::op::operator+(ValueHandle lhs, ValueHandle rhs) {
|
||||
return createBinaryHandle<AddIOp, AddFOp>(
|
||||
lhs, rhs, [](AffineExpr d0, AffineExpr d1) { return d0 + d1; });
|
||||
}
|
||||
|
||||
ValueHandle mlir::edsc::op::operator-(ValueHandle lhs, ValueHandle rhs) {
|
||||
return createBinaryHandle<SubIOp, SubFOp>(
|
||||
lhs, rhs, [](AffineExpr d0, AffineExpr d1) { return d0 - d1; });
|
||||
}
|
||||
|
||||
ValueHandle mlir::edsc::op::operator*(ValueHandle lhs, ValueHandle rhs) {
|
||||
return createBinaryHandle<MulIOp, MulFOp>(
|
||||
lhs, rhs, [](AffineExpr d0, AffineExpr d1) { return d0 * d1; });
|
||||
}
|
||||
|
||||
ValueHandle mlir::edsc::op::operator/(ValueHandle lhs, ValueHandle rhs) {
|
||||
return createBinaryHandle<DivISOp, DivFOp>(
|
||||
lhs, rhs, [](AffineExpr d0, AffineExpr d1) -> AffineExpr {
|
||||
llvm_unreachable("only exprs of non-index type support operator/");
|
||||
});
|
||||
}
|
||||
|
||||
ValueHandle mlir::edsc::op::operator%(ValueHandle lhs, ValueHandle rhs) {
|
||||
return createBinaryHandle<RemISOp, RemFOp>(
|
||||
lhs, rhs, [](AffineExpr d0, AffineExpr d1) { return d0 % d1; });
|
||||
}
|
||||
|
||||
ValueHandle mlir::edsc::op::floorDiv(ValueHandle lhs, ValueHandle rhs) {
|
||||
return createBinaryIndexHandle(
|
||||
lhs, rhs, [](AffineExpr d0, AffineExpr d1) { return d0.floorDiv(d1); });
|
||||
}
|
||||
|
||||
ValueHandle mlir::edsc::op::ceilDiv(ValueHandle lhs, ValueHandle rhs) {
|
||||
return createBinaryIndexHandle(
|
||||
lhs, rhs, [](AffineExpr d0, AffineExpr d1) { return d0.ceilDiv(d1); });
|
||||
}
|
||||
|
|
@ -0,0 +1,30 @@
|
|||
//===- Intrinsics.cpp - MLIR Operations for Declarative Builders *- C++ -*-===//
|
||||
//
|
||||
// Copyright 2019 The MLIR Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
// =============================================================================
|
||||
|
||||
#include "mlir/EDSC/Intrinsics.h"
|
||||
#include "mlir/EDSC/Builders.h"
|
||||
#include "mlir/IR/AffineExpr.h"
|
||||
|
||||
#include "llvm/ADT/STLExtras.h"
|
||||
|
||||
using namespace mlir;
|
||||
using namespace mlir::edsc;
|
||||
|
||||
ValueHandle mlir::edsc::intrinsics::RETURN(ArrayRef<ValueHandle> operands) {
|
||||
SmallVector<Value *, 4> ops(operands.begin(), operands.end());
|
||||
return ValueHandle::create<ReturnOp>(ops);
|
||||
}
|
||||
|
|
@ -233,9 +233,8 @@ TEST_FUNC(max_min_for) {
|
|||
.emitStmt(loop);
|
||||
|
||||
// clang-format off
|
||||
// CHECK-LABEL: func @max_min_for(%arg0: index, %arg1: index, %arg2: index,
|
||||
// %arg3: index) { CHECK: for %i0 = max (d0, d1) -> (d0, d1)(%arg0, %arg1) to
|
||||
// min (d0, d1) -> (d0, d1)(%arg2, %arg3) {
|
||||
// CHECK-LABEL: func @max_min_for(%arg0: index, %arg1: index, %arg2: index, %arg3: index) {
|
||||
// CHECK: for %i0 = max (d0, d1) -> (d0, d1)(%arg0, %arg1) to min (d0, d1) -> (d0, d1)(%arg2, %arg3) {
|
||||
// clang-format on
|
||||
f->print(llvm::outs());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,156 @@
|
|||
//===- builder-api-test.cpp - Tests for Declarative Builder APIs
|
||||
//-----------===//
|
||||
//
|
||||
// Copyright 2019 The MLIR Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
// =============================================================================
|
||||
|
||||
// RUN: %p/builder-api-test | FileCheck %s
|
||||
|
||||
#include "mlir/AffineOps/AffineOps.h"
|
||||
#include "mlir/EDSC/Builders.h"
|
||||
#include "mlir/EDSC/Intrinsics.h"
|
||||
#include "mlir/EDSC/MLIREmitter.h"
|
||||
#include "mlir/EDSC/Types.h"
|
||||
#include "mlir/IR/Builders.h"
|
||||
#include "mlir/IR/MLIRContext.h"
|
||||
#include "mlir/IR/Module.h"
|
||||
#include "mlir/IR/StandardTypes.h"
|
||||
#include "mlir/IR/Types.h"
|
||||
#include "mlir/Pass/Pass.h"
|
||||
#include "mlir/StandardOps/Ops.h"
|
||||
#include "mlir/Transforms/LoopUtils.h"
|
||||
|
||||
#include "Test.h"
|
||||
|
||||
#include "llvm/Support/raw_ostream.h"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
static MLIRContext &globalContext() {
|
||||
static thread_local MLIRContext context;
|
||||
return context;
|
||||
}
|
||||
|
||||
static std::unique_ptr<Function> makeFunction(StringRef name,
|
||||
ArrayRef<Type> results = {},
|
||||
ArrayRef<Type> args = {}) {
|
||||
auto &ctx = globalContext();
|
||||
auto function = llvm::make_unique<Function>(
|
||||
UnknownLoc::get(&ctx), name, FunctionType::get(args, results, &ctx));
|
||||
function->addEntryBlock();
|
||||
return function;
|
||||
}
|
||||
|
||||
TEST_FUNC(builder_dynamic_for_func_args) {
|
||||
using namespace edsc;
|
||||
using namespace edsc::op;
|
||||
using namespace edsc::intrinsics;
|
||||
auto indexType = IndexType::get(&globalContext());
|
||||
auto f32Type = FloatType::getF32(&globalContext());
|
||||
auto f =
|
||||
makeFunction("builder_dynamic_for_func_args", {}, {indexType, indexType});
|
||||
|
||||
ScopedContext scope(f.get());
|
||||
ValueHandle i(indexType), j(indexType), lb(f->getArgument(0)),
|
||||
ub(f->getArgument(1));
|
||||
ValueHandle f7(
|
||||
ValueHandle::create<ConstantFloatOp>(llvm::APFloat(7.0f), f32Type));
|
||||
ValueHandle f13(
|
||||
ValueHandle::create<ConstantFloatOp>(llvm::APFloat(13.0f), f32Type));
|
||||
ValueHandle i7(ValueHandle::create<ConstantIntOp>(7, 32));
|
||||
ValueHandle i13(ValueHandle::create<ConstantIntOp>(13, 32));
|
||||
LoopBuilder(&i, lb, ub, 3)({
|
||||
lb * index_t(3) + ub,
|
||||
lb + index_t(3),
|
||||
LoopBuilder(&j, lb, ub, 2)({
|
||||
ceilDiv(index_t(31) * floorDiv(i + j * index_t(3), index_t(32)),
|
||||
index_t(32)),
|
||||
((f7 + f13) / f7) % f13 - f7 * f13,
|
||||
((i7 + i13) / i7) % i13 - i7 * i13,
|
||||
}),
|
||||
});
|
||||
|
||||
// clang-format off
|
||||
// CHECK-LABEL: func @builder_dynamic_for_func_args(%arg0: index, %arg1: index) {
|
||||
// CHECK: for %i0 = (d0) -> (d0)(%arg0) to (d0) -> (d0)(%arg1) step 3 {
|
||||
// CHECK: {{.*}} = affine.apply (d0) -> (d0 * 3)(%arg0)
|
||||
// CHECK: {{.*}} = affine.apply (d0, d1) -> (d0 * 3 + d1)(%arg0, %arg1)
|
||||
// CHECK: {{.*}} = affine.apply (d0) -> (d0 + 3)(%arg0)
|
||||
// CHECK: for %i1 = (d0) -> (d0)(%arg0) to (d0) -> (d0)(%arg1) step 2 {
|
||||
// CHECK: {{.*}} = affine.apply (d0, d1) -> ((d0 + d1 * 3) floordiv 32)(%i0, %i1)
|
||||
// CHECK: {{.*}} = affine.apply (d0, d1) -> (((d0 + d1 * 3) floordiv 32) * 31)(%i0, %i1)
|
||||
// CHECK: {{.*}} = affine.apply (d0, d1) -> ((((d0 + d1 * 3) floordiv 32) * 31) ceildiv 32)(%i0, %i1)
|
||||
// CHECK: [[rf1:%[0-9]+]] = addf {{.*}}, {{.*}} : f32
|
||||
// CHECK: [[rf2:%[0-9]+]] = divf [[rf1]], {{.*}} : f32
|
||||
// CHECK: [[rf3:%[0-9]+]] = remf [[rf2]], {{.*}} : f32
|
||||
// CHECK: [[rf4:%[0-9]+]] = mulf {{.*}}, {{.*}} : f32
|
||||
// CHECK: {{.*}} = subf [[rf3]], [[rf4]] : f32
|
||||
// CHECK: [[ri1:%[0-9]+]] = addi {{.*}}, {{.*}} : i32
|
||||
// CHECK: [[ri2:%[0-9]+]] = divis [[ri1]], {{.*}} : i32
|
||||
// CHECK: [[ri3:%[0-9]+]] = remis [[ri2]], {{.*}} : i32
|
||||
// CHECK: [[ri4:%[0-9]+]] = muli {{.*}}, {{.*}} : i32
|
||||
// CHECK: {{.*}} = subi [[ri3]], [[ri4]] : i32
|
||||
// clang-format on
|
||||
f->print(llvm::outs());
|
||||
}
|
||||
|
||||
TEST_FUNC(builder_dynamic_for) {
|
||||
using namespace edsc;
|
||||
using namespace edsc::op;
|
||||
using namespace edsc::intrinsics;
|
||||
auto indexType = IndexType::get(&globalContext());
|
||||
auto f = makeFunction("builder_dynamic_for", {},
|
||||
{indexType, indexType, indexType, indexType});
|
||||
|
||||
ScopedContext scope(f.get());
|
||||
ValueHandle i(indexType), a(f->getArgument(0)), b(f->getArgument(1)),
|
||||
c(f->getArgument(2)), d(f->getArgument(3));
|
||||
LoopBuilder(&i, a - b, c + d, 2)({});
|
||||
|
||||
// clang-format off
|
||||
// CHECK-LABEL: func @builder_dynamic_for(%arg0: index, %arg1: index, %arg2: index, %arg3: index) {
|
||||
// CHECK: %0 = affine.apply (d0, d1) -> (d0 - d1)(%arg0, %arg1)
|
||||
// CHECK-NEXT: %1 = affine.apply (d0, d1) -> (d0 + d1)(%arg2, %arg3)
|
||||
// CHECK-NEXT: for %i0 = (d0) -> (d0)(%0) to (d0) -> (d0)(%1) step 2 {
|
||||
// clang-format on
|
||||
f->print(llvm::outs());
|
||||
}
|
||||
|
||||
TEST_FUNC(builder_max_min_for) {
|
||||
using namespace edsc;
|
||||
using namespace edsc::op;
|
||||
using namespace edsc::intrinsics;
|
||||
auto indexType = IndexType::get(&globalContext());
|
||||
auto f = makeFunction("builder_max_min_for", {},
|
||||
{indexType, indexType, indexType, indexType});
|
||||
|
||||
ScopedContext scope(f.get());
|
||||
ValueHandle i(indexType), lb1(f->getArgument(0)), lb2(f->getArgument(1)),
|
||||
ub1(f->getArgument(2)), ub2(f->getArgument(3));
|
||||
LoopBuilder(&i, {lb1, lb2}, {ub1, ub2}, 1)({});
|
||||
RETURN({});
|
||||
|
||||
// clang-format off
|
||||
// CHECK-LABEL: func @builder_max_min_for(%arg0: index, %arg1: index, %arg2: index, %arg3: index) {
|
||||
// CHECK: for %i0 = max (d0, d1) -> (d0, d1)(%arg0, %arg1) to min (d0, d1) -> (d0, d1)(%arg2, %arg3) {
|
||||
// CHECK: return
|
||||
// clang-format on
|
||||
f->print(llvm::outs());
|
||||
}
|
||||
|
||||
int main() {
|
||||
RUN_TESTS();
|
||||
return 0;
|
||||
}
|
||||
Loading…
Reference in New Issue