maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir][Affine] NFC - Drop Affine EDSC usage 

Drop the Affine dialect EDSC subdirectory and update all uses.

Differential Revision: https://reviews.llvm.org/D102878
This commit is contained in:
Nicolas Vasilache 2021-05-20 21:35:25 +00:00
parent 33b71ec9c6
commit e84a9b9bb3
14 changed files with 57 additions and 491 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

95
mlir/include/mlir/Dialect/Affine/EDSC/Builders.h
View File

@ -1,95 +0,0 @@
//===- Builders.h - MLIR Declarative Builder Classes ------------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Provides intuitive composable interfaces for building structured MLIR
// snippets in a declarative fashion.
//
//===----------------------------------------------------------------------===//
#ifndef MLIR_DIALECT_AFFINE_EDSC_BUILDERS_H_
#define MLIR_DIALECT_AFFINE_EDSC_BUILDERS_H_
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/EDSC/Builders.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Types.h"
namespace mlir {
namespace edsc {
/// Creates a perfect nest of affine "for" loops, given the list of lower
/// bounds, upper bounds and steps. The three lists are expected to contain the
/// same number of elements. Uses the OpBuilder and Location stored in
/// ScopedContext and assumes they are non-null. The optional "bodyBuilderFn"
/// callback is called to construct the body of the innermost loop and is passed
/// the list of loop induction variables, in order from outermost to innermost.
/// The function is expected to use the builder and location stored in
/// ScopedContext at the moment of the call. The function should not create
/// the affine terminator op, which will be added regardless of the
/// "bodyBuilderFn" being present.
void affineLoopNestBuilder(
ValueRange lbs, ValueRange ubs, ArrayRef<int64_t> steps,
function_ref<void(ValueRange)> bodyBuilderFn = nullptr);
/// Creates a single affine "for" loop, iterating from max(lbs) to min(ubs) with
/// the given step. Uses the OpBuilder and Location stored in ScopedContext and
/// assumes they are non-null. The optional "bodyBuilderFn" callback is called
/// to construct the body of the loop and is passed the induction variable. The
/// function is expected to use the builder and location stored in ScopedContext
/// at the moment of the call. The function should not create the affine
/// terminator op, which will be added regardless of the "bodyBuilderFn" being
/// present.
void affineLoopBuilder(ValueRange lbs, ValueRange ubs, int64_t step,
function_ref<void(Value)> bodyBuilderFn = nullptr);
/// Creates a single affine "for" loop, iterating from max(lbs) to min(ubs) with
/// the given step. Uses the OpBuilder and Location stored in ScopedContext and
/// assumes they are non-null. "iterArgs" is used to specify the initial values
/// of the result affine "for" might yield. The optional "bodyBuilderFn"
/// callback is called to construct the body of the loop and is passed the
/// induction variable and the iteration arguments. The function is expected to
/// use the builder and location stored in ScopedContext at the moment of the
/// call. The function will create the affine terminator op in case "iterArgs"
/// is empty and "bodyBuilderFn" is not present.
void affineLoopBuilder(
ValueRange lbs, ValueRange ubs, int64_t step, ValueRange iterArgs,
function_ref<void(Value, ValueRange)> bodyBuilderFn = nullptr);
namespace op {
Value operator+(Value lhs, Value rhs);
Value operator-(Value lhs, Value rhs);
Value operator*(Value lhs, Value rhs);
Value operator/(Value lhs, Value rhs);
Value operator%(Value lhs, Value rhs);
Value floorDiv(Value lhs, Value rhs);
Value ceilDiv(Value lhs, Value rhs);
/// Logical operator overloadings.
Value negate(Value value);
Value operator&&(Value lhs, Value rhs);
Value operator||(Value lhs, Value rhs);
Value operator^(Value lhs, Value rhs);
/// Comparison operator overloadings.
Value eq(Value lhs, Value rhs);
Value ne(Value lhs, Value rhs);
Value slt(Value lhs, Value rhs);
Value sle(Value lhs, Value rhs);
Value sgt(Value lhs, Value rhs);
Value sge(Value lhs, Value rhs);
Value ult(Value lhs, Value rhs);
Value ule(Value lhs, Value rhs);
Value ugt(Value lhs, Value rhs);
Value uge(Value lhs, Value rhs);
} // namespace op
} // namespace edsc
} // namespace mlir
#endif // MLIR_DIALECT_AFFINE_EDSC_BUILDERS_H_

28
mlir/include/mlir/Dialect/Affine/EDSC/Intrinsics.h
View File

@ -1,28 +0,0 @@
//===- Intrinsics.h - MLIR EDSC Intrinsics for AffineOps --------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
#ifndef MLIR_DIALECT_AFFINE_EDSC_INTRINSICS_H_
#define MLIR_DIALECT_AFFINE_EDSC_INTRINSICS_H_
#include "mlir/Dialect/Affine/EDSC/Builders.h"
namespace mlir {
namespace edsc {
namespace intrinsics {
using affine_apply = ValueBuilder<AffineApplyOp>;
using affine_if = OperationBuilder<AffineIfOp>;
using affine_load = ValueBuilder<AffineLoadOp>;
using affine_min = ValueBuilder<AffineMinOp>;
using affine_max = ValueBuilder<AffineMaxOp>;
using affine_store = OperationBuilder<AffineStoreOp>;
} // namespace intrinsics
} // namespace edsc
} // namespace mlir
#endif // MLIR_DIALECT_STANDARDOPS_EDSC_INTRINSICS_H_

1
mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
View File

@ -9,7 +9,6 @@
#ifndef MLIR_DIALECT_LINALG_UTILS_H_
#define MLIR_DIALECT_LINALG_UTILS_H_
#include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
#include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/SCF/SCF.h"

16
mlir/include/mlir/Dialect/StandardOps/Utils/Utils.h
View File

@ -72,6 +72,22 @@ public:
}
};
/// Helper struct to build simple arithmetic quantities with minimal type
/// inference support.
struct ArithBuilder {
ArithBuilder(OpBuilder &b, Location loc) : b(b), loc(loc) {}
Value _and(Value lhs, Value rhs);
Value add(Value lhs, Value rhs);
Value mul(Value lhs, Value rhs);
Value select(Value cmp, Value lhs, Value rhs);
Value sgt(Value lhs, Value rhs);
Value slt(Value lhs, Value rhs);
private:
OpBuilder &b;
Location loc;
};
} // end namespace mlir
#endif // MLIR_DIALECT_STANDARDOPS_UTILS_UTILS_H

1
mlir/lib/Conversion/VectorToSCF/CMakeLists.txt
View File

@ -9,7 +9,6 @@ add_mlir_conversion_library(MLIRVectorToSCF
LINK_LIBS PUBLIC
MLIREDSC
MLIRAffineEDSC
MLIRLLVMIR
MLIRMemRef
MLIRTransforms

1
mlir/lib/Dialect/Affine/CMakeLists.txt
View File

@ -1,4 +1,3 @@
add_subdirectory(IR)
add_subdirectory(EDSC)
add_subdirectory(Transforms)
add_subdirectory(Utils)

296
mlir/lib/Dialect/Affine/EDSC/Builders.cpp
View File

@ -1,296 +0,0 @@
//===- 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/Dialect/Affine/EDSC/Builders.h"
#include "mlir/Dialect/StandardOps/EDSC/Builders.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
using namespace mlir;
using namespace mlir::edsc;
void mlir::edsc::affineLoopNestBuilder(
ValueRange lbs, ValueRange ubs, ArrayRef<int64_t> steps,
function_ref<void(ValueRange)> bodyBuilderFn) {
assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
// Wrap the body builder function into an interface compatible with the main
// builder.
auto wrappedBuilderFn = [&](OpBuilder &nestedBuilder, Location nestedLoc,
ValueRange ivs) {
ScopedContext context(nestedBuilder, nestedLoc);
bodyBuilderFn(ivs);
};
function_ref<void(OpBuilder &, Location, ValueRange)> wrapper;
if (bodyBuilderFn)
wrapper = wrappedBuilderFn;
// Extract the builder, location and construct the loop nest.
OpBuilder &builder = ScopedContext::getBuilderRef();
Location loc = ScopedContext::getLocation();
buildAffineLoopNest(builder, loc, lbs, ubs, steps, wrapper);
}
void mlir::edsc::affineLoopBuilder(ValueRange lbs, ValueRange ubs, int64_t step,
function_ref<void(Value)> bodyBuilderFn) {
// Fetch the builder and location.
assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
OpBuilder &builder = ScopedContext::getBuilderRef();
Location loc = ScopedContext::getLocation();
// Create the actual loop and call the body builder, if provided, after
// updating the scoped context.
builder.create<AffineForOp>(
loc, lbs, builder.getMultiDimIdentityMap(lbs.size()), ubs,
builder.getMultiDimIdentityMap(ubs.size()), step, llvm::None,
[&](OpBuilder &nestedBuilder, Location nestedLoc, Value iv,
ValueRange itrArgs) {
if (bodyBuilderFn) {
ScopedContext nestedContext(nestedBuilder, nestedLoc);
OpBuilder::InsertionGuard guard(nestedBuilder);
bodyBuilderFn(iv);
}
nestedBuilder.create<AffineYieldOp>(nestedLoc);
});
}
void mlir::edsc::affineLoopBuilder(
ValueRange lbs, ValueRange ubs, int64_t step, ValueRange iterArgs,
function_ref<void(Value, ValueRange)> bodyBuilderFn) {
// Fetch the builder and location.
assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
OpBuilder &builder = ScopedContext::getBuilderRef();
Location loc = ScopedContext::getLocation();
// Create the actual loop and call the body builder, if provided, after
// updating the scoped context.
builder.create<AffineForOp>(
loc, lbs, builder.getMultiDimIdentityMap(lbs.size()), ubs,
builder.getMultiDimIdentityMap(ubs.size()), step, iterArgs,
[&](OpBuilder &nestedBuilder, Location nestedLoc, Value iv,
ValueRange itrArgs) {
if (bodyBuilderFn) {
ScopedContext nestedContext(nestedBuilder, nestedLoc);
OpBuilder::InsertionGuard guard(nestedBuilder);
bodyBuilderFn(iv, itrArgs);
} else if (itrArgs.empty())
nestedBuilder.create<AffineYieldOp>(nestedLoc);
});
}
static std::pair<AffineExpr, Value>
categorizeValueByAffineType(MLIRContext *context, Value val, unsigned &numDims,
unsigned &numSymbols) {
AffineExpr d;
Value resultVal = nullptr;
if (auto constant = val.getDefiningOp<ConstantIndexOp>()) {
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 Value createBinaryIndexHandle(
Value lhs, Value 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, numDims, numSymbols);
std::tie(d1, v1) =
categorizeValueByAffineType(context, rhs, 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.
Operation *op =
makeComposedAffineApply(ScopedContext::getBuilderRef(),
ScopedContext::getLocation(), map, operands)
.getOperation();
assert(op->getNumResults() == 1 && "Expected single result AffineApply");
return op->getResult(0);
}
template <typename IOp, typename FOp>
static Value createBinaryHandle(
Value lhs, Value rhs,
function_ref<AffineExpr(AffineExpr, AffineExpr)> affCombiner) {
auto thisType = lhs.getType();
auto thatType = rhs.getType();
assert(thisType == thatType && "cannot mix types in operators");
(void)thisType;
(void)thatType;
if (thisType.isIndex()) {
return createBinaryIndexHandle(lhs, rhs, affCombiner);
} else if (thisType.isSignlessInteger()) {
return ValueBuilder<IOp>(lhs, rhs);
} else if (thisType.isa<FloatType>()) {
return ValueBuilder<FOp>(lhs, rhs);
} else if (thisType.isa<VectorType, TensorType>()) {
auto aggregateType = thisType.cast<ShapedType>();
if (aggregateType.getElementType().isSignlessInteger())
return ValueBuilder<IOp>(lhs, rhs);
else if (aggregateType.getElementType().isa<FloatType>())
return ValueBuilder<FOp>(lhs, rhs);
}
llvm_unreachable("failed to create a Value");
}
Value mlir::edsc::op::operator+(Value lhs, Value rhs) {
return createBinaryHandle<AddIOp, AddFOp>(
lhs, rhs, [](AffineExpr d0, AffineExpr d1) { return d0 + d1; });
}
Value mlir::edsc::op::operator-(Value lhs, Value rhs) {
return createBinaryHandle<SubIOp, SubFOp>(
lhs, rhs, [](AffineExpr d0, AffineExpr d1) { return d0 - d1; });
}
Value mlir::edsc::op::operator*(Value lhs, Value rhs) {
return createBinaryHandle<MulIOp, MulFOp>(
lhs, rhs, [](AffineExpr d0, AffineExpr d1) { return d0 * d1; });
}
Value mlir::edsc::op::operator/(Value lhs, Value rhs) {
return createBinaryHandle<SignedDivIOp, DivFOp>(
lhs, rhs, [](AffineExpr d0, AffineExpr d1) -> AffineExpr {
llvm_unreachable("only exprs of non-index type support operator/");
});
}
Value mlir::edsc::op::operator%(Value lhs, Value rhs) {
return createBinaryHandle<SignedRemIOp, RemFOp>(
lhs, rhs, [](AffineExpr d0, AffineExpr d1) { return d0 % d1; });
}
Value mlir::edsc::op::floorDiv(Value lhs, Value rhs) {
return createBinaryIndexHandle(
lhs, rhs, [](AffineExpr d0, AffineExpr d1) { return d0.floorDiv(d1); });
}
Value mlir::edsc::op::ceilDiv(Value lhs, Value rhs) {
return createBinaryIndexHandle(
lhs, rhs, [](AffineExpr d0, AffineExpr d1) { return d0.ceilDiv(d1); });
}
Value mlir::edsc::op::negate(Value value) {
assert(value.getType().isInteger(1) && "expected boolean expression");
return ValueBuilder<ConstantIntOp>(1, 1) - value;
}
Value mlir::edsc::op::operator&&(Value lhs, Value rhs) {
assert(lhs.getType().isInteger(1) && "expected boolean expression on LHS");
assert(rhs.getType().isInteger(1) && "expected boolean expression on RHS");
return ValueBuilder<AndOp>(lhs, rhs);
}
Value mlir::edsc::op::operator||(Value lhs, Value rhs) {
assert(lhs.getType().isInteger(1) && "expected boolean expression on LHS");
assert(rhs.getType().isInteger(1) && "expected boolean expression on RHS");
return ValueBuilder<OrOp>(lhs, rhs);
}
static Value createIComparisonExpr(CmpIPredicate predicate, Value lhs,
Value 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.isSignlessInteger()) &&
"only integer comparisons are supported");
return ScopedContext::getBuilderRef().create<CmpIOp>(
ScopedContext::getLocation(), predicate, lhs, rhs);
}
static Value createFComparisonExpr(CmpFPredicate predicate, Value lhs,
Value 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");
return ScopedContext::getBuilderRef().create<CmpFOp>(
ScopedContext::getLocation(), predicate, lhs, rhs);
}
// All floating point comparison are ordered through EDSL
Value mlir::edsc::op::eq(Value lhs, Value rhs) {
auto type = lhs.getType();
return type.isa<FloatType>()
? createFComparisonExpr(CmpFPredicate::OEQ, lhs, rhs)
: createIComparisonExpr(CmpIPredicate::eq, lhs, rhs);
}
Value mlir::edsc::op::ne(Value lhs, Value rhs) {
auto type = lhs.getType();
return type.isa<FloatType>()
? createFComparisonExpr(CmpFPredicate::ONE, lhs, rhs)
: createIComparisonExpr(CmpIPredicate::ne, lhs, rhs);
}
Value mlir::edsc::op::slt(Value lhs, Value rhs) {
auto type = lhs.getType();
return type.isa<FloatType>()
? createFComparisonExpr(CmpFPredicate::OLT, lhs, rhs)
: createIComparisonExpr(CmpIPredicate::slt, lhs, rhs);
}
Value mlir::edsc::op::sle(Value lhs, Value rhs) {
auto type = lhs.getType();
return type.isa<FloatType>()
? createFComparisonExpr(CmpFPredicate::OLE, lhs, rhs)
: createIComparisonExpr(CmpIPredicate::sle, lhs, rhs);
}
Value mlir::edsc::op::sgt(Value lhs, Value rhs) {
auto type = lhs.getType();
return type.isa<FloatType>()
? createFComparisonExpr(CmpFPredicate::OGT, lhs, rhs)
: createIComparisonExpr(CmpIPredicate::sgt, lhs, rhs);
}
Value mlir::edsc::op::sge(Value lhs, Value rhs) {
auto type = lhs.getType();
return type.isa<FloatType>()
? createFComparisonExpr(CmpFPredicate::OGE, lhs, rhs)
: createIComparisonExpr(CmpIPredicate::sge, lhs, rhs);
}
Value mlir::edsc::op::ult(Value lhs, Value rhs) {
auto type = lhs.getType();
return type.isa<FloatType>()
? createFComparisonExpr(CmpFPredicate::OLT, lhs, rhs)
: createIComparisonExpr(CmpIPredicate::ult, lhs, rhs);
}
Value mlir::edsc::op::ule(Value lhs, Value rhs) {
auto type = lhs.getType();
return type.isa<FloatType>()
? createFComparisonExpr(CmpFPredicate::OLE, lhs, rhs)
: createIComparisonExpr(CmpIPredicate::ule, lhs, rhs);
}
Value mlir::edsc::op::ugt(Value lhs, Value rhs) {
auto type = lhs.getType();
return type.isa<FloatType>()
? createFComparisonExpr(CmpFPredicate::OGT, lhs, rhs)
: createIComparisonExpr(CmpIPredicate::ugt, lhs, rhs);
}
Value mlir::edsc::op::uge(Value lhs, Value rhs) {
auto type = lhs.getType();
return type.isa<FloatType>()
? createFComparisonExpr(CmpFPredicate::OGE, lhs, rhs)
: createIComparisonExpr(CmpIPredicate::uge, lhs, rhs);
}

17
mlir/lib/Dialect/Affine/EDSC/CMakeLists.txt
View File

@ -1,17 +0,0 @@
add_mlir_dialect_library(MLIRAffineEDSC
Builders.cpp
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
DEPENDS
MLIRAffineOpsIncGen
LINK_LIBS PUBLIC
MLIRAffine
MLIREDSC
MLIRIR
MLIRLoopLikeInterface
MLIRSideEffectInterfaces
MLIRStandard
)

36
mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
View File

@ -12,6 +12,7 @@
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/StandardOps/Utils/Utils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BlockAndValueMapping.h"
@ -24,41 +25,6 @@
using namespace mlir;
using namespace mlir::linalg;
namespace {
/// Helper struct to build simple arithmetic quantities with minimal type
/// inference support.
struct ArithBuilder {
ArithBuilder(OpBuilder &b, Location loc) : b(b), loc(loc) {}
Value select(Value cmp, Value lhs, Value rhs) {
return b.create<SelectOp>(loc, cmp, lhs, rhs);
}
Value slt(Value lhs, Value rhs) {
if (lhs.getType().isa<IntegerType>())
return b.create<CmpIOp>(loc, CmpIPredicate::slt, lhs, rhs);
return b.create<CmpFOp>(loc, CmpFPredicate::OLT, lhs, rhs);
}
Value sgt(Value lhs, Value rhs) {
if (lhs.getType().isa<IntegerType>())
return b.create<CmpIOp>(loc, CmpIPredicate::sgt, lhs, rhs);
return b.create<CmpFOp>(loc, CmpFPredicate::OGT, lhs, rhs);
}
Value add(Value lhs, Value rhs) {
if (lhs.getType().isa<IntegerType>())
return b.create<AddIOp>(loc, lhs, rhs);
return b.create<AddFOp>(loc, lhs, rhs);
}
Value mul(Value lhs, Value rhs) {
if (lhs.getType().isa<IntegerType>())
return b.create<MulIOp>(loc, lhs, rhs);
return b.create<MulFOp>(loc, lhs, rhs);
}
OpBuilder &b;
Location loc;
};
} // namespace
static SmallVector<Value> makeCanonicalAffineApplies(OpBuilder &b, Location loc,
AffineMap map,
ArrayRef<Value> vals) {

1
mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
View File

@ -11,7 +11,6 @@
//===----------------------------------------------------------------------===//
#include "PassDetail.h"
#include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
#include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"

1
mlir/lib/Dialect/Linalg/Utils/CMakeLists.txt
View File

@ -6,7 +6,6 @@ add_mlir_dialect_library(MLIRLinalgUtils
LINK_LIBS PUBLIC
MLIRAffine
MLIRAffineEDSC
MLIRIR
MLIRLinalg
MLIRSCF

26
mlir/lib/Dialect/Linalg/Utils/Utils.cpp
View File

@ -12,13 +12,13 @@
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
#include "mlir/Dialect/SCF/SCF.h"
#include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/Dialect/StandardOps/Utils/Utils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineExprVisitor.h"
#include "mlir/IR/AffineMap.h"
@ -311,10 +311,11 @@ void GenerateLoopNest<TiledLoopOp>::doit(
void updateBoundsForCyclicDistribution(OpBuilder &b, Location loc, Value procId,
Value nprocs, Value &lb, Value &ub,
Value &step) {
using edsc::op::operator+;
using edsc::op::operator*;
lb = lb + (procId * step);
step = nprocs * step;
AffineExpr d0, d1;
bindDims(b.getContext(), d0, d1);
AffineExpr s0 = getAffineSymbolExpr(0, b.getContext());
lb = makeComposedAffineApply(b, loc, d0 + d1 * s0, {lb, procId, step});
step = makeComposedAffineApply(b, loc, d0 * s0, {nprocs, step});
}
/// Generates a loop nest consisting of scf.parallel and scf.for, depending
@ -413,11 +414,10 @@ static void generateParallelLoopNest(
}
case DistributionMethod::CyclicNumProcsGeNumIters: {
// Check (for the processed loops) that the iteration is in-bounds.
using edsc::op::slt;
using edsc::op::operator&&;
Value cond = slt(lbs[0], ubs[0]);
ArithBuilder ab(b, loc);
Value cond = ab.slt(lbs[0], ubs[0]);
for (unsigned i = 1; i < numProcessed; ++i)
cond = cond && slt(lbs[i], ubs[i]);
cond = ab._and(cond, ab.slt(lbs[i], ubs[i]));
ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
b.create<scf::IfOp>(loc, cond, [&](OpBuilder &b, Location loc) {
generateParallelLoopNest(
@ -517,8 +517,6 @@ SmallVector<Value, 4> makeTiledShapes(OpBuilder &b, Location loc,
[](Value v) { return !isZero(v); })) &&
"expected as many ivs as non-zero sizes");
using namespace edsc::op;
// Construct (potentially temporary) mins and maxes on which to apply maps
// that define tile subshapes.
SmallVector<Value, 8> lbs, subShapeSizes;
@ -529,7 +527,8 @@ SmallVector<Value, 4> makeTiledShapes(OpBuilder &b, Location loc,
: (Value)b.create<ConstantIndexOp>(loc, 0));
// Before composing, we need to make range a closed interval.
Value size = isTiled ? tileSizes[idx] : sizeBounds[idx];
subShapeSizes.push_back(size - b.create<ConstantIndexOp>(loc, 1));
AffineExpr d0 = getAffineDimExpr(0, b.getContext());
subShapeSizes.push_back(makeComposedAffineApply(b, loc, d0 - 1, size));
LLVM_DEBUG(llvm::dbgs() << "lb: " << lbs.back() << "\n");
LLVM_DEBUG(llvm::dbgs() << "size: " << subShapeSizes.back() << "\n");
}
@ -577,7 +576,8 @@ SmallVector<Value, 4> makeTiledShapes(OpBuilder &b, Location loc,
offsets.push_back(offset);
auto closedIntSize = applyMapToValues(b, loc, m, subShapeSizes).front();
// Resulting size needs to be made half open interval again.
auto size = closedIntSize + b.create<ConstantIndexOp>(loc, 1);
AffineExpr s0 = getAffineSymbolExpr(0, b.getContext());
Value size = makeComposedAffineApply(b, loc, s0 + 1, closedIntSize);
LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: raw size: " << size << "\n");
// The size of the subview / subtensor should be trimmed to avoid

27
mlir/lib/Dialect/StandardOps/Utils/Utils.cpp
View File

@ -48,3 +48,30 @@ void mlir::getPositionsOfShapeOne(
}
}
}
Value ArithBuilder::_and(Value lhs, Value rhs) {
return b.create<AndOp>(loc, lhs, rhs);
}
Value ArithBuilder::add(Value lhs, Value rhs) {
if (lhs.getType().isa<IntegerType>())
return b.create<AddIOp>(loc, lhs, rhs);
return b.create<AddFOp>(loc, lhs, rhs);
}
Value ArithBuilder::mul(Value lhs, Value rhs) {
if (lhs.getType().isa<IntegerType>())
return b.create<MulIOp>(loc, lhs, rhs);
return b.create<MulFOp>(loc, lhs, rhs);
}
Value ArithBuilder::sgt(Value lhs, Value rhs) {
if (lhs.getType().isa<IndexType, IntegerType>())
return b.create<CmpIOp>(loc, CmpIPredicate::sgt, lhs, rhs);
return b.create<CmpFOp>(loc, CmpFPredicate::OGT, lhs, rhs);
}
Value ArithBuilder::slt(Value lhs, Value rhs) {
if (lhs.getType().isa<IndexType, IntegerType>())
return b.create<CmpIOp>(loc, CmpIPredicate::slt, lhs, rhs);
return b.create<CmpFOp>(loc, CmpFPredicate::OLT, lhs, rhs);
}
Value ArithBuilder::select(Value cmp, Value lhs, Value rhs) {
return b.create<SelectOp>(loc, cmp, lhs, rhs);
}

2
mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir
View File

@ -203,8 +203,6 @@ func @conv_tensors_dynamic(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?x?x?x
return %for0 : tensor<?x?x?x?xf32>
}
// -----
// CHECK: #[[BOUND8_MAP:.+]] = affine_map<(d0)[s0] -> (8, -d0 + s0)>
// CHECK: #[[BOUND8_MAP_2:.+]] = affine_map<(d0)[s0, s1] -> (-d0 + s0, 8, -d0 + s1)>
// CHECK: #[[BOUND8_MAP_3:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 8)>