set testpath on pytest

* enable rangeify const folding

* renumber ranges for kernel deduping

examples/mlperf/scripts/stable_diffusion_downloads.sh

@@ -0,0 +1,57 @@
+#!/usr/bin/env bash
+# adapted from https://github.com/mlcommons/training/blob/4bdf5c8ed218ad76565a2ba1ac27c919ccc6d689/stable_diffusion/README.md
+
+# setup dirs
+
+DATA=/raid/datasets/stable_diffusion
+
+LAION=$DATA/laion-400m/webdataset-moments-filtered 
+COCO=$DATA/coco2014
+mkdir -p $LAION $COCO
+
+CKPT=/raid/weights/stable_diffusion
+mkdir -p $CKPT/clip $CKPT/sd $CKPT/inception
+
+# download data
+
+# if rclone isn't installed system-wide / in your PATH, put the executable path in quotes below
+#RCLONE=""
+RCLONE="rclone"
+
+## VAE-encoded image latents, from 6.1M image subset of laion-400m
+## about 1 TB for whole download
+$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
+$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/ ${LAION} --include="*.tar" -P
+$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/sha512sums.txt ${LAION} -P
+cd $LAION && grep -E '\.tar$' sha512sums.txt | sha512sum -c --quiet - && \
+  echo "All .tar files verified" || { echo "Checksum failure when validating downloaded Laion moments"; exit 1; }
+
+## prompts and FID statistics from 30k image subset of coco2014
+## 33 MB
+$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
+$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k.tsv ${COCO} -P
+
+$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
+$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k_stats.npz ${COCO} -P
+
+# download checkpoints
+
+## clip (needed for text and vision encoders for validation)
+CLIP_WEIGHTS_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin"
+CLIP_WEIGHTS_SHA256="9a78ef8e8c73fd0df621682e7a8e8eb36c6916cb3c16b291a082ecd52ab79cc4"
+CLIP_CONFIG_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/raw/main/open_clip_config.json"
+wget -N -P ${CKPT}/clip ${CLIP_WEIGHTS_URL}
+wget -N -P ${CKPT}/clip ${CLIP_CONFIG_URL}
+echo "${CLIP_WEIGHTS_SHA256}  ${CKPT}/clip/open_clip_pytorch_model.bin"                    | sha256sum -c
+
+## sd (needed for latent->image decoder for validation, also has clip text encoder for training)
+SD_WEIGHTS_URL='https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt'
+SD_WEIGHTS_SHA256="d635794c1fedfdfa261e065370bea59c651fc9bfa65dc6d67ad29e11869a1824"
+wget -N -P ${CKPT}/sd ${SD_WEIGHTS_URL}
+echo "${SD_WEIGHTS_SHA256}  ${CKPT}/sd/512-base-ema.ckpt"                    | sha256sum -c
+
+## inception (needed for validation)
+FID_WEIGHTS_URL='https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'
+FID_WEIGHTS_SHA1="bd836944fd6db519dfd8d924aa457f5b3c8357ff"
+wget -N -P ${CKPT}/inception ${FID_WEIGHTS_URL}
+echo "${FID_WEIGHTS_SHA1}  ${CKPT}/inception/pt_inception-2015-12-05-6726825d.pth"                    | sha1sum -c

extra/models/inception.py

@@ -270,8 +270,10 @@ class FidInceptionV3:
     self.Mixed_7b = inception.Mixed_7b
     self.Mixed_7c = inception.Mixed_7c
 
-  def load_from_pretrained(self):
-    state_dict = torch_load(str(fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")))
+  def load_from_pretrained(self, path=None):
+    if path is None:
+      path = fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")
+    state_dict = torch_load(str(path))
     for k,v in state_dict.items():
       if k.endswith(".num_batches_tracked"):
         state_dict[k] = v.reshape(1)

pytest.ini

@@ -3,3 +3,4 @@ norecursedirs = extra
 timeout = 180
 timeout_method = thread
 timeout_func_only = true
+testpaths = test

test/test_symbolic_jit.py

@@ -16,6 +16,18 @@ class TestSymbolicJit(unittest.TestCase):
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
+  @unittest.expectedFailure # TODO: fix, this works without jit
+  def test_plus1_pad(self):
+    def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = jf(a[:, :vi]).numpy()
+      expected = f(a[:, :i]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
   def test_add(self):
     def f(a, b): return (a+b).realize()
     jf = TinyJit(f)

test/test_symbolic_ops.py

@@ -17,6 +17,15 @@ class TestSymbolicOps(unittest.TestCase):
       expected = f(a[:, :i]).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
+  def test_plus1_pad(self):
+    def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
+    a = Tensor.rand(3, 10)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = f(a[:, :vi]).numpy()
+      expected = f(a[:, :i]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
   def test_add(self):
     def f(a, b): return (a+b).realize()
     a = Tensor.rand(3, 10)

tinygrad/schedule/rangeify.py

@@ -458,6 +458,7 @@ class LocalAddBufferContext:
   dg:int = 0
   map:dict = field(default_factory=dict)
   vars:dict = field(default_factory=dict)
+  range:int = 0
 
 def debuf(ctx:LocalAddBufferContext, buf:UOp):
   ret = UOp(Ops.DEFINE_GLOBAL, buf.dtype.ptr(buf.arg), arg=ctx.dg)
@@ -477,6 +478,12 @@ def handle_assign(ctx:LocalAddBufferContext, assign:UOp):
   ctx.map[buf] = assign
   return buf
 
+def renumber_range(ctx:LocalAddBufferContext, r:UOp):
+  if r.tag is not None: return None
+  ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=())
+  ctx.range += 1
+  return ret
+
 to_define_global = PatternMatcher([
   (UPat(Ops.BUFFER, name="buf"), debuf),
   (UPat(Ops.BIND, name="b"), unbind_kernel),
@@ -485,6 +492,9 @@ to_define_global = PatternMatcher([
   # HACK in case any CONSTs were replaced
   # this is only needed if you are using symbolic
   #(UPat(Ops.CONST, name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
+
+  # renumber the ranges starting with 0 so that kernel deduping works
+  (UPat(Ops.RANGE, name="r"), renumber_range),
 ])
 
 rangeify_codegen = PatternMatcher([