[Mlir-commits] [mlir] aca7b8d - [MLIR][Shape] Lower `shape.shape_eq` to `scf`
Frederik Gossen
llvmlistbot at llvm.org
Thu Jul 16 07:45:18 PDT 2020
Author: Frederik Gossen
Date: 2020-07-16T14:44:29Z
New Revision: aca7b8dd63c8be4fc9c301a6079adc04add23c22
URL: https://github.com/llvm/llvm-project/commit/aca7b8dd63c8be4fc9c301a6079adc04add23c22
DIFF: https://github.com/llvm/llvm-project/commit/aca7b8dd63c8be4fc9c301a6079adc04add23c22.diff
LOG: [MLIR][Shape] Lower `shape.shape_eq` to `scf`
Lower `shape.shape_eq` to the `scf` (and `std`) dialect. For now, this lowering
is limited to extent tensor operands.
Differential Revision: https://reviews.llvm.org/D82530
Added:
Modified:
mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp
mlir/test/Conversion/ShapeToSCF/shape-to-scf.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp b/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp
index f050c4eb2278..7986aaaa6816 100644
--- a/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp
+++ b/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp
@@ -19,6 +19,92 @@ using namespace mlir;
using namespace mlir::shape;
using namespace mlir::scf;
+namespace {
+/// Converts `shape.shape_eq` to an `scf.for` loop. For now, the lowering is
+/// only defined on `tensor<?xindex>` operands. The test for equality first
+/// compares their size and, if equal, checks every extent for equality.
+///
+/// Example:
+///
+/// %result = shape.shape_eq %a, %b : tensor<?xindex>, tensor<?xindex>
+///
+/// becomes
+///
+/// %c0 = constant 0 : index
+/// %0 = dim %arg0, %c0 : tensor<?xindex>
+/// %1 = dim %arg1, %c0 : tensor<?xindex>
+/// %2 = cmpi "eq", %0, %1 : index
+/// %result = scf.if %2 -> (i1) {
+/// %c1 = constant 1 : index
+/// %true = constant true
+/// %4 = scf.for %arg2 = %c0 to %0 step %c1 iter_args(%arg3 = %true) -> (i1) {
+/// %5 = extract_element %arg0[%arg2] : tensor<?xindex>
+/// %6 = extract_element %arg1[%arg2] : tensor<?xindex>
+/// %7 = cmpi "eq", %5, %6 : index
+/// %8 = and %arg3, %7 : i1
+/// scf.yield %8 : i1
+/// }
+/// scf.yield %4 : i1
+/// } else {
+/// %false = constant false
+/// scf.yield %false : i1
+/// }
+///
+struct ShapeEqOpConverter : public OpConversionPattern<ShapeEqOp> {
+ using OpConversionPattern<ShapeEqOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(ShapeEqOp op, ArrayRef<Value> operands,
+ ConversionPatternRewriter &rewriter) const override;
+};
+} // namespace
+
+LogicalResult
+ShapeEqOpConverter::matchAndRewrite(ShapeEqOp op, ArrayRef<Value> operands,
+ ConversionPatternRewriter &rewriter) const {
+ // For now, this lowering is only defined on `tensor<?xindex>` operands, not
+ // on shapes.
+ if (op.lhs().getType().isa<ShapeType>() ||
+ op.rhs().getType().isa<ShapeType>()) {
+ return failure();
+ }
+
+ ShapeEqOp::Adaptor transformed(operands);
+ auto loc = op.getLoc();
+ Type indexTy = rewriter.getIndexType();
+ Value zero = rewriter.create<ConstantIndexOp>(loc, 0);
+ Value lhsRank = rewriter.create<DimOp>(loc, indexTy, transformed.lhs(), zero);
+ Value rhsRank = rewriter.create<DimOp>(loc, indexTy, transformed.rhs(), zero);
+ Value eqRank =
+ rewriter.create<CmpIOp>(loc, CmpIPredicate::eq, lhsRank, rhsRank);
+ Type i1Ty = rewriter.getI1Type();
+ rewriter.replaceOpWithNewOp<IfOp>(
+ op, i1Ty, eqRank,
+ [&](OpBuilder &b, Location loc) {
+ Value one = b.create<ConstantIndexOp>(loc, 1);
+ Value init = b.create<ConstantOp>(loc, i1Ty, b.getBoolAttr(true));
+ auto loop = b.create<scf::ForOp>(
+ loc, zero, lhsRank, one, ValueRange{init},
+ [&](OpBuilder &b, Location nestedLoc, Value iv, ValueRange args) {
+ Value conj = args[0];
+ Value lhsExtent =
+ b.create<ExtractElementOp>(loc, transformed.lhs(), iv);
+ Value rhsExtent =
+ b.create<ExtractElementOp>(loc, transformed.rhs(), iv);
+ Value eqExtent = b.create<CmpIOp>(loc, CmpIPredicate::eq,
+ lhsExtent, rhsExtent);
+ Value conjNext = b.create<AndOp>(loc, conj, eqExtent);
+ b.create<scf::YieldOp>(loc, ValueRange({conjNext}));
+ });
+ b.create<scf::YieldOp>(loc, loop.getResults());
+ },
+ [&](OpBuilder &b, Location loc) {
+ Value result = b.create<ConstantOp>(loc, i1Ty, b.getBoolAttr(false));
+ b.create<scf::YieldOp>(loc, result);
+ });
+ return success();
+}
+
namespace {
/// Converts `shape.reduce` to `scf.for`.
struct ReduceOpConverter : public OpConversionPattern<shape::ReduceOp> {
@@ -148,7 +234,12 @@ void ConvertShapeToSCFPass::runOnFunction() {
void mlir::populateShapeToSCFConversionPatterns(
OwningRewritePatternList &patterns, MLIRContext *ctx) {
- patterns.insert<ReduceOpConverter, ShapeOfOpConverter>(ctx);
+ // clang-format off
+ patterns.insert<
+ ShapeEqOpConverter,
+ ReduceOpConverter,
+ ShapeOfOpConverter>(ctx);
+ // clang-format on
}
std::unique_ptr<FunctionPass> mlir::createConvertShapeToSCFPass() {
diff --git a/mlir/test/Conversion/ShapeToSCF/shape-to-scf.mlir b/mlir/test/Conversion/ShapeToSCF/shape-to-scf.mlir
index 6ba630aa4aa6..2e5a45c4cc11 100644
--- a/mlir/test/Conversion/ShapeToSCF/shape-to-scf.mlir
+++ b/mlir/test/Conversion/ShapeToSCF/shape-to-scf.mlir
@@ -43,3 +43,31 @@ func @shape_of_unranked(%arg : tensor<*xf32>) {
return
}
+// -----
+
+// CHECK-LABEL: @shape_eq
+// CHECK-SAME: (%[[A:.*]]: tensor<?xindex>, %[[B:.*]]: tensor<?xindex>) -> i1
+func @shape_eq(%a : tensor<?xindex>, %b : tensor<?xindex>) -> i1 {
+ // CHECK: %[[C0:.*]] = constant 0 : index
+ // CHECK: %[[RANK_A:.*]] = dim %[[A]], %[[C0]] : tensor<?xindex>
+ // CHECK: %[[RANK_B:.*]] = dim %[[B]], %[[C0]] : tensor<?xindex>
+ // CHECK: %[[RANK_EQ:.*]] = cmpi "eq", %[[RANK_A]], %[[RANK_B]]
+ // CHECK: %[[SHAPE_EQ:.*]] = scf.if %[[RANK_EQ]] -> (i1) {
+ // CHECK: %[[C1:.*]] = constant 1 : index
+ // CHECK: %[[INIT:.*]] = constant true
+ // CHECK: %[[SHAPE_EQ_INNER:.*]] = scf.for %[[I:.*]] = %[[C0]] to %[[RANK_A]] step %[[C1]] iter_args(%[[CONJ:.*]] = %[[INIT]]) -> (i1) {
+ // CHECK: %[[EXTENT_A:.*]] = extract_element %[[A]][%[[I]]] : tensor<?xindex>
+ // CHECK: %[[EXTENT_B:.*]] = extract_element %[[B]][%[[I]]] : tensor<?xindex>
+ // CHECK: %[[EXTENT_EQ:.*]] = cmpi "eq", %[[EXTENT_A]], %[[EXTENT_B]]
+ // CHECK: %[[CONJ_NEXT:.*]] = and %[[CONJ]], %[[EXTENT_EQ]]
+ // CHECK: scf.yield %[[CONJ_NEXT]] : i1
+ // CHECK: }
+ // CHECK: scf.yield %[[SHAPE_EQ_INNER]] : i1
+ // CHECK: } else {
+ // CHECK: %[[SHAPE_EQ_INNER:.*]] = constant false
+ // CHECK: scf.yield %[[SHAPE_EQ_INNER]] : i1
+ // CHECK: }
+ // CHECK: return %[[SHAPE_EQ]] : i1
+ %result = shape.shape_eq %a, %b : tensor<?xindex>, tensor<?xindex>
+ return %result : i1
+}
More information about the Mlir-commits
mailing list