diff --git a/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp b/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp index 88a8b7fb185c5..13a084407e53f 100644 --- a/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp +++ b/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp @@ -138,11 +138,12 @@ struct BranchOpLowering : public ConvertOpToLLVMPattern { TypeRange(adaptor.getOperands())); if (failed(convertedBlock)) return failure(); + DictionaryAttr attrs = op->getAttrDictionary(); Operation *newOp = rewriter.replaceOpWithNewOp( 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 { TypeRange(adaptor.getFalseDestOperands())); if (failed(convertedFalseBlock)) return failure(); + DictionaryAttr attrs = op->getAttrDictionary(); auto newOp = rewriter.replaceOpWithNewOp( - op, adaptor.getCondition(), *convertedTrueBlock, - adaptor.getTrueDestOperands(), *convertedFalseBlock, - adaptor.getFalseDestOperands()); - ArrayRef 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(); } };