reorder expand (tinygrad#9051)

* reorder expand

* symbolic ops needs resolve here

* s/arg/st + whitespace

* viz

---------

Co-authored-by: qazal <qazal.software@gmail.com>

Author: geohot

test/external/external_test_softmax_fusion.py

@@ -0,0 +1,12 @@
+from tinygrad import Tensor, GlobalCounters, Context
+
+if __name__ == "__main__":
+  with Context(TRACK_MATCH_STATS=0): test = Tensor.ones(32, 10).contiguous().realize()
+  GlobalCounters.reset()
+
+  # this is the softmax from scaled_dot_product_attention
+  # it becomes 3 kernels
+  print("*** softmax ***")
+  with Context(NOOPT=1):
+    out = test.softmax(-1)
+    out.realize()

test/test_schedule.py

@@ -2527,7 +2527,6 @@ def test_new_flat_buffer(self):
 
   # sometimes we prefer to perform an op before movement ops, in this case we should stack the mops on top of the new buffer
 
-  @unittest.expectedFailure
   def test_reorder_expand(self):
     a = Tensor.empty(4, 1)
     b = a.expand(4, 4).reciprocal()

tinygrad/engine/schedule.py

@@ -92,6 +92,9 @@ def replace_contiguous(ctx:dict[UOp, UOp], alu:UOp):
   # substitute BITCAST/CONTIGUOUS with BUFFER_VIEW on DISK
   (UPat((Ops.BITCAST, Ops.CONTIGUOUS), name="root"),
    lambda root: root.replace(op=Ops.BUFFER_VIEW) if isinstance(root.device, str) and root.device.startswith("DISK") else None),
+  # put UnaryOps before EXPANDs
+  (UPat(GroupOp.Unary, src=UPat(Ops.VIEW, src=(UPat.var("inp"),), name="v"), name="alu"),
+   lambda inp,v,alu: inp.alu(alu.op).view(v.st) if resolve(prod(alu.shape) > v.st.real_size()) else None),
 ])
 
 remove_movement_ops = merge_views+PatternMatcher([