[MLIR][Shape] Lower `shape.shape_eq` to `scf`

[frgossen]

Summary

Lower shape.shape_eq to the scf (and std) dialect.

Depends On D82529

[pifon2a]

mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp

42

if (!operandTy.isa<ShapeType>())
      return failure();

74

nit: replace auto with Value or Type where possible.

mlir/test/Conversion/ShapeToSCF/shape-to-scf.mlir

28

i bet you don't need most of these -DAGs.

[frgossen]

Thanks

mlir/test/Conversion/ShapeToSCF/shape-to-scf.mlir

28

I like them because they make it clear where order matters and where it doesn't.
Should I change this?

[frgossen]

Address comments

[pifon2a]

mlir/test/Conversion/ShapeToSCF/shape-to-scf.mlir

28

yes, please. It shows what deterministic and what not.

[frgossen]

Remove CHECK-DAGs

[frgossen]

Remove CHECK-DAGs

[herhut]

mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp

45

nit: size -> shape?

63

A shame we cannot model early abort in SCF. This is ok for now, though. If we get to the point that shape comparisons are performance critical, we got quite far :)

[rriddle]

mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp

45

Please add a bigger comment detailing how this lowering works. Having a before/after code sample is also generally acceptable.

[frgossen]

Add documentation

[frgossen]

Update to extent tensor variant

mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp

 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::shape;
 using namespace mlir::scf;
 
 /// Conversion patterns.
 namespace {
+/// Converts `shape.eq` to `scf.for` for occurrences that test shape equality.
+struct EqOpConverter : public OpConversionPattern<EqOp> {
+  using OpConversionPattern<EqOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(EqOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+} // namespace
+
+LogicalResult
+EqOpConverter::matchAndRewrite(EqOp op, ArrayRef<Value> operands,
+                               ConversionPatternRewriter &rewriter) const {
+  Type operandTy = op.lhs().getType();
+
+  // Size equality is lowered to `std` and the pattern can be found in the
+  // corresponding pass.
+  if (!operandTy.isa<ShapeType>())
+    return failure();
+
+  // Lower size equality.
+  EqOp::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);
+        Attribute trueAttr = b.getI1IntegerAttr(true);
+        Value init = b.create<ConstantOp>(loc, i1Ty, trueAttr);
+        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) {
+        Attribute falseAttr = b.getI1IntegerAttr(false);
+        Value result = b.create<ConstantOp>(loc, i1Ty, falseAttr);
+        b.create<scf::YieldOp>(loc, result);
+      });
+  return success();
+}
+
+namespace {
 /// Converts `shape.reduce` to `scf.for`.
 struct ReduceOpConverter : public OpConversionPattern<shape::ReduceOp> {
 public:
   using OpConversionPattern<shape::ReduceOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(shape::ReduceOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
return typeConverter.isSignatureLegal(op.getType()) &&
            typeConverter.isLegal(&op.getBody());
   });
   if (failed(applyPartialConversion(getFunction(), target, patterns)))
     signalPassFailure();
 }
 
 void mlir::populateShapeToSCFConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
-  patterns.insert<ReduceOpConverter>(ctx);
+  patterns.insert<EqOpConverter, ReduceOpConverter>(ctx);
 }
 
 std::unique_ptr<FunctionPass> mlir::createConvertShapeToSCFPass() {
   return std::make_unique<ConvertShapeToSCFPass>();
 }

mlir/test/Conversion/ShapeToSCF/shape-to-scf.mlir

 // CHECK:      %[[SIZE:.*]] = dim %[[SHAPE]], %[[C0]] : tensor<?xindex>
 // CHECK:      %[[C1:.*]] = constant 1 : index
 // CHECK:      %[[RESULT:.*]] = scf.for %[[I:.*]] = %[[C0]] to %[[SIZE]]
 // CHECK-SAME:     step %[[C1]] iter_args(%[[ACC:.*]] = %[[INIT]]) -> (index) {
 // CHECK:        %[[EXTENT:.*]] = extract_element %[[SHAPE]][%[[I]]]
 // CHECK:        %[[ACC_NEXT:.*]] = muli %[[ACC]], %[[EXTENT]]
 // CHECK:        scf.yield %[[ACC_NEXT]] : index
 // CHECK:      }
+
+// -----
+
+// CHECK-LABEL:  @eq_shape
+// CHECK-SAME:   (%[[A:.*]]: tensor<?xindex>, %[[B:.*]]: tensor<?xindex>) -> i1
+func @eq_shape(%a : !shape.shape, %b : !shape.shape) -> 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]]
+  // CHECK-SAME:       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.eq %a, %b : !shape.shape
+  return %result : i1
+}