Add sort memory (#1113)

* Add sort memory

* Update Compiler.jl

* extra while concat

* Update Project.toml

---------

Co-authored-by: William Moses <[email protected]>

Author: giordano

Project.toml

@@ -87,7 +87,7 @@ PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.9"
-Reactant_jll = "0.0.117"
+Reactant_jll = "0.0.119"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

src/Compiler.jl

@@ -457,6 +457,8 @@ function create_result(
     return Meta.quot(tocopy)
 end
 
+const WHILE_CONCAT = Ref(false)
+
 # Optimization passes via transform dialect
 function optimization_passes(;
     no_nan::Bool=false,
@@ -638,8 +640,13 @@ function optimization_passes(;
         "broadcastindim_is_reshape",
         "slice_reduce_window<1>",
         "while_deadresult",
+        "while_dus",
     ]
 
+    if WHILE_CONCAT[]
+        push!(transform_passes_list, "while_concat")
+    end
+
     if reshape_propagate === :up
         append!(
             transform_passes_list,
@@ -1038,7 +1045,7 @@ function compile_mlir!(
     elseif raise
         # Raise enabled but use default passes
         # TODO remove redundant libdevice raise after fixing phase ordering
-        "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,libdevice-funcs-raise,canonicalize,affine-cfg,canonicalize,func.func(canonicalize-loops),canonicalize,llvm-to-affine-access,canonicalize,delinearize-indexing,canonicalize,simplify-affine-exprs,affine-cfg,canonicalize,func.func(affine-loop-invariant-code-motion),canonicalize,raise-affine-to-stablehlo{prefer_while_raising=false dump_failed_lockstep=$(DUMP_FAILED_LOCKSTEP[])},canonicalize,arith-raise{stablehlo=true}," *
+        "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,libdevice-funcs-raise,canonicalize,affine-cfg,canonicalize,func.func(canonicalize-loops),canonicalize,llvm-to-affine-access,canonicalize,delinearize-indexing,canonicalize,simplify-affine-exprs,affine-cfg,canonicalize,func.func(affine-loop-invariant-code-motion),canonicalize,sort-memory,raise-affine-to-stablehlo{prefer_while_raising=false dump_failed_lockstep=$(DUMP_FAILED_LOCKSTEP[])},canonicalize,arith-raise{stablehlo=true}," *
         opt_passes2
     else
         "canonicalize"