[mlir][cf] Do not access erased operation in cf.cond_br lowering (#148358)

Do not access the erased `cf.cond_br` operation in the lowering pattern.
That won't work anymore in a One-Shot Dialect Conversion and triggers a
use-after-free sanitizer error.

After the One-Shot Dialect Conversion refactoring, a
`ConversionPatternRewriter` will behave more like a normal
`PatternRewriter`.

---------

Co-authored-by: Markus Böck <markus.boeck02@gmail.com>

Author: matthias-springer

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

@@ -138,11 +138,12 @@ struct BranchOpLowering : public ConvertOpToLLVMPattern<cf::BranchOp> {
                           TypeRange(adaptor.getOperands()));
     if (failed(convertedBlock))
       return failure();
+    DictionaryAttr attrs = op->getAttrDictionary();
     Operation *newOp = rewriter.replaceOpWithNewOp<LLVM::BrOp>(
         op, adaptor.getOperands(), *convertedBlock);
     // TODO: We should not just forward all attributes like that. But there are
     // existing Flang tests that depend on this behavior.
-    newOp->setAttrs(op->getAttrDictionary());
+    newOp->setAttrs(attrs);
     return success();
   }
 };
@@ -166,18 +167,14 @@ struct CondBranchOpLowering : public ConvertOpToLLVMPattern<cf::CondBranchOp> {
                           TypeRange(adaptor.getFalseDestOperands()));
     if (failed(convertedFalseBlock))
       return failure();
+    DictionaryAttr attrs = op->getAttrDictionary();
     auto newOp = rewriter.replaceOpWithNewOp<LLVM::CondBrOp>(
-        op, adaptor.getCondition(), *convertedTrueBlock,
-        adaptor.getTrueDestOperands(), *convertedFalseBlock,
-        adaptor.getFalseDestOperands());
-    ArrayRef<int32_t> weights = op.getWeights();
-    if (!weights.empty()) {
-      newOp.setWeights(weights);
-      op.removeBranchWeightsAttr();
-    }
+        op, adaptor.getCondition(), adaptor.getTrueDestOperands(),
+        adaptor.getFalseDestOperands(), op.getBranchWeightsAttr(),
+        *convertedTrueBlock, *convertedFalseBlock);
     // TODO: We should not just forward all attributes like that. But there are
     // existing Flang tests that depend on this behavior.
-    newOp->setAttrs(op->getAttrDictionary());
+    newOp->setAttrs(attrs);
     return success();
   }
 };