broadcast iota (#1178)

* broadcast iota

* select comp iota to dus

* fix

* Update Compiler.jl

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* fix

* add broadcast compare

* add not compare

* add broadcast iota

* Update Reactant_jll to v0.0.142

* Update src/Compiler.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* more passes

* fix

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Mosè Giordano <[email protected]>

Author: 3 people

src/CompileOptions.jl

@@ -5,7 +5,7 @@
     rotate_to_pad_comm::Int = 1
     wrap_comm::Int = 0
     extend_comm::Int = 0
-    dus_to_pad_manual_comp_comm::Int = 0
+    dus_to_pad_manual_comp_comm::Int = 0 # 2
     dus_to_pad_comm::Int = 1
     concat_two_operands_comm::Int = 0
     concat_to_pad_comm::Int = 1

src/Compiler.jl

@@ -683,6 +683,15 @@ function optimization_passes(;
         "concat_multipad",
         "concat_concat_to_dus",
         "speculate_if_pad_to_select",
+        "broadcast_iota_simplify",
+        "select_comp_iota_to_dus",
+        "compare_cleanup",
+        "broadcast_compare",
+        "not_compare",
+        "broadcast_iota",
+	"cse_iota",
+	"compare_iota_const_simplify",
+	"reshuffle_ands_compares",
         # TODO we want to enable but may cause an infinite compile time
         # "concat_to_onedim_dusslice",
     ]
@@ -1047,6 +1056,9 @@ function get_optimize_comms_passes(options::Reactant.OptimizeCommunicationOption
         "enzyme-hlo-generate-td{patterns=lower_rotate;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}",
+        "transform-interpreter",
+        "enzyme-hlo-remove-transform",
         options_str,
         "enzyme-hlo-generate-td{patterns=lower_rotate;lower_wrap;lower_extend}",
         "transform-interpreter",