Multipad

Author: wsmoses

src/CompileOptions.jl

@@ -25,6 +25,7 @@ communication.
     updatewithoutcorners_to_select::Int = 1
     multirotate_custom_call::Int = 1
     multislice_custom_call::Int = 1
+    multipad_custom_call::Int = 1
     wrap_custom_call::Int = 1
 end
 

src/Compiler.jl

@@ -1840,11 +1840,11 @@ function get_optimize_comms_passes(options::OptimizeCommunicationOptions)
         "enzyme-hlo-generate-td{patterns=concat_to_onedim_dus;concat_to_onedim_dusslice;concatreshape_to_onedim_dus}",
         "transform-interpreter",
         "enzyme-hlo-remove-transform",
-        "enzyme-hlo-generate-td{patterns=reshape_to_broadcast;recognize_multirotate;use_multirotate_neutral_result;recognize_multislice}",
+        "enzyme-hlo-generate-td{patterns=reshape_to_broadcast;recognize_multirotate;use_multirotate_neutral_result;recognize_multislice;recognize_multipad}",
         "transform-interpreter",
         "enzyme-hlo-remove-transform",
         options_str,
-        "enzyme-hlo-generate-td{patterns=lower_rotate;lower_wrap;lower_extend;lower_updatewithoutcorners;lower_multislice}",
+        "enzyme-hlo-generate-td{patterns=lower_rotate;lower_wrap;lower_extend;lower_updatewithoutcorners;lower_multislice;lower_multipad}",
         "transform-interpreter",
         "enzyme-hlo-remove-transform",
         options_str,