fuse LARS

[Qazalin]

631 -> 472 in BACKWARD=1 python3 examples/handcode_resnet50_opt.py
lazy.py allows fusion of assign like:

#include <metal_stdlib>
using namespace metal;
kernel void E_4_4(device float* data0, const device float* data1, const device float* data2, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {
  int gidx0 = gid.y; /* 4 */
  int gidx1 = gid.x; /* 4 */
  float val0 = *(data1+gidx0);
  float val1 = *(data2+gidx1);
  *(data0+(gidx0*4)+gidx1) = (val0*val1);
}

#include <metal_stdlib>
using namespace metal;
kernel void E_16n1(device float* data0, const device float* data1, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {
  int gidx0 = gid.x; /* 16 */
  float val0 = *(data0+gidx0);
  float val1 = *(data1+gidx0);
  *(data0+gidx0) = ((0.9f*val0)+val1);
}

and this is wrong:

#include <metal_stdlib>
using namespace metal;
kernel void E_4_4n3(device float* data0, const device float* data1, const device float* data2, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {
  int gidx0 = gid.y; /* 4 */
  int gidx1 = gid.x; /* 4 */
  int alu0 = ((gidx0*4)+gidx1);
  float val0 = *(data0+alu0);
  float val1 = *(data1+gidx0);
  float val2 = *(data2+gidx1);
  *(data0+alu0) = ((0.9f*val0)+(val1*val2));
}

[github-actions]

This branch currently is behind tinygrad/master. The line count difference bot is disabled.