[mlir][Shape] Lower cstr_eq to shape_eq + assert

mlir/lib/Conversion/ShapeToStandard/ConvertShapeConstraints.cpp

LogicalResult matchAndRewrite(shape::CstrRequireOp op,
     return success();
   }
 };
 } // namespace
 
 void mlir::populateConvertShapeConstraintsConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   patterns.insert<CstrBroadcastableToRequire>(ctx);
+  patterns.insert<CstrEqToRequire>(ctx);
   patterns.insert<ConvertCstrRequireOp>(ctx);
 }
 
 namespace {
 // This pass eliminates shape constraints from the program, converting them to
 // eager (side-effecting) error handling code. After eager error handling code
 // is emitted, witnesses are satisfied, so they are replace with
 // `shape.const_witness true`.

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.td

def BroadcastableStringAttr : NativeCodeCall<[{
   $_builder.getStringAttr("required broadcastable shapes")
 }]>;
 
 def CstrBroadcastableToRequire : Pat<(Shape_CstrBroadcastableOp $shapes),
             (Shape_CstrRequireOp
               (Shape_IsBroadcastableOp $shapes),
               (BroadcastableStringAttr))>;
 
+def EqStringAttr : NativeCodeCall<[{
+  $_builder.getStringAttr("required equal shapes")
+}]>;
+
+def CstrEqToRequire : Pat<(Shape_CstrEqOp $shapes),
+  (Shape_CstrRequireOp (Shape_ShapeEqOp $shapes), (EqStringAttr))>;
+
 #endif // MLIR_CONVERSION_SHAPETOSTANDARD_TD

mlir/test/Conversion/ShapeToStandard/convert-shape-constraints.mlir

 // RUN: mlir-opt -convert-shape-constraints <%s | FileCheck %s
 
 // There's not very much useful to check here other than pasting the output.
 // CHECK-LABEL:   func @cstr_broadcastable(
 // CHECK-SAME:                             %[[LHS:.*]]: tensor<?xindex>,
 // CHECK-SAME:                             %[[RHS:.*]]: tensor<?xindex>) -> !shape.witness {
 // CHECK:           %[[RET:.*]] = shape.const_witness true
 // CHECK:           %[[BROADCAST_IS_VALID:.*]] = shape.is_broadcastable %[[LHS]], %[[RHS]]
 // CHECK:           assert %[[BROADCAST_IS_VALID]], "required broadcastable shapes"
 // CHECK:           return %[[RET]] : !shape.witness
 // CHECK:         }
 func @cstr_broadcastable(%arg0: tensor<?xindex>, %arg1: tensor<?xindex>) -> !shape.witness {
   %witness = shape.cstr_broadcastable %arg0, %arg1 : tensor<?xindex>, tensor<?xindex>
   return %witness : !shape.witness
 }
 
+// CHECK-LABEL:   func @cstr_eq(
+// CHECK-SAME:                             %[[LHS:.*]]: tensor<?xindex>,
+// CHECK-SAME:                             %[[RHS:.*]]: tensor<?xindex>) -> !shape.witness {
+// CHECK:           %[[RET:.*]] = shape.const_witness true
+// CHECK:           %[[EQUAL_IS_VALID:.*]] = shape.shape_eq %[[LHS]], %[[RHS]]
+// CHECK:           assert %[[EQUAL_IS_VALID]], "required equal shapes"
+// CHECK:           return %[[RET]] : !shape.witness
+// CHECK:         }
+func @cstr_eq(%arg0: tensor<?xindex>, %arg1: tensor<?xindex>) -> !shape.witness {
+  %witness = shape.cstr_eq %arg0, %arg1 : tensor<?xindex>, tensor<?xindex>
+  return %witness : !shape.witness
+}
+
 // CHECK-LABEL: func @cstr_require
 func @cstr_require(%arg0: i1) -> !shape.witness {
   // CHECK: %[[RET:.*]] = shape.const_witness true
   // CHECK: assert %arg0, "msg"
   // CHECK: return %[[RET]]
   %witness = shape.cstr_require %arg0, "msg"
   return %witness : !shape.witness
 }