forked from OSchip/llvm-project
117 lines
4.4 KiB
C++
117 lines
4.4 KiB
C++
//===- Builders.cpp - MLIR Declarative Builder Classes --------------------===//
|
|
//
|
|
// Part of the LLVM 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/EDSC/Builders.h"
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/IR/AffineMap.h"
|
|
|
|
#include "llvm/ADT/Optional.h"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::edsc;
|
|
|
|
mlir::edsc::ScopedContext::ScopedContext(OpBuilder &b)
|
|
: ScopedContext(b, b.getInsertionPoint()->getLoc()) {}
|
|
|
|
mlir::edsc::ScopedContext::ScopedContext(OpBuilder &b, Location location)
|
|
: builder(b), guard(builder), location(location),
|
|
enclosingScopedContext(ScopedContext::getCurrentScopedContext()) {
|
|
getCurrentScopedContext() = this;
|
|
}
|
|
|
|
/// Sets the insertion point of the builder to 'newInsertPt' for the duration
|
|
/// of the scope. The existing insertion point of the builder is restored on
|
|
/// destruction.
|
|
mlir::edsc::ScopedContext::ScopedContext(OpBuilder &b,
|
|
OpBuilder::InsertPoint newInsertPt,
|
|
Location location)
|
|
: builder(b), guard(builder), location(location),
|
|
enclosingScopedContext(ScopedContext::getCurrentScopedContext()) {
|
|
getCurrentScopedContext() = this;
|
|
builder.restoreInsertionPoint(newInsertPt);
|
|
}
|
|
|
|
mlir::edsc::ScopedContext::~ScopedContext() {
|
|
getCurrentScopedContext() = enclosingScopedContext;
|
|
}
|
|
|
|
ScopedContext *&mlir::edsc::ScopedContext::getCurrentScopedContext() {
|
|
thread_local ScopedContext *context = nullptr;
|
|
return context;
|
|
}
|
|
|
|
OpBuilder &mlir::edsc::ScopedContext::getBuilderRef() {
|
|
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() {
|
|
return getBuilderRef().getContext();
|
|
}
|
|
|
|
Block *mlir::edsc::createBlock(TypeRange argTypes) {
|
|
assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
|
|
OpBuilder &builder = ScopedContext::getBuilderRef();
|
|
Block *block = builder.getInsertionBlock();
|
|
assert(block != nullptr &&
|
|
"insertion point not set up in the builder within ScopedContext");
|
|
|
|
return createBlockInRegion(*block->getParent(), argTypes);
|
|
}
|
|
|
|
Block *mlir::edsc::createBlockInRegion(Region ®ion, TypeRange argTypes) {
|
|
assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
|
|
OpBuilder &builder = ScopedContext::getBuilderRef();
|
|
|
|
OpBuilder::InsertionGuard guard(builder);
|
|
return builder.createBlock(®ion, {}, argTypes);
|
|
}
|
|
|
|
void mlir::edsc::appendToBlock(Block *block,
|
|
function_ref<void(ValueRange)> builderFn) {
|
|
assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
|
|
OpBuilder &builder = ScopedContext::getBuilderRef();
|
|
|
|
OpBuilder::InsertionGuard guard(builder);
|
|
if (block->empty() || !block->back().mightHaveTrait<OpTrait::IsTerminator>())
|
|
builder.setInsertionPointToEnd(block);
|
|
else
|
|
builder.setInsertionPoint(&block->back());
|
|
builderFn(block->getArguments());
|
|
}
|
|
|
|
Block *mlir::edsc::buildInNewBlock(TypeRange argTypes,
|
|
function_ref<void(ValueRange)> builderFn) {
|
|
assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
|
|
OpBuilder &builder = ScopedContext::getBuilderRef();
|
|
Block *block = builder.getInsertionBlock();
|
|
assert(block != nullptr &&
|
|
"insertion point not set up in the builder within ScopedContext");
|
|
return buildInNewBlock(*block->getParent(), argTypes, builderFn);
|
|
}
|
|
|
|
Block *mlir::edsc::buildInNewBlock(Region ®ion, TypeRange argTypes,
|
|
function_ref<void(ValueRange)> builderFn) {
|
|
assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
|
|
OpBuilder &builder = ScopedContext::getBuilderRef();
|
|
|
|
Block *block = createBlockInRegion(region, argTypes);
|
|
OpBuilder::InsertionGuard guard(builder);
|
|
builder.setInsertionPointToStart(block);
|
|
builderFn(block->getArguments());
|
|
return block;
|
|
}
|
|
|