[TorchToArith] Add a lowering for AtenEqFloat (#4022)

Addresses an issue introduced by <#3945> in an external test suite.
llvm · Feb 12, 2025 · aa74936 · aa74936
1 parent c9694c6
commit aa74936
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/lib/Conversion/TorchToArith/TorchToArith.cpp b/lib/Conversion/TorchToArith/TorchToArith.cpp
@@ -454,8 +454,11 @@ class ConvertTorchToArith
     patterns.add<
         ConvertAtenIntComparisonOp<AtenLeIntOp, arith::CmpIPredicate::sle>>(
         typeConverter, context);
-    target.addIllegalOp<AtenGeFloatOp, AtenGtFloatOp, AtenGeFloatIntOp,
-                        AtenNeFloatIntOp, AtenGtFloatIntOp>();
+    target.addIllegalOp<AtenEqFloatOp, AtenGeFloatOp, AtenGtFloatOp,
+                        AtenGeFloatIntOp, AtenNeFloatIntOp, AtenGtFloatIntOp>();
+    patterns.add<
+        ConvertAtenFloatComparisonOp<AtenEqFloatOp, arith::CmpFPredicate::UEQ>>(
+        typeConverter, context);
     patterns.add<
         ConvertAtenFloatComparisonOp<AtenGeFloatOp, arith::CmpFPredicate::UGE>>(
         typeConverter, context);