[MLIR][Shape] Lower `shape.size_eq` to the standard dialect

[frgossen]

Summary

Lower shape.size_eq to the standard dialect's cmpi operation.

Depends On D82628

[pifon2a]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

197

shouldn't this check be a part of verify()?

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

200

just CHECK or CHECK-NEXT?

[pifon2a]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

197

oh, i saw another PR

[pifon2a]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

197

in general, i would prefer having here

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

and

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

in the SCF counter-part.

[frgossen]

Address comments

[frgossen]

Thanks

[frgossen]

Rebase

[frgossen]

Rebase

[frgossen]

ni-t

[pifon2a]

@frgossen this one is obsolete, right? we can probably close it.

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

class GetExtentOpConverter : public OpConversionPattern<GetExtentOp> {
   LogicalResult
   matchAndRewrite(GetExtentOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 } // namespace
 
 LogicalResult GetExtentOpConverter::matchAndRewrite(
     GetExtentOp op, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   GetExtentOp::Adaptor transformed(operands);
 
   // For now, only error-free types are supported by this lowering.
   if (op.getType().isa<SizeType>())
     return failure();
 
   // Derive shape extent directly from shape origin if possible. This
   // circumvents the necessity to materialize the shape in memory.
   if (auto shapeOfOp = op.shape().getDefiningOp<ShapeOfOp>()) {
     if (shapeOfOp.arg().getType().isa<ShapedType>()) {
       rewriter.replaceOpWithNewOp<DimOp>(op, shapeOfOp.arg(),
                                          transformed.dim());
       return success();
     }
   }
 
   rewriter.replaceOpWithNewOp<ExtractElementOp>(op, rewriter.getIndexType(),
                                                 transformed.shape(),
                                                 ValueRange{transformed.dim()});
   return success();
 }
 
 namespace {
+class SizeEqOpConverter : public OpConversionPattern<SizeEqOp> {
+public:
+  using OpConversionPattern<SizeEqOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(SizeEqOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+} // namespace
+
+LogicalResult
+SizeEqOpConverter::matchAndRewrite(SizeEqOp op, ArrayRef<Value> operands,
+                                   ConversionPatternRewriter &rewriter) const {
+
+  // For now, only error-free types are supported by this lowering.
+  if (!op.lhs().getType().isa<IndexType>() ||
+      !op.rhs().getType().isa<IndexType>())
+    return failure();
+
+  SizeEqOp::Adaptor transformed(operands);
+  rewriter.replaceOpWithNewOp<CmpIOp>(op, CmpIPredicate::eq, transformed.lhs(),
+                                      transformed.rhs());
+  return success();
+}
+
+namespace {
 class RankOpConverter : public OpConversionPattern<shape::RankOp> {
 public:
   using OpConversionPattern<shape::RankOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(shape::RankOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
patterns.insert<
       AnyOpConversion,
       BinaryOpConversion<AddOp, AddIOp>,
       ConstShapeOpConverter,
       BinaryOpConversion<MulOp, MulIOp>,
       ConstSizeOpConversion,
       GetExtentOpConverter,
       RankOpConverter,
       ShapeOfOpConversion,
+      SizeEqOpConverter,
       ToExtentTensorOpConversion>(ctx);
   // clang-format on
 }
 
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createConvertShapeToStandardPass() {
   return std::make_unique<ConvertShapeToStandardPass>();
 }

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

 // CHECK-SAME: (%[[ARG:.*]]: tensor<?xindex>
 func @to_extent_tensor(%arg: tensor<?xindex>) -> tensor<3xindex> {
   // CHECK-NOT: to_extent_tensor
   // CHECK: %[[RES:.*]] = tensor_cast %[[ARG]] : tensor<?xindex> to tensor<3xindex
   %casted = shape.to_extent_tensor %arg : tensor<?xindex> -> tensor<3xindex>
   // CHECK: return %[[RES]]
   return %casted : tensor<3xindex>
 }
+
+// -----
+
+// Lower `size_eq` to `cmpi`.
+// CHECK-LABEL: @size_eq
+// CHECK-SAME: (%[[A:.*]]: index, %[[B:.*]]: index) -> i1
+func @size_eq(%a : index, %b : index) -> i1 {
+  // CHECK: %[[RESULT:.*]] = cmpi "eq", %[[A]], %[[B]]
+  // CHECK: return %[[RESULT]] : i1
+  %result = shape.size_eq %a, %b : index, index
+  return %result : i1
+}
+