From cd61e81f551416330f7375d256989ca46a92838e Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Wed, 16 Oct 2024 23:00:05 +0800
Subject: [PATCH] beautiful mnist works on windows (#7100)

* beautiful mnist works on windows [pr]

* add comment for that (no pr)
---
 tinygrad/dtype.py            | 4 ++--
 tinygrad/helpers.py          | 4 ++++
 tinygrad/runtime/ops_disk.py | 8 +++++---
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py
index 1b376f44..429419f6 100644
--- a/tinygrad/dtype.py
+++ b/tinygrad/dtype.py
@@ -91,8 +91,8 @@ class dtypes:
   uint16: Final[DType] = DType(4, 2, "unsigned short", 'H', 1)
   int32: Final[DType] = DType(5, 4, "int", 'i', 1)
   uint32: Final[DType] = DType(6, 4, "unsigned int", 'I', 1)
-  int64: Final[DType] = DType(7, 8, "long", 'l', 1)
-  uint64: Final[DType] = DType(8, 8, "unsigned long", 'L', 1)
+  int64: Final[DType] = DType(7, 8, "long", 'q', 1)
+  uint64: Final[DType] = DType(8, 8, "unsigned long", 'Q', 1)
   float16: Final[DType] = DType(9, 2, "half", 'e', 1)
   # bfloat16 has higher priority than float16, so least_upper_dtype(dtypes.int64, dtypes.uint64) = dtypes.float16
   bfloat16: Final[DType] = DType(10, 2, "__bf16", None, 1)
diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py
index 53c74e6b..b4b26061 100644
--- a/tinygrad/helpers.py
+++ b/tinygrad/helpers.py
@@ -13,8 +13,12 @@ def prod(x:Iterable[T]) -> Union[T,int]: return functools.reduce(operator.mul, x
 
 # NOTE: helpers is not allowed to import from anything else in tinygrad
 OSX = platform.system() == "Darwin"
+WIN = platform.system() == "Windows"
 CI = os.getenv("CI", "") != ""
 
+# fix colors on Windows, https://stackoverflow.com/questions/12492810/python-how-can-i-make-the-ansi-escape-codes-to-work-also-in-windows
+if WIN: os.system("")
+
 def dedup(x:Iterable[T]): return list(dict.fromkeys(x))   # retains list order
 def argfix(*x):
   if x and x[0].__class__ in (tuple, list):
diff --git a/tinygrad/runtime/ops_disk.py b/tinygrad/runtime/ops_disk.py
index 8d60a534..a78b0d96 100644
--- a/tinygrad/runtime/ops_disk.py
+++ b/tinygrad/runtime/ops_disk.py
@@ -1,9 +1,11 @@
 from __future__ import annotations
-import os, sys, mmap, _posixshmem, io, ctypes, ctypes.util, platform, contextlib
+import os, sys, mmap, io, ctypes, ctypes.util, platform, contextlib
 from typing import Optional, Generator, Tuple, Callable, List
 from tinygrad.helpers import OSX, round_up
 from tinygrad.device import Compiled, Allocator
-from tinygrad.runtime.autogen import io_uring, libc
+with contextlib.suppress(ImportError):
+  import _posixshmem
+  from tinygrad.runtime.autogen import io_uring, libc
 
 class DiskBuffer:
   def __init__(self, device:DiskDevice, size:int, offset=0):
@@ -87,7 +89,7 @@ class DiskDevice(Compiled):
       self.mem = mmap.mmap(fd, self.size, mmap.MAP_SHARED | MAP_POPULATE | MAP_LOCKED)
       os.close(fd)
     else:
-      try: self.fd = os.open(filename, os.O_RDWR|os.O_CREAT|(0 if OSX else os.O_DIRECT))
+      try: self.fd = os.open(filename, os.O_RDWR|os.O_CREAT|getattr(os, "O_DIRECT", 0))
       except OSError: self.fd = os.open(filename, os.O_RDWR|os.O_CREAT)
       if os.fstat(self.fd).st_size < self.size: os.ftruncate(self.fd, self.size)
       self.mem = mmap.mmap(self.fd, self.size)