[mlir] Fix bug in use of rewiter for AtomicRMWOp.

[tpopp]

Summary

The legalization did not forward the listener which prevents dynamic
legalization and prevents rollbacks. This handled that and then changed
the associated pass to support all other std ops to support partial
conversion.

Previously, this lowering was failing, but due to the
initial bug, the op's modifications were not reverted, and thus the
pattern matching succeeded.

mlir/lib/Dialect/StandardOps/Transforms/ExpandAtomic.cpp

case AtomicRMWKind::minf:
       break;
     default:
       return failure();
     }
 
     auto loc = op.getLoc();
     auto genericOp =
         rewriter.create<GenericAtomicRMWOp>(loc, op.memref(), op.indices());
-    OpBuilder bodyBuilder = OpBuilder::atBlockEnd(genericOp.getBody());
+    OpBuilder bodyBuilder =
+        OpBuilder::atBlockEnd(genericOp.getBody(), rewriter.getListener());
 
     Value lhs = genericOp.getCurrentValue();
     Value rhs = op.value();
     Value cmp = bodyBuilder.create<CmpFOp>(loc, predicate, lhs, rhs);
     Value select = bodyBuilder.create<SelectOp>(loc, cmp, lhs, rhs);
     bodyBuilder.create<AtomicYieldOp>(loc, select);
 
     rewriter.replaceOp(op, genericOp.getResult());
     return success();
   }
 };
 
 struct ExpandAtomic : public ExpandAtomicBase<ExpandAtomic> {
   void runOnFunction() override {
     OwningRewritePatternList patterns;
     patterns.insert<AtomicRMWOpConverter>(&getContext());
 
     ConversionTarget target(getContext());
-    target.addLegalOp<GenericAtomicRMWOp>();
+    target.addLegalDialect<StandardOpsDialect>();
     target.addDynamicallyLegalOp<AtomicRMWOp>([](AtomicRMWOp op) {
       return op.kind() != AtomicRMWKind::maxf &&
              op.kind() != AtomicRMWKind::minf;
     });
     if (failed(mlir::applyPartialConversion(getFunction(), target,
                                             std::move(patterns))))
       signalPassFailure();
   }
 };
 
 } // namespace
 
 std::unique_ptr<Pass> mlir::createExpandAtomicPass() {
   return std::make_unique<ExpandAtomic>();
 }