Softmax toy example runs fine once per each clean install from source on Metal

[Laz4rz]

Ok, I do not think I've seen this one here anywhere. I am running a toy example of tinygrad softmax on MBP Pro 16 2023 (M2 Pro). Built from source. I'll happily provide you with any needed details.

from tinygrad import Tensor, Device

p = Tensor([0.50, 0.60, 0.70, 0.30, 0.25])
s = p.softmax()

print(s.numpy())

this runs fine exactly ONCE after clean install from source. It runs (exacltly once) both on native system python and conda python.

On any other consecutive run I get:

➜  tracing_comparison git:(main) ✗ python3 trace_tinygrad_softmax.py
Traceback (most recent call last):
  File "/Users/mikolajboronski/github/jax/digression/tracing_comparison/trace_tinygrad_softmax.py", line 6, in <module>
    print(s.numpy())
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/tensor.py", line 194, in numpy
    return np.frombuffer(self._data(), dtype=self.dtype.np).reshape(self.shape)
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/tensor.py", line 177, in _data
    cpu = self.cast(self.dtype.scalar()).contiguous().to("CLANG").realize()
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/tensor.py", line 144, in realize
    Tensor.corealize([self])
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/tensor.py", line 141, in corealize
    run_schedule(create_schedule(flatten([x.lazydata.lbs if isinstance(x.lazydata, MultiLazyBuffer) else [x.lazydata] for x in lst])))
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/engine/realize.py", line 32, in run_schedule
    prg = lower_schedule_item(si)
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/engine/realize.py", line 16, in lower_schedule_item
    if si.ast[0].op is BufferOps.STORE: return Device[si.outputs[0].device].get_runner(*si.ast)
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/device.py", line 257, in get_runner
    def get_runner(self, *ast:LazyOp) -> CompiledASTRunner: return self.to_program(self.get_linearizer(*ast))
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/device.py", line 221, in to_program
    ret = CompiledASTRunner(k.name, self.compiler.render(to_function_name(k.name), k.uops), self.dname, k.global_size, k.local_size,
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/device.py", line 172, in __init__
    self.lib, self.clprg, self.outcount = lib, self.device.runtime(self.name, lib), outcount
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/runtime/ops_metal.py", line 39, in __init__
    assert lib[:4] == b"MTLB", "Invalid Metal library. Could be due to using conda. Try system python or METAL_XCODE=1 DISABLE_COMPILER_CACHE=1."
AssertionError: Invalid Metal library. Could be due to using conda. Try system python or METAL_XCODE=1 DISABLE_COMPILER_CACHE=1.

after setting:

export METAL_XCODE=1
export DISABLE_COMPILER_CACHE=1

On conda: nothing changes.
On native:

➜  tracing_comparison git:(main) ✗ python3 trace_tinygrad_softmax.py
<tinygrad.device._Device object at 0x104592d00>
xcrun: error: unable to find utility "metal", not a developer tool or in PATH
Traceback (most recent call last):
  File "/Users/mikolajboronski/github/jax/digression/tracing_comparison/trace_tinygrad_softmax.py", line 9, in <module>
    print(s.numpy())
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/tensor.py", line 194, in numpy
    return np.frombuffer(self._data(), dtype=self.dtype.np).reshape(self.shape)
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/tensor.py", line 177, in _data
    cpu = self.cast(self.dtype.scalar()).contiguous().to("CLANG").realize()
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/tensor.py", line 144, in realize
    Tensor.corealize([self])
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/tensor.py", line 141, in corealize
    run_schedule(create_schedule(flatten([x.lazydata.lbs if isinstance(x.lazydata, MultiLazyBuffer) else [x.lazydata] for x in lst])))
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/engine/realize.py", line 32, in run_schedule
    prg = lower_schedule_item(si)
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/engine/realize.py", line 16, in lower_schedule_item
    if si.ast[0].op is BufferOps.STORE: return Device[si.outputs[0].device].get_runner(*si.ast)
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/device.py", line 257, in get_runner
    def get_runner(self, *ast:LazyOp) -> CompiledASTRunner: return self.to_program(self.get_linearizer(*ast))
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/device.py", line 221, in to_program
    ret = CompiledASTRunner(k.name, self.compiler.render(to_function_name(k.name), k.uops), self.dname, k.global_size, k.local_size,
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/device.py", line 171, in __init__
    lib:bytes = precompiled if precompiled is not None else self.device.compiler.compile_cached(prg)
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/device.py", line 157, in compile_cached
    lib = self.compile(src)
  File "/Users/mikolajboronski/github/tinygrad/tinygrad/tinygrad/runtime/ops_metal.py", line 23, in compile
    air = subprocess.check_output(['xcrun', '-sdk', 'macosx', 'metal', '-x', 'metal', '-c', '-', '-o', '-'], input=src.encode('utf-8'))
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.9/lib/python3.9/subprocess.py", line 424, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.9/lib/python3.9/subprocess.py", line 528, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['xcrun', '-sdk', 'macosx', 'metal', '-x', 'metal', '-c', '-', '-o', '-']' returned non-zero exit status 72.

[ignaciosica]

it seems it doesn't find metal (xcrun: error: unable to find utility "metal", not a developer tool or in PATH)

try the following to reinstall the Command Line Tools
sudo rm -rf /Library/Developer/CommandLineTools
xcode-select --install

this may also fix your issue
xcode-select --switch /Applications/Xcode.app/Contents/Developer

[Laz4rz]

Ok, thanks for the tip but none of these helped. Or maybe first?

It is now working without run limits. What I did after the initial changes:

1. tried

xcode-select --switch /Applications/Xcode.app/Contents/Developer

but got no Xcode.app path, so nothing changed

2. tried

sudo rm -rf /Library/Developer/CommandLineTools
xcode-select --install

and rerun softmax example, resulted in:

xcrun: error: unable to find utility "metal", not a developer tool or in PATH

again.

3. Changed METAL_XCODE back to 0.

This fixed all of the previous problems.

So my current situation is DISABLE_COMPILER_CACHE=1, METAL_XCODE=0 and clean CommandLineTools install, running on native Python. I'd maybe add the information that setting METAL_XCODE to 1 may break the working of metal to the troubleshooting info, something like:

AssertionError: Invalid Metal library. Could be due to using conda. Try system python or METAL_XCODE=1 DISABLE_COMPILER_CACHE=1. 
Setting METAL_XCODE=1 may result in system not being able to find metals. 

[Laz4rz]

Restarted everything. What seems to be fixing the issue reproductively is:

1. Deactivating Conda
2. DISABLE_COMPILER_CACHE=1
3. Run by pointing to the exact bin of native python with tinygrad installed from source