Pass pipeline for unopt

[ftynse]

inline{default-pipeline=canonicalize max-iterations=4},propagate-constant-bounds,sroa-wrappers{dump_prellvm=true dump_postllvm=true instcombine=false instsimplify=true},canonicalize,sroa-wrappers{instcombine=false instsimplify=true},libdevice-funcs-raise,canonicalize,remove-duplicate-func-def,canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform,enzyme-batch,inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform

then

enzyme{postpasses=\"arith-raise{stablehlo=true},canonicalize,cse,canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math,canonicalize,cse,canonicalize\"},arith-raise{stablehlo=true}

then

canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math,inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform,lower-kernel{backend=cpu},canonicalize,canonicalize,llvm-to-memref-access,canonicalize,convert-llvm-to-cf,canonicalize,enzyme-lift-cf-to-scf,canonicalize,func.func(canonicalize-loops),canonicalize-scf-for,canonicalize,affine-cfg,canonicalize,func.func(canonicalize-loops),canonicalize,llvm-to-affine-access,canonicalize,delinearize-indexing,canonicalize,simplify-affine-exprs,affine-cfg,canonicalize,raise-affine-to-stablehlo,arith-raise{stablehlo=true},inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform

[wsmoses]

(base) wmoses@hydra:~/git/Enzyme-JaX$ cat fullpipe.sh 
#!/usr/bin/env sh

PIPELINE1="--pass-pipeline=builtin.module(inline{default-pipeline=canonicalize max-iterations=4},propagate-constant-bounds,sroa-wrappers{instcombine=false instsimplify=true},canonicalize,sroa-wrappers{instcombine=false instsimplify=true},libdevice-funcs-raise,canonicalize,remove-duplicate-func-def,canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform,enzyme-batch,inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform)"

PIPELINE2="--pass-pipeline=builtin.module(enzyme{postpasses=\"arith-raise{stablehlo=true},canonicalize,cse,canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math,canonicalize,cse,canonicalize\"},arith-raise{stablehlo=true})"

PIPELINE3="--pass-pipeline=builtin.module(canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math,inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform,lower-kernel{backend=cpu},canonicalize,canonicalize,llvm-to-memref-access,canonicalize,convert-llvm-to-cf,canonicalize,enzyme-lift-cf-to-scf,canonicalize,func.func(canonicalize-loops),canonicalize-scf-for,canonicalize,affine-cfg,canonicalize,func.func(canonicalize-loops),canonicalize,llvm-to-affine-access,canonicalize,delinearize-indexing,canonicalize,simplify-affine-exprs,affine-cfg,canonicalize,raise-affine-to-stablehlo,arith-raise{stablehlo=true},inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform)"


# PIPELINE3="--pass-pipeline=builtin.module(lower-kernel{backend=cpu},canonicalize,llvm-to-memref-access,canonicalize,convert-llvm-to-cf,canonicalize,enzyme-lift-cf-to-scf,canonicalize,func.func(canonicalize-loops),canonicalize-scf-for,canonicalize,affine-cfg,canonicalize)" # ,func.func(canonicalize-loops),canonicalize,llvm-to-affine-access)"

./bazel-bin/enzymexlamlir-opt unopt.mlir "$PIPELINE1" -o mid1.mlir
./bazel-bin/enzymexlamlir-opt mid1.mlir "$PIPELINE2" -o mid2.mlir
./bazel-bin/enzymexlamlir-opt mid2.mlir "$PIPELINE3" -o mid3.mlir

[ivanradanov]

To stop just before raise-affine-to-stablehlo

#!/usr/bin/env sh

set -e

PIPELINE1="--pass-pipeline=builtin.module(inline{default-pipeline=canonicalize max-iterations=4},propagate-constant-bounds,sroa-wrappers{instcombine=false instsimplify=true},canonicalize,sroa-wrappers{instcombine=false instsimplify=true},libdevice-funcs-raise,canonicalize,remove-duplicate-func-def,canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform,enzyme-batch,inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform)"

PIPELINE2="--pass-pipeline=builtin.module(enzyme{postpasses=\"arith-raise{stablehlo=true},canonicalize,cse,canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math,canonicalize,cse,canonicalize\"},arith-raise{stablehlo=true})"

PIPELINE3="--pass-pipeline=builtin.module(canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math,inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform,lower-kernel{backend=cpu},canonicalize,canonicalize,llvm-to-memref-access,canonicalize,convert-llvm-to-cf,canonicalize,enzyme-lift-cf-to-scf,canonicalize,func.func(canonicalize-loops),canonicalize-scf-for,canonicalize,affine-cfg,canonicalize,func.func(canonicalize-loops),canonicalize,func.func(llvm-to-affine-access),canonicalize,delinearize-indexing,canonicalize,simplify-affine-exprs,affine-cfg,canonicalize)"

PIPELINE4="--pass-pipeline=builtin.module(func.func(affine-loop-unroll{unroll-full=true}),canonicalize,raise-affine-to-stablehlo,arith-raise{stablehlo=true},inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform)"

./bazel-bin/enzymexlamlir-opt unopt.mlir "$PIPELINE1" -o mid1.mlir
./bazel-bin/enzymexlamlir-opt mid1.mlir "$PIPELINE2" -o mid2.mlir
./bazel-bin/enzymexlamlir-opt mid2.mlir "$PIPELINE3" -o mid3.mlir
./bazel-bin/enzymexlamlir-opt mid3.mlir "$PIPELINE4" -o mid4.mlir

[ivanradanov]

With affine loop unroll

#!/usr/bin/env sh

PIPELINE1="--pass-pipeline=builtin.module(inline{default-pipeline=canonicalize max-iterations=4},propagate-constant-bounds,sroa-wrappers{instcombine=false instsimplify=true},canonicalize,sroa-wrappers{instcombine=false instsimplify=true},libdevice-funcs-raise,canonicalize,remove-duplicate-func-def,canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform,enzyme-batch,inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform)"

PIPELINE2="--pass-pipeline=builtin.module(enzyme{postpasses=\"arith-raise{stablehlo=true},canonicalize,cse,canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math,canonicalize,cse,canonicalize\"},arith-raise{stablehlo=true})"

PIPELINE3="--pass-pipeline=builtin.module(canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math,inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform,lower-kernel{backend=cpu},canonicalize,canonicalize,llvm-to-memref-access,canonicalize,convert-llvm-to-cf,canonicalize,enzyme-lift-cf-to-scf,canonicalize,func.func(canonicalize-loops),canonicalize-scf-for,canonicalize,affine-cfg,canonicalize,func.func(canonicalize-loops),canonicalize,llvm-to-affine-access,canonicalize,delinearize-indexing,canonicalize,simplify-affine-exprs,affine-cfg,canonicalize)"

PIPELINE4="--pass-pipeline=builtin.module(func.func(affine-loop-unroll{unroll-full=true}),canonicalize,raise-affine-to-stablehlo,arith-raise{stablehlo=true},inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;gather_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;reorder_elementwise_and_shape_op<16>;chlo_inf_const_prop<16>;gamma_const_prop<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_pad<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;negate_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;sqrt_simplify<16>;cos_simplify<16>;sin_simplify<16>;noop_slice<16>;noop_reverse<16>;const_prop_through_barrier<16>;slice_slice<16>;shift_right_logical_simplify<16>;pad_simplify<16>;negative_pad_to_slice<16>;tanh_simplify<16>;exp_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;gather_simplify<16>;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;slice_pad<1>;dot_reshape_dot<1>;concat_const_prop<1>;concat_fuse<1>;pad_reshape_pad<1>;pad_pad<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;scatter_to_dynamic_update_slice<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;reshape_iota<16>;slice_reshape_slice<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;add_pad_pad_to_concat<1>;broadcast_reshape<1>;slice_reshape_concat<1>;slice_reshape_elementwise<1>;slice_reshape_transpose<1>;slice_reshape_dot_general<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;slice_reshape_pad<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;transpose_pad<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_einsum<1>;einsum_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;concat_to_pad<1>;concat_appending_reshape<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;dot_reshape_pad<1>;pad_dot_general<1>(0);dot_reshape_pad<1>;pad_dot_general<1>(1);if_inline<1>;if_to_select<1>;dynamic_update_slice_const_prop;dynamic_gather_op_is_not_dynamic<16>;divide_sqrt_to_multiply_rsqrt<16>;binary_op_transpose_simplify_add;binary_op_transpose_simplify_sub;binary_op_transpose_simplify_mul;binary_op_transpose_simplify_div;binary_op_transpose_simplify_min;binary_op_transpose_simplify_max;binary_op_transpose_simplify_pow;binary_op_transpose_simplify_rem;binary_op_transpose_simplify_or;binary_op_transpose_simplify_and;binary_op_transpose_simplify_xor;associative_binary_op_reordering<1>;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;scatter_indices_are_unique;transpose_reduce_simplify;replace_neg_add_with_subtract;log_const_prop<1>;log_plus_one_const_prop<1>;binop_const_simplify;transpose_broadcast_in_dim_to_broadcast_in_dim;not_select_simplify;scatter_update_computation_const_prop;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>;scatter_update_computation_const_prop;if_remove_unused;no_nan_add_sub_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform)"

./bazel-bin/enzymexlamlir-opt unopt.mlir "$PIPELINE1" -o mid1.mlir
./bazel-bin/enzymexlamlir-opt mid1.mlir "$PIPELINE2" -o mid2.mlir
./bazel-bin/enzymexlamlir-opt mid2.mlir "$PIPELINE3" -o mid3.mlir
./bazel-bin/enzymexlamlir-opt mid3.mlir "$PIPELINE4" -o mid4.mlir

[wsmoses]

we should generally try to keep that in sync with the pipeline run by default, here: https://github.com/EnzymeAD/Reactant.jl/blob/b36f525f525aa3fc7cd1fa68344a3441680100fa/src/Compiler.jl#L778