Revert "Add .warmup() for triton.jit()" (#682)

Reverts openai/triton#671

It seems like for some reason this caused out-of-memory errors on some
of our internal workloads. I'm reverting this so that HEAD can be used
in production at OpenAI, and I will work on digging into this issue
asynchronously.
This commit is contained in:
Philippe Tillet
2022-09-20 16:05:14 -07:00
committed by GitHub
parent 48f30550f1
commit 7dc2a70edb
4 changed files with 19 additions and 90 deletions

View File

@@ -150,25 +150,3 @@ def test_constexpr_not_callable() -> None:
except BaseException:
error = True
assert error is True
def test_jit_warmup_cache() -> None:
@triton.jit
def kernel_add(a, b, o, N: tl.constexpr):
idx = tl.arange(0, N)
tl.store(o + idx,
tl.load(a + idx) + tl.load(b + idx))
args = [
torch.randn(32, dtype=torch.float32, device="cuda"),
torch.randn(32, dtype=torch.float32, device="cuda"),
torch.randn(32, dtype=torch.float32, device="cuda"),
32,
]
assert len(kernel_add.cache) == 0
kernel_add[(1,)].warmup(torch.float32, torch.float32, torch.float32, 32)
assert len(kernel_add.cache) == 1
kernel_add[(1,)].warmup(*args)
assert len(kernel_add.cache) == 1
kernel_add[(1,)](*args)
assert len(kernel_add.cache) == 1