merge consecutive adjust_scale ops

PiperOrigin-RevId: 952797976

Author: j2kun

lib/Dialect/Mgmt/IR/MgmtCanonicalization.td

@@ -52,4 +52,10 @@ def MergeModReduce : Pat<
   [(HasOneUse:$modReduceOp1)]
 >;
 
+def MergeAdjustScale : Pat<
+  (Mgmt_AdjustScaleOp:$adjustScaleOp2 (Mgmt_AdjustScaleOp:$adjustScaleOp1 $input, $id1), $id2),
+  (Mgmt_AdjustScaleOp $input, $id2),
+  [(HasOneUse:$adjustScaleOp1)]
+>;
+
 #endif  // LIB_DIALECT_MGMT_IR_MGMTCANONICALIZATION_TD_

lib/Dialect/Mgmt/IR/MgmtOps.cpp

@@ -31,7 +31,7 @@ void LevelReduceOp::getCanonicalizationPatterns(RewritePatternSet& results,
 
 void AdjustScaleOp::getCanonicalizationPatterns(RewritePatternSet& results,
                                                 MLIRContext* context) {
-  results.add<ModReduceAfterAdjustScale>(context);
+  results.add<ModReduceAfterAdjustScale, MergeAdjustScale>(context);
 }
 
 void LevelReduceMinOp::getCanonicalizationPatterns(RewritePatternSet& results,