std::unique_ptr<BuilderType> builder;
};
-enum class IterType { Parallel, Reduction };
-
-inline StringRef toString(IterType t) {
- switch (t) {
- case IterType::Parallel:
- return getParallelIteratorTypeName();
- case IterType::Reduction:
- return getReductionIteratorTypeName();
- }
- llvm_unreachable("Unsupported IterType");
-}
-
-/// A StructuredIndexed represents an indexable quantity that is either:
-/// 1. a captured value, which is suitable for buffer and tensor operands, or;
-/// 2. a captured type, which is suitable for tensor return values.
-///
-/// A StructuredIndexed itself is indexed and passed to `makeGenericLinalgOp`.
-/// It enable an idiomatic syntax for index expressions such as:
-///
-/// ```
-/// StructuredIndexed A(buffer_or_tensor_value), B(buffer_or_tensor_value),
-/// C(buffer_value_or_tensor_type);
-/// makeGenericLinalgOp({A({m, n}), B({k, n})}, {C({m, n})}, ... );
-/// ```
-struct StructuredIndexed : public ValueHandle {
- StructuredIndexed(Type type) : ValueHandle(type) {}
- StructuredIndexed(Value value) : ValueHandle(value) {}
- StructuredIndexed(ValueHandle valueHandle) : ValueHandle(valueHandle) {}
- StructuredIndexed operator()(ArrayRef<AffineExpr> indexings) {
- return StructuredIndexed(*this, indexings);
- }
-
- ArrayRef<AffineExpr> getExprs() { return exprs; }
-
-private:
- StructuredIndexed(Type t, ArrayRef<AffineExpr> indexings)
- : ValueHandle(t), exprs(indexings.begin(), indexings.end()) {
- assert(t.isa<RankedTensorType>() && "RankedTensor expected");
- }
- StructuredIndexed(Value v, ArrayRef<AffineExpr> indexings)
- : ValueHandle(v), exprs(indexings.begin(), indexings.end()) {
- assert((v.getType().isa<MemRefType>() ||
- v.getType().isa<RankedTensorType>()) &&
- "MemRef or RankedTensor expected");
- }
- StructuredIndexed(ValueHandle vh, ArrayRef<AffineExpr> indexings)
- : ValueHandle(vh), exprs(indexings.begin(), indexings.end()) {}
-
- SmallVector<AffineExpr, 4> exprs;
-};
-
inline void defaultRegionBuilder(ArrayRef<BlockArgument> args) {}
/// Build a `linalg.generic` op with the specified `inputs`, `outputs` and
/// restriction output tensor results would need to be reordered, which would
/// result in surprising behavior when combined with region definition.
Operation *makeGenericLinalgOp(
- ArrayRef<IterType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
+ ArrayRef<IteratorType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
ArrayRef<StructuredIndexed> outputs,
function_ref<void(ArrayRef<BlockArgument>)> regionBuilder =
defaultRegionBuilder,
return res;
}
+/// Typed representation for loop type strings.
+enum class IteratorType { Parallel, Reduction };
+
+inline StringRef toString(IteratorType t) {
+ switch (t) {
+ case IteratorType::Parallel:
+ return getParallelIteratorTypeName();
+ case IteratorType::Reduction:
+ return getReductionIteratorTypeName();
+ }
+ llvm_unreachable("Unsupported IteratorType");
+}
+
} // end namespace mlir
#endif // MLIR_UTILS_STRUCTUREDOPSUTILS_H
#include "mlir/Dialect/AffineOps/AffineOps.h"
#include "mlir/Dialect/LoopOps/LoopOps.h"
#include "mlir/Dialect/StandardOps/Ops.h"
+#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/Builders.h"
#include "mlir/Transforms/FoldUtils.h"
mlir::Block *block;
};
+/// A StructuredIndexed represents an indexable quantity that is either:
+/// 1. a captured value, which is suitable for buffer and tensor operands, or;
+/// 2. a captured type, which is suitable for tensor return values.
+///
+/// A StructuredIndexed itself is indexed and passed to `makeGenericLinalgOp`.
+/// It enable an idiomatic syntax for index expressions such as:
+///
+/// ```
+/// StructuredIndexed A(buffer_or_tensor_value), B(buffer_or_tensor_value),
+/// C(buffer_value_or_tensor_type);
+/// makeGenericLinalgOp({A({m, n}), B({k, n})}, {C({m, n})}, ... );
+/// ```
+struct StructuredIndexed : public ValueHandle {
+ StructuredIndexed(Type type) : ValueHandle(type) {}
+ StructuredIndexed(Value value) : ValueHandle(value) {}
+ StructuredIndexed(ValueHandle valueHandle) : ValueHandle(valueHandle) {}
+ StructuredIndexed operator()(ArrayRef<AffineExpr> indexings) {
+ return this->hasValue() ? StructuredIndexed(this->getValue(), indexings)
+ : StructuredIndexed(this->getType(), indexings);
+ }
+
+ StructuredIndexed(Type t, ArrayRef<AffineExpr> indexings)
+ : ValueHandle(t), exprs(indexings.begin(), indexings.end()) {
+ assert(t.isa<RankedTensorType>() && "RankedTensor expected");
+ }
+ StructuredIndexed(Value v, ArrayRef<AffineExpr> indexings)
+ : ValueHandle(v), exprs(indexings.begin(), indexings.end()) {
+ assert((v.getType().isa<MemRefType>() ||
+ v.getType().isa<RankedTensorType>()) &&
+ "MemRef or RankedTensor expected");
+ }
+ StructuredIndexed(ValueHandle vh, ArrayRef<AffineExpr> indexings)
+ : ValueHandle(vh), exprs(indexings.begin(), indexings.end()) {}
+
+ ArrayRef<AffineExpr> getExprs() { return exprs; }
+
+private:
+ SmallVector<AffineExpr, 4> exprs;
+};
+
template <typename Op, typename... Args>
OperationHandle OperationHandle::create(Args... args) {
return OperationHandle(ScopedContext::getBuilder()
#include "mlir/Dialect/Linalg/EDSC/Builders.h"
#include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/EDSC/Builders.h"
#include "mlir/EDSC/Intrinsics.h"
#include "mlir/IR/AffineExpr.h"
}
Operation *mlir::edsc::makeGenericLinalgOp(
- ArrayRef<IterType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
+ ArrayRef<IteratorType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
ArrayRef<StructuredIndexed> outputs,
function_ref<void(ArrayRef<BlockArgument>)> regionBuilder,
ArrayRef<Value> otherValues, ArrayRef<Attribute> otherAttributes) {
Operation *mlir::edsc::ops::linalg_pointwise(UnaryPointwiseOpBuilder unaryOp,
StructuredIndexed I,
StructuredIndexed O) {
- SmallVector<edsc::IterType, 4> iterTypes(O.getExprs().size(),
- edsc::IterType::Parallel);
+ SmallVector<IteratorType, 4> iterTypes(O.getExprs().size(),
+ IteratorType::Parallel);
if (O.getType().isa<RankedTensorType>()) {
auto fun = [&unaryOp](ArrayRef<BlockArgument> args) {
assert(args.size() == 1 && "expected 1 block arguments");
StructuredIndexed I1,
StructuredIndexed I2,
StructuredIndexed O) {
- SmallVector<edsc::IterType, 4> iterTypes(O.getExprs().size(),
- edsc::IterType::Parallel);
+ SmallVector<IteratorType, 4> iterTypes(O.getExprs().size(),
+ IteratorType::Parallel);
if (O.getType().isa<RankedTensorType>()) {
auto fun = [&binaryOp](ArrayRef<BlockArgument> args) {
assert(args.size() == 2 && "expected 2 block arguments");
bindDims(ScopedContext::getContext(), m, n, k);
StructuredIndexed A(vA), B(vB), C(vC);
return makeGenericLinalgOp(
- {IterType::Parallel, IterType::Parallel, IterType::Reduction},
+ {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
{A({m, k}), B({k, n})},
{C({m, n})},
macRegionBuilder);
bindDims(ScopedContext::getContext(), m, n, k);
StructuredIndexed A(vA), B(vB), C(tC);
return makeGenericLinalgOp(
- {IterType::Parallel, IterType::Parallel, IterType::Reduction},
+ {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
{A({m, k}), B({k, n})},
{C({m, n})},
mulRegionBuilder);
bindDims(ScopedContext::getContext(), m, n, k);
StructuredIndexed A(vA), B(vB), C(vC), D(tD);
return makeGenericLinalgOp(
- {IterType::Parallel, IterType::Parallel, IterType::Reduction},
+ {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
{A({m, k}), B({k, n}), C({m, n})},
{D({m, n})},
macRegionBuilder);
assert((strides.empty() || strides.size() == 2) && "only 2-D conv atm");
// Some short names.
- auto par = IterType::Parallel;
- auto red = IterType::Reduction;
+ auto par = IteratorType::Parallel;
+ auto red = IteratorType::Reduction;
auto s = strides;
auto d = dilations;
assert((strides.empty() || strides.size() == 2) && "only 2-D conv atm");
// Some short names.
- auto par = IterType::Parallel;
- auto red = IterType::Reduction;
+ auto par = IteratorType::Parallel;
+ auto red = IteratorType::Reduction;
auto s = strides;
auto d = dilations;