mlir/lib/EDSC/Builders.cpp - llvm-project - Git at Google

 //===- Builders.cpp - MLIR Declarative Builder Classes --------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//

 #include "mlir/EDSC/Builders.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/IR/AffineExpr.h"

 #include "llvm/ADT/Optional.h"

 using namespace mlir;
 using namespace mlir::edsc;

 mlir::edsc::ScopedContext::ScopedContext(OpBuilder &builder, Location location)
     : builder(builder), location(location),
       enclosingScopedContext(ScopedContext::getCurrentScopedContext()),
       nestedBuilder(nullptr) {
   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 &builder,
                                          OpBuilder::InsertPoint newInsertPt,
                                          Location location)
     : builder(builder), prevBuilderInsertPoint(builder.saveInsertionPoint()),
       location(location),
       enclosingScopedContext(ScopedContext::getCurrentScopedContext()),
       nestedBuilder(nullptr) {
   getCurrentScopedContext() = this;
   builder.restoreInsertionPoint(newInsertPt);
 }

 mlir::edsc::ScopedContext::~ScopedContext() {
   assert(!nestedBuilder &&
          "Active NestedBuilder must have been exited at this point!");
   if (prevBuilderInsertPoint)
     builder.restoreInsertionPoint(*prevBuilderInsertPoint);
   getCurrentScopedContext() = enclosingScopedContext;
 }

 ScopedContext *&mlir::edsc::ScopedContext::getCurrentScopedContext() {
   thread_local ScopedContext *context = nullptr;
   return context;
 }

 OpBuilder &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() {
   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) {
   Operation *op =
       makeComposedAffineApply(ScopedContext::getBuilder(),
                               ScopedContext::getLocation(), map, operands)
           .getOperation();
   assert(op->getNumResults() == 1 && "Not a single result AffineApply");
   return ValueHandle(op->getResult(0));
 }

 ValueHandle ValueHandle::create(StringRef name, ArrayRef<ValueHandle> operands,
                                 ArrayRef<Type> resultTypes,
                                 ArrayRef<NamedAttribute> attributes) {
   Operation *op =
       OperationHandle::create(name, operands, resultTypes, attributes);
   if (op->getNumResults() == 1) {
     return ValueHandle(op->getResult(0));
   }
   if (auto f = dyn_cast<AffineForOp>(op)) {
     return ValueHandle(f.getInductionVar());
   }
   llvm_unreachable("unsupported operation, use an OperationHandle instead");
 }

 OperationHandle OperationHandle::create(StringRef name,
                                         ArrayRef<ValueHandle> operands,
                                         ArrayRef<Type> resultTypes,
                                         ArrayRef<NamedAttribute> attributes) {
   OperationState state(ScopedContext::getLocation(), name);
   SmallVector<Value, 4> ops(operands.begin(), operands.end());
   state.addOperands(ops);
   state.addTypes(resultTypes);
   for (const auto &attr : attributes) {
     state.addAttribute(attr.first, attr.second);
   }
   return OperationHandle(ScopedContext::getBuilder().createOperation(state));
 }

 BlockHandle mlir::edsc::BlockHandle::create(ArrayRef<Type> argTypes) {
   auto &currentB = ScopedContext::getBuilder();
   auto *ib = currentB.getInsertionBlock();
   auto ip = currentB.getInsertionPoint();
   BlockHandle res;
   res.block = ScopedContext::getBuilder().createBlock(ib->getParent());
   // createBlock sets the insertion point inside the block.
   // We do not want this behavior when using declarative builders with nesting.
   currentB.setInsertionPoint(ib, ip);
   for (auto t : argTypes) {
     res.block->addArgument(t);
   }
   return res;
 }

 static Optional<ValueHandle> emitStaticFor(ArrayRef<ValueHandle> lbs,
                                            ArrayRef<ValueHandle> ubs,
                                            int64_t step) {
   if (lbs.size() != 1 || ubs.size() != 1)
     return Optional<ValueHandle>();

   auto *lbDef = lbs.front().getValue().getDefiningOp();
   auto *ubDef = ubs.front().getValue().getDefiningOp();
   if (!lbDef || !ubDef)
     return Optional<ValueHandle>();

   auto lbConst = dyn_cast<ConstantIndexOp>(lbDef);
   auto ubConst = dyn_cast<ConstantIndexOp>(ubDef);
   if (!lbConst || !ubConst)
     return Optional<ValueHandle>();

   return ValueHandle::create<AffineForOp>(lbConst.getValue(),
                                           ubConst.getValue(), step);
 }

 mlir::edsc::LoopBuilder mlir::edsc::LoopBuilder::makeAffine(
     ValueHandle *iv, ArrayRef<ValueHandle> lbHandles,
     ArrayRef<ValueHandle> ubHandles, int64_t step) {
   mlir::edsc::LoopBuilder result;
   if (auto staticFor = emitStaticFor(lbHandles, ubHandles, step)) {
     *iv = staticFor.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();
   result.enter(body, /*prev=*/1);
   return result;
 }

 mlir::edsc::LoopBuilder
 mlir::edsc::LoopBuilder::makeLoop(ValueHandle *iv, ValueHandle lbHandle,
                                   ValueHandle ubHandle,
                                   ValueHandle stepHandle) {
   mlir::edsc::LoopBuilder result;
   auto forOp =
       OperationHandle::createOp<loop::ForOp>(lbHandle, ubHandle, stepHandle);
   *iv = ValueHandle(forOp.getInductionVar());
   auto *body = loop::getForInductionVarOwner(iv->getValue()).getBody();
   result.enter(body, /*prev=*/1);
   return result;
 }

 void mlir::edsc::LoopBuilder::operator()(function_ref<void(void)> fun) {
   // 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,
   ///      }),
   ///    });
   /// ```
   if (fun)
     fun();
   exit();
 }

 mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(
     ValueHandle *iv, ArrayRef<ValueHandle> lbs, ArrayRef<ValueHandle> ubs,
     int64_t step) {
   loops.emplace_back(LoopBuilder::makeAffine(iv, lbs, ubs, step));
 }

 mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(
     ArrayRef<ValueHandle *> ivs, ArrayRef<ValueHandle> lbs,
     ArrayRef<ValueHandle> ubs, ArrayRef<int64_t> steps) {
   assert(ivs.size() == lbs.size() && "Mismatch in number of arguments");
   assert(ivs.size() == ubs.size() && "Mismatch in number of arguments");
   assert(ivs.size() == steps.size() && "Mismatch in number of arguments");
   for (auto it : llvm::zip(ivs, lbs, ubs, steps))
     loops.emplace_back(LoopBuilder::makeAffine(
         std::get<0>(it), std::get<1>(it), std::get<2>(it), std::get<3>(it)));
 }

 void mlir::edsc::AffineLoopNestBuilder::operator()(
     function_ref<void(void)> fun) {
   if (fun)
     fun();
   // Iterate on the calling operator() on all the loops in the nest.
   // The iteration order is from innermost to outermost because enter/exit needs
   // to be asymmetric (i.e. enter() occurs on LoopBuilder construction, exit()
   // occurs on calling operator()). The asymmetry is required for properly
   // nesting imperfectly nested regions (see LoopBuilder::operator()).
   for (auto lit = loops.rbegin(), eit = loops.rend(); lit != eit; ++lit)
     (*lit)();
 }

 mlir::edsc::LoopNestBuilder::LoopNestBuilder(ArrayRef<ValueHandle *> ivs,
                                              ArrayRef<ValueHandle> lbs,
                                              ArrayRef<ValueHandle> ubs,
                                              ArrayRef<ValueHandle> steps) {
   assert(ivs.size() == lbs.size() && "expected size of ivs and lbs to match");
   assert(ivs.size() == ubs.size() && "expected size of ivs and ubs to match");
   assert(ivs.size() == steps.size() &&
          "expected size of ivs and steps to match");
   loops.reserve(ivs.size());
   for (auto it : llvm::zip(ivs, lbs, ubs, steps)) {
     loops.emplace_back(LoopBuilder::makeLoop(std::get<0>(it), std::get<1>(it),
                                              std::get<2>(it), std::get<3>(it)));
   }
   assert(loops.size() == ivs.size() && "Mismatch loops vs ivs size");
 }

 void LoopNestBuilder::LoopNestBuilder::operator()(
     std::function<void(void)> fun) {
   if (fun)
     fun();
   for (auto &lit : reverse(loops))
     lit({});
 }

 mlir::edsc::BlockBuilder::BlockBuilder(BlockHandle bh, Append) {
   assert(bh && "Expected already captured BlockHandle");
   enter(bh.getBlock());
 }

 mlir::edsc::BlockBuilder::BlockBuilder(BlockHandle *bh,
                                        ArrayRef<ValueHandle *> args) {
   assert(!*bh && "BlockHandle already captures a block, use "
                  "the explicit BockBuilder(bh, Append())({}) syntax instead.");
   SmallVector<Type, 8> types;
   for (auto *a : args) {
     assert(!a->hasValue() &&
            "Expected delayed ValueHandle that has not yet captured.");
     types.push_back(a->getType());
   }
   *bh = BlockHandle::create(types);
   for (auto it : llvm::zip(args, bh->getBlock()->getArguments())) {
     *(std::get<0>(it)) = ValueHandle(std::get<1>(it));
   }
   enter(bh->getBlock());
 }

 /// Only serves as an ordering point between entering nested block and creating
 /// stmts.
 void mlir::edsc::BlockBuilder::operator()(function_ref<void(void)> fun) {
   // Call to `exit` must be explicit and asymmetric (cannot happen in the
   // destructor) because of ordering wrt comma operator.
   if (fun)
     fun();
   exit();
 }

 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;
   if (auto constant = dyn_cast_or_null<ConstantIndexOp>(val.getDefiningOp())) {
     d = getAffineConstantExpr(constant.getValue(), context);
   } else if (isValidSymbol(val) && !isValidDim(val)) {
     d = getAffineSymbolExpr(numSymbols++, context);
     resultVal = val;
   } else {
     d = getAffineDimExpr(numDims++, context);
     resultVal = val;
   }
   return std::make_pair(d, resultVal);
 }

 static ValueHandle createBinaryIndexHandle(
     ValueHandle lhs, ValueHandle rhs,
     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,
     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 (thisType.isa<VectorType>() || thisType.isa<TensorType>()) {
     auto aggregateType = thisType.cast<ShapedType>();
     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<SignedDivIOp, 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<SignedRemIOp, 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); });
 }

 ValueHandle mlir::edsc::op::operator!(ValueHandle value) {
   assert(value.getType().isInteger(1) && "expected boolean expression");
   return ValueHandle::create<ConstantIntOp>(1, 1) - value;
 }

 ValueHandle mlir::edsc::op::operator&&(ValueHandle lhs, ValueHandle rhs) {
   assert(lhs.getType().isInteger(1) && "expected boolean expression on LHS");
   assert(rhs.getType().isInteger(1) && "expected boolean expression on RHS");
   return lhs * rhs;
 }

 ValueHandle mlir::edsc::op::operator||(ValueHandle lhs, ValueHandle rhs) {
   return !(!lhs && !rhs);
 }

 static ValueHandle createIComparisonExpr(CmpIPredicate predicate,
                                          ValueHandle lhs, ValueHandle rhs) {
   auto lhsType = lhs.getType();
   auto rhsType = rhs.getType();
   (void)lhsType;
   (void)rhsType;
   assert(lhsType == rhsType && "cannot mix types in operators");
   assert((lhsType.isa<IndexType>() || lhsType.isa<IntegerType>()) &&
          "only integer comparisons are supported");

   auto op = ScopedContext::getBuilder().create<CmpIOp>(
       ScopedContext::getLocation(), predicate, lhs.getValue(), rhs.getValue());
   return ValueHandle(op.getResult());
 }

 static ValueHandle createFComparisonExpr(CmpFPredicate predicate,
                                          ValueHandle lhs, ValueHandle rhs) {
   auto lhsType = lhs.getType();
   auto rhsType = rhs.getType();
   (void)lhsType;
   (void)rhsType;
   assert(lhsType == rhsType && "cannot mix types in operators");
   assert(lhsType.isa<FloatType>() && "only float comparisons are supported");

   auto op = ScopedContext::getBuilder().create<CmpFOp>(
       ScopedContext::getLocation(), predicate, lhs.getValue(), rhs.getValue());
   return ValueHandle(op.getResult());
 }

 // All floating point comparison are ordered through EDSL
 ValueHandle mlir::edsc::op::operator==(ValueHandle lhs, ValueHandle rhs) {
   auto type = lhs.getType();
   return type.isa<FloatType>()
              ? createFComparisonExpr(CmpFPredicate::OEQ, lhs, rhs)
              : createIComparisonExpr(CmpIPredicate::eq, lhs, rhs);
 }
 ValueHandle mlir::edsc::op::operator!=(ValueHandle lhs, ValueHandle rhs) {
   auto type = lhs.getType();
   return type.isa<FloatType>()
              ? createFComparisonExpr(CmpFPredicate::ONE, lhs, rhs)
              : createIComparisonExpr(CmpIPredicate::ne, lhs, rhs);
 }
 ValueHandle mlir::edsc::op::operator<(ValueHandle lhs, ValueHandle rhs) {
   auto type = lhs.getType();
   return type.isa<FloatType>()
              ? createFComparisonExpr(CmpFPredicate::OLT, lhs, rhs)
              :
              // TODO(ntv,zinenko): signed by default, how about unsigned?
              createIComparisonExpr(CmpIPredicate::slt, lhs, rhs);
 }
 ValueHandle mlir::edsc::op::operator<=(ValueHandle lhs, ValueHandle rhs) {
   auto type = lhs.getType();
   return type.isa<FloatType>()
              ? createFComparisonExpr(CmpFPredicate::OLE, lhs, rhs)
              : createIComparisonExpr(CmpIPredicate::sle, lhs, rhs);
 }
 ValueHandle mlir::edsc::op::operator>(ValueHandle lhs, ValueHandle rhs) {
   auto type = lhs.getType();
   return type.isa<FloatType>()
              ? createFComparisonExpr(CmpFPredicate::OGT, lhs, rhs)
              : createIComparisonExpr(CmpIPredicate::sgt, lhs, rhs);
 }
 ValueHandle mlir::edsc::op::operator>=(ValueHandle lhs, ValueHandle rhs) {
   auto type = lhs.getType();
   return type.isa<FloatType>()
              ? createFComparisonExpr(CmpFPredicate::OGE, lhs, rhs)
              : createIComparisonExpr(CmpIPredicate::sge, lhs, rhs);
 }
	//===- Builders.cpp - MLIR Declarative Builder Classes --------------------===//
	//
	// 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
	//
	//===----------------------------------------------------------------------===//

	#include "mlir/EDSC/Builders.h"
	#include "mlir/Dialect/StandardOps/Ops.h"
	#include "mlir/IR/AffineExpr.h"

	#include "llvm/ADT/Optional.h"

	using namespace mlir;
	using namespace mlir::edsc;

	mlir::edsc::ScopedContext::ScopedContext(OpBuilder &builder, Location location)
	: builder(builder), location(location),
	enclosingScopedContext(ScopedContext::getCurrentScopedContext()),
	nestedBuilder(nullptr) {
	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 &builder,
	OpBuilder::InsertPoint newInsertPt,
	Location location)
	: builder(builder), prevBuilderInsertPoint(builder.saveInsertionPoint()),
	location(location),
	enclosingScopedContext(ScopedContext::getCurrentScopedContext()),
	nestedBuilder(nullptr) {
	getCurrentScopedContext() = this;
	builder.restoreInsertionPoint(newInsertPt);
	}

	mlir::edsc::ScopedContext::~ScopedContext() {
	assert(!nestedBuilder &&
	"Active NestedBuilder must have been exited at this point!");
	if (prevBuilderInsertPoint)
	builder.restoreInsertionPoint(*prevBuilderInsertPoint);
	getCurrentScopedContext() = enclosingScopedContext;
	}

	ScopedContext *&mlir::edsc::ScopedContext::getCurrentScopedContext() {
	thread_local ScopedContext *context = nullptr;
	return context;
	}

	OpBuilder &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() {
	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) {
	Operation *op =
	makeComposedAffineApply(ScopedContext::getBuilder(),
	ScopedContext::getLocation(), map, operands)
	.getOperation();
	assert(op->getNumResults() == 1 && "Not a single result AffineApply");
	return ValueHandle(op->getResult(0));
	}

	ValueHandle ValueHandle::create(StringRef name, ArrayRef<ValueHandle> operands,
	ArrayRef<Type> resultTypes,
	ArrayRef<NamedAttribute> attributes) {
	Operation *op =
	OperationHandle::create(name, operands, resultTypes, attributes);
	if (op->getNumResults() == 1) {
	return ValueHandle(op->getResult(0));
	}
	if (auto f = dyn_cast<AffineForOp>(op)) {
	return ValueHandle(f.getInductionVar());
	}
	llvm_unreachable("unsupported operation, use an OperationHandle instead");
	}

	OperationHandle OperationHandle::create(StringRef name,
	ArrayRef<ValueHandle> operands,
	ArrayRef<Type> resultTypes,
	ArrayRef<NamedAttribute> attributes) {
	OperationState state(ScopedContext::getLocation(), name);
	SmallVector<Value, 4> ops(operands.begin(), operands.end());
	state.addOperands(ops);
	state.addTypes(resultTypes);
	for (const auto &attr : attributes) {
	state.addAttribute(attr.first, attr.second);
	}
	return OperationHandle(ScopedContext::getBuilder().createOperation(state));
	}

	BlockHandle mlir::edsc::BlockHandle::create(ArrayRef<Type> argTypes) {
	auto &currentB = ScopedContext::getBuilder();
	auto *ib = currentB.getInsertionBlock();
	auto ip = currentB.getInsertionPoint();
	BlockHandle res;
	res.block = ScopedContext::getBuilder().createBlock(ib->getParent());
	// createBlock sets the insertion point inside the block.
	// We do not want this behavior when using declarative builders with nesting.
	currentB.setInsertionPoint(ib, ip);
	for (auto t : argTypes) {
	res.block->addArgument(t);
	}
	return res;
	}

	static Optional<ValueHandle> emitStaticFor(ArrayRef<ValueHandle> lbs,
	ArrayRef<ValueHandle> ubs,
	int64_t step) {
	if (lbs.size() != 1 \|\| ubs.size() != 1)
	return Optional<ValueHandle>();

	auto *lbDef = lbs.front().getValue().getDefiningOp();
	auto *ubDef = ubs.front().getValue().getDefiningOp();
	if (!lbDef \|\| !ubDef)
	return Optional<ValueHandle>();

	auto lbConst = dyn_cast<ConstantIndexOp>(lbDef);
	auto ubConst = dyn_cast<ConstantIndexOp>(ubDef);
	if (!lbConst \|\| !ubConst)
	return Optional<ValueHandle>();

	return ValueHandle::create<AffineForOp>(lbConst.getValue(),
	ubConst.getValue(), step);
	}

	mlir::edsc::LoopBuilder mlir::edsc::LoopBuilder::makeAffine(
	ValueHandle *iv, ArrayRef<ValueHandle> lbHandles,
	ArrayRef<ValueHandle> ubHandles, int64_t step) {
	mlir::edsc::LoopBuilder result;
	if (auto staticFor = emitStaticFor(lbHandles, ubHandles, step)) {
	*iv = staticFor.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();
	result.enter(body, /prev=/1);
	return result;
	}

	mlir::edsc::LoopBuilder
	mlir::edsc::LoopBuilder::makeLoop(ValueHandle *iv, ValueHandle lbHandle,
	ValueHandle ubHandle,
	ValueHandle stepHandle) {
	mlir::edsc::LoopBuilder result;
	auto forOp =
	OperationHandle::createOp<loop::ForOp>(lbHandle, ubHandle, stepHandle);
	*iv = ValueHandle(forOp.getInductionVar());
	auto *body = loop::getForInductionVarOwner(iv->getValue()).getBody();
	result.enter(body, /prev=/1);
	return result;
	}

	void mlir::edsc::LoopBuilder::operator()(function_ref<void(void)> fun) {
	// 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,
	/// }),
	/// });
	/// ```
	if (fun)
	fun();
	exit();
	}

	mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(
	ValueHandle *iv, ArrayRef<ValueHandle> lbs, ArrayRef<ValueHandle> ubs,
	int64_t step) {
	loops.emplace_back(LoopBuilder::makeAffine(iv, lbs, ubs, step));
	}

	mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(
	ArrayRef<ValueHandle *> ivs, ArrayRef<ValueHandle> lbs,
	ArrayRef<ValueHandle> ubs, ArrayRef<int64_t> steps) {
	assert(ivs.size() == lbs.size() && "Mismatch in number of arguments");
	assert(ivs.size() == ubs.size() && "Mismatch in number of arguments");
	assert(ivs.size() == steps.size() && "Mismatch in number of arguments");
	for (auto it : llvm::zip(ivs, lbs, ubs, steps))
	loops.emplace_back(LoopBuilder::makeAffine(
	std::get<0>(it), std::get<1>(it), std::get<2>(it), std::get<3>(it)));
	}

	void mlir::edsc::AffineLoopNestBuilder::operator()(
	function_ref<void(void)> fun) {
	if (fun)
	fun();
	// Iterate on the calling operator() on all the loops in the nest.
	// The iteration order is from innermost to outermost because enter/exit needs
	// to be asymmetric (i.e. enter() occurs on LoopBuilder construction, exit()
	// occurs on calling operator()). The asymmetry is required for properly
	// nesting imperfectly nested regions (see LoopBuilder::operator()).
	for (auto lit = loops.rbegin(), eit = loops.rend(); lit != eit; ++lit)
	(*lit)();
	}

	mlir::edsc::LoopNestBuilder::LoopNestBuilder(ArrayRef<ValueHandle *> ivs,
	ArrayRef<ValueHandle> lbs,
	ArrayRef<ValueHandle> ubs,
	ArrayRef<ValueHandle> steps) {
	assert(ivs.size() == lbs.size() && "expected size of ivs and lbs to match");
	assert(ivs.size() == ubs.size() && "expected size of ivs and ubs to match");
	assert(ivs.size() == steps.size() &&
	"expected size of ivs and steps to match");
	loops.reserve(ivs.size());
	for (auto it : llvm::zip(ivs, lbs, ubs, steps)) {
	loops.emplace_back(LoopBuilder::makeLoop(std::get<0>(it), std::get<1>(it),
	std::get<2>(it), std::get<3>(it)));
	}
	assert(loops.size() == ivs.size() && "Mismatch loops vs ivs size");
	}

	void LoopNestBuilder::LoopNestBuilder::operator()(
	std::function<void(void)> fun) {
	if (fun)
	fun();
	for (auto &lit : reverse(loops))
	lit({});
	}

	mlir::edsc::BlockBuilder::BlockBuilder(BlockHandle bh, Append) {
	assert(bh && "Expected already captured BlockHandle");
	enter(bh.getBlock());
	}

	mlir::edsc::BlockBuilder::BlockBuilder(BlockHandle *bh,
	ArrayRef<ValueHandle *> args) {
	assert(!*bh && "BlockHandle already captures a block, use "
	"the explicit BockBuilder(bh, Append())({}) syntax instead.");
	SmallVector<Type, 8> types;
	for (auto *a : args) {
	assert(!a->hasValue() &&
	"Expected delayed ValueHandle that has not yet captured.");
	types.push_back(a->getType());
	}
	*bh = BlockHandle::create(types);
	for (auto it : llvm::zip(args, bh->getBlock()->getArguments())) {
	*(std::get<0>(it)) = ValueHandle(std::get<1>(it));
	}
	enter(bh->getBlock());
	}

	/// Only serves as an ordering point between entering nested block and creating
	/// stmts.
	void mlir::edsc::BlockBuilder::operator()(function_ref<void(void)> fun) {
	// Call to `exit` must be explicit and asymmetric (cannot happen in the
	// destructor) because of ordering wrt comma operator.
	if (fun)
	fun();
	exit();
	}

	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;
	if (auto constant = dyn_cast_or_null<ConstantIndexOp>(val.getDefiningOp())) {
	d = getAffineConstantExpr(constant.getValue(), context);
	} else if (isValidSymbol(val) && !isValidDim(val)) {
	d = getAffineSymbolExpr(numSymbols++, context);
	resultVal = val;
	} else {
	d = getAffineDimExpr(numDims++, context);
	resultVal = val;
	}
	return std::make_pair(d, resultVal);
	}

	static ValueHandle createBinaryIndexHandle(
	ValueHandle lhs, ValueHandle rhs,
	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,
	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 (thisType.isa<VectorType>() \|\| thisType.isa<TensorType>()) {
	auto aggregateType = thisType.cast<ShapedType>();
	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<SignedDivIOp, 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<SignedRemIOp, 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); });
	}

	ValueHandle mlir::edsc::op::operator!(ValueHandle value) {
	assert(value.getType().isInteger(1) && "expected boolean expression");
	return ValueHandle::create<ConstantIntOp>(1, 1) - value;
	}

	ValueHandle mlir::edsc::op::operator&&(ValueHandle lhs, ValueHandle rhs) {
	assert(lhs.getType().isInteger(1) && "expected boolean expression on LHS");
	assert(rhs.getType().isInteger(1) && "expected boolean expression on RHS");
	return lhs * rhs;
	}

	ValueHandle mlir::edsc::op::operator\|\|(ValueHandle lhs, ValueHandle rhs) {
	return !(!lhs && !rhs);
	}

	static ValueHandle createIComparisonExpr(CmpIPredicate predicate,
	ValueHandle lhs, ValueHandle rhs) {
	auto lhsType = lhs.getType();
	auto rhsType = rhs.getType();
	(void)lhsType;
	(void)rhsType;
	assert(lhsType == rhsType && "cannot mix types in operators");
	assert((lhsType.isa<IndexType>() \|\| lhsType.isa<IntegerType>()) &&
	"only integer comparisons are supported");

	auto op = ScopedContext::getBuilder().create<CmpIOp>(
	ScopedContext::getLocation(), predicate, lhs.getValue(), rhs.getValue());
	return ValueHandle(op.getResult());
	}

	static ValueHandle createFComparisonExpr(CmpFPredicate predicate,
	ValueHandle lhs, ValueHandle rhs) {
	auto lhsType = lhs.getType();
	auto rhsType = rhs.getType();
	(void)lhsType;
	(void)rhsType;
	assert(lhsType == rhsType && "cannot mix types in operators");
	assert(lhsType.isa<FloatType>() && "only float comparisons are supported");

	auto op = ScopedContext::getBuilder().create<CmpFOp>(
	ScopedContext::getLocation(), predicate, lhs.getValue(), rhs.getValue());
	return ValueHandle(op.getResult());
	}

	// All floating point comparison are ordered through EDSL
	ValueHandle mlir::edsc::op::operator==(ValueHandle lhs, ValueHandle rhs) {
	auto type = lhs.getType();
	return type.isa<FloatType>()
	? createFComparisonExpr(CmpFPredicate::OEQ, lhs, rhs)
	: createIComparisonExpr(CmpIPredicate::eq, lhs, rhs);
	}
	ValueHandle mlir::edsc::op::operator!=(ValueHandle lhs, ValueHandle rhs) {
	auto type = lhs.getType();
	return type.isa<FloatType>()
	? createFComparisonExpr(CmpFPredicate::ONE, lhs, rhs)
	: createIComparisonExpr(CmpIPredicate::ne, lhs, rhs);
	}
	ValueHandle mlir::edsc::op::operator<(ValueHandle lhs, ValueHandle rhs) {
	auto type = lhs.getType();
	return type.isa<FloatType>()
	? createFComparisonExpr(CmpFPredicate::OLT, lhs, rhs)
	:
	// TODO(ntv,zinenko): signed by default, how about unsigned?
	createIComparisonExpr(CmpIPredicate::slt, lhs, rhs);
	}
	ValueHandle mlir::edsc::op::operator<=(ValueHandle lhs, ValueHandle rhs) {
	auto type = lhs.getType();
	return type.isa<FloatType>()
	? createFComparisonExpr(CmpFPredicate::OLE, lhs, rhs)
	: createIComparisonExpr(CmpIPredicate::sle, lhs, rhs);
	}
	ValueHandle mlir::edsc::op::operator>(ValueHandle lhs, ValueHandle rhs) {
	auto type = lhs.getType();
	return type.isa<FloatType>()
	? createFComparisonExpr(CmpFPredicate::OGT, lhs, rhs)
	: createIComparisonExpr(CmpIPredicate::sgt, lhs, rhs);
	}
	ValueHandle mlir::edsc::op::operator>=(ValueHandle lhs, ValueHandle rhs) {
	auto type = lhs.getType();
	return type.isa<FloatType>()
	? createFComparisonExpr(CmpFPredicate::OGE, lhs, rhs)
	: createIComparisonExpr(CmpIPredicate::sge, lhs, rhs);
	}