feat: dot_general_licm (#1550)

Author: avik-pal

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -25018,6 +25018,13 @@ void mlir::transform::addSliceLICM(RewritePatternSet &patterns,
   patterns.insert<LICM<stablehlo::SliceOp>>(single_user, &context, benefit);
 }
 
+void mlir::transform::addDotGeneralLICM(RewritePatternSet &patterns,
+                                        bool single_user, MLIRContext &context,
+                                        PatternBenefit benefit) {
+  patterns.insert<LICM<stablehlo::DotGeneralOp>>(single_user, &context,
+                                                 benefit);
+}
+
 void mlir::transform::addDUSLICM(RewritePatternSet &patterns, bool single_user,
                                  MLIRContext &context, PatternBenefit benefit) {
   patterns.insert<LICM<stablehlo::DynamicUpdateSliceOp>>(single_user, &context,

src/enzyme_ad/jax/Passes/EnzymeHLOPatterns.h

@@ -59,6 +59,8 @@ void addWhileLICM(RewritePatternSet &patterns, bool hoist_all,
                   MLIRContext &context, PatternBenefit benefit);
 void addSliceLICM(RewritePatternSet &patterns, bool single_user,
                   MLIRContext &context, PatternBenefit benefit);
+void addDotGeneralLICM(RewritePatternSet &patterns, bool single_user,
+                       MLIRContext &context, PatternBenefit benefit);
 void addDUSLICM(RewritePatternSet &patterns, bool single_user,
                 MLIRContext &context, PatternBenefit benefit);
 void addPadLICM(RewritePatternSet &patterns, bool single_user,

src/enzyme_ad/jax/TransformOps/GenerateApplyPatterns.cpp

@@ -149,11 +149,11 @@ LogicalResult parseTransform(OpBuilder &builder, Location loc,
     if (parameter != -1) {
       if (opName == "no_nan_add_sub_simplify" || opName == "while_simplify" ||
           opName == "sum_to_conv" || opName == "while_licm" ||
-          opName == "slice_licm" || opName == "dus_licm" ||
-          opName == "pad_licm" || opName == "elementwise_licm" ||
-          opName == "concatenate_licm" || opName == "broadcastindim_licm" ||
-          opName == "reshape_licm" || opName == "transpose_licm" ||
-          opName == "transpose_elementwise" ||
+          opName == "slice_licm" || opName == "dot_general_licm" ||
+          opName == "dus_licm" || opName == "pad_licm" ||
+          opName == "elementwise_licm" || opName == "concatenate_licm" ||
+          opName == "broadcastindim_licm" || opName == "reshape_licm" ||
+          opName == "transpose_licm" || opName == "transpose_elementwise" ||
           opName == "reshape_elementwise" || opName == "reshape_slice" ||
           opName == "reshape_dynamic_slice" ||
           opName == "extend_unary_elementwise" ||

src/enzyme_ad/jax/TransformOps/TransformOps.cpp

@@ -80,6 +80,11 @@ void ApplySliceLICMPatterns::populatePatterns(RewritePatternSet &patterns) {
   addSliceLICM(patterns, getParameter(), *getContext(),
                PatternBenefit(getBenefit().value_or(1)));
 }
+void ApplyDotGeneralLICMPatterns::populatePatterns(
+    RewritePatternSet &patterns) {
+  addDotGeneralLICM(patterns, getParameter(), *getContext(),
+                    PatternBenefit(getBenefit().value_or(1)));
+}
 void ApplyDUSLICMPatterns::populatePatterns(RewritePatternSet &patterns) {
   addDUSLICM(patterns, getParameter(), *getContext(),
              PatternBenefit(getBenefit().value_or(1)));

src/enzyme_ad/jax/TransformOps/TransformOps.td

@@ -1121,6 +1121,21 @@ def ApplySliceLICMPatterns : EnzymeHLOParameterizedPatternOp<
   }];
 }
 
+def ApplyDotGeneralLICMPatterns : EnzymeHLOParameterizedPatternOp<
+    "dot_general_licm"> {
+  let arguments = (ins OptionalAttr<I64Attr>:$benefit, BoolAttr:$parameter);
+  let assemblyFormat = "attr-dict";
+  // TODO: this should be made better searchable.
+  let extraClassDeclaration = [{
+    ::llvm::SmallVector<::mlir::DictionaryAttr>
+    static getPossibleAttrCombinations(::mlir::Builder &builder) {
+      return {builder.getDictionaryAttr(
+                  builder.getNamedAttr("parameter",
+                                       builder.getBoolAttr(true)))};
+    }
+  }];
+}
+
 def ApplyDUSLICMPatterns : EnzymeHLOParameterizedPatternOp<
     "dus_licm"> {
   let arguments = (ins OptionalAttr<I64Attr>:$benefit, BoolAttr:$parameter);

src/enzyme_ad/jax/primitives.py

@@ -283,6 +283,7 @@ def optimization_passes(
         "slice_dus_to_concat",
         "while_induction_reduction",
         "slice_licm(0)",
+        "dot_general_licm(0)",
         "pad_licm(0)",
         "elementwise_licm(0)",
         "concatenate_licm(0)",