[POC] emit indexing code directly

[geohotstan]

There are a bunch of rough edges here I haven't ironed out, but here's a POC.
I wrote a bunch of hacks just to get it working.

I've tried doing symbolic index tensor with a new symbolic object, but I found myself needing a new UOps to deal with the loading for the render_ops
I saw how setitem was doing the store part with ShapeTracker, so I thought on a conceptual level the advanced indexing stuff should also be apart of the ShapeTracker.

The basic idea is just creating a non-mergable special View (IndexedVIew) that contains the Tensor indices which is then global_loaded and used to determine the offsets of the Buffer that holds the IndexedView when this Buffer is loaded/stored. The kernel code just indexes directly instead of arange masking.

Currently things that work:

• Embedding
• Advanced indexing
• Multi index advanced indexing

Things that don't work:

• Backward pass :)
• Advanced setitem (bufs are not being loaded but kernel code looks correct)
• More complex operations where advanced indexing has operations before it or after it (like basic indexing + advanced indexing mix)
• Image dtype indexing? (the codepath looks wrong but I haven't tested)
• And more, lol

There are a lot of garbage hacks in this currently
The main thing that works is the index() function in linearizer that correctly determines the offset given buffers used for indexing (doesn't work in all cases yet... bunch of hacks in Tensor getitem to get tests passing)
It's gonna take some more time for me to get this to a point where it's mergeable since I still don't fully get schedule and shapetracker. Or idk is this pursuable?

[geohotstan]

speaking of leaky abstractions
how's shoving a LB into a View 😎

[github-actions]

Changes

Name                              Lines    Diff    Tokens/Line    Diff
------------------------------  -------  ------  -------------  ------
tinygrad/tensor.py                  802      +3           21.6    -0.1
tinygrad/codegen/linearizer.py      376     +39           18.9    -0.5
tinygrad/shape/view.py              260     +23           17.8    -0.4
tinygrad/engine/schedule.py         247      +4           14.0    -0.0
tinygrad/lazy.py                    160      +1           19.0    +0.0
tinygrad/nn/__init__.py              97      +1           19.3    -0.1
tinygrad/shape/shapetracker.py       95      +2           18.6    +0.3


total lines changes: +73

[chenyuxyz]

does this support backward for setitem?

[geohotstan]

does this support backward for setitem?

I haven't tried yet.
But if this were to support backward, I think I'll need to implement new functions in function.py for getitem (and maybe setitem too?). Though afterwards all indexing-based ops can use direct indexing.