[inductor] `float64` KeyError

[anijain2305]

🐛 Describe the bug

Following repro fails

Error logs

Traceback (most recent call last):
  File "/scratch/anijain/work/pytorch/repro.py", line 45, in <module>
    compiled = compile_fx_inner(mod, args)
  File "/scratch/anijain/work/pytorch/torch/_dynamo/debug_utils.py", line 471, in debug_wrapper
    compiled_fn = compiler_fn(gm, example_inputs, **kwargs)
  File "/scratch/anijain/work/pytorch/torch/_inductor/debug.py", line 178, in inner
    return fn(*args, **kwargs)
  File "/scratch/anijain/work/env/lib/python3.9/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/scratch/anijain/work/pytorch/torch/_inductor/compile_fx.py", line 123, in compile_fx_inner
    compiled_fn = graph.compile_to_fn()
  File "/scratch/anijain/work/pytorch/torch/_inductor/graph.py", line 338, in compile_to_fn
    return self.compile_to_module().call
  File "/scratch/anijain/work/pytorch/torch/_dynamo/utils.py", line 85, in time_wrapper
    r = func(*args, **kwargs)
  File "/scratch/anijain/work/pytorch/torch/_inductor/graph.py", line 328, in compile_to_module
    mod = PyCodeCache.load(code)
  File "/scratch/anijain/work/pytorch/torch/_inductor/codecache.py", line 212, in load
    exec(code, mod.__dict__, mod.__dict__)
  File "/tmp/torchinductor_anijain/h3/ch3eodunx3i3wmnkeoqrc35vsmeyvk464ff3agm4pzftyiopn277.py", line 41, in <module>
    async_compile.wait(globals())
  File "/scratch/anijain/work/pytorch/torch/_inductor/codecache.py", line 352, in wait
    scope[key] = result.result()
  File "/scratch/anijain/work/pytorch/torch/_inductor/codecache.py", line 259, in result
    self.future.result()
  File "/scratch/anijain/work/env/lib/python3.9/concurrent/futures/_base.py", line 446, in result
    return self.__get_result()
  File "/scratch/anijain/work/env/lib/python3.9/concurrent/futures/_base.py", line 391, in __get_result
    raise self._exception
KeyError: 'fp64'

Minified repro

import torch._inductor.overrides

import torch
from torch import tensor, device
import torch.fx as fx
from torch._dynamo.testing import rand_strided
from math import inf
from torch.fx.experimental.proxy_tensor import make_fx

# torch version: 1.14.0a0+git8a8cd09
# torch cuda version: 11.6
# torch git version: 8a8cd092c8537b226f5c38ed88bc07e181b0946c


# CUDA Info:
# nvcc: NVIDIA (R) Cuda compiler driver
# Copyright (c) 2005-2022 NVIDIA Corporation
# Built on Thu_Feb_10_18:23:41_PST_2022
# Cuda compilation tools, release 11.6, V11.6.112
# Build cuda_11.6.r11.6/compiler.30978841_0

# GPU Hardware Info:
# NVIDIA A100-SXM4-40GB : 8


from torch.nn import *
class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()



    def forward(self, arg104_1, bmm):
        div = torch.ops.aten.div.Tensor(bmm, arg104_1);  bmm = arg104_1 = None
        return (div,)

args = [((), (), torch.float64, 'cpu'), ((80, 204, 204), (41616, 204, 1), torch.float32, 'cuda')]
args = [rand_strided(sh, st, dt, dev) for (sh, st, dt, dev) in args]
mod = make_fx(Repro().to(device="cuda"))(*args)
mod(*args)

from torch._inductor.compile_fx import compile_fx_inner
from torch._dynamo.debug_utils import same_two_models

compiled = compile_fx_inner(mod, args)
compiled(args)

[anijain2305]

cc @ngimel @yanboliang

[ngimel]

Hm can't repro on master
Edit: I didn't have @yanboliang's 0d cpu PR on my build, I suspect it's because of it.

[ngimel]

THe fix for this should be in triton, we should add "fp64": "double" to ty_to_cpp, @yanboliang can you please submit that PR?

[yanboliang]

THe fix for this should be in triton, we should add "fp64": "double" to ty_to_cpp, @yanboliang can you please submit that PR?

Yes, it's because of ty_to_cpp, will send a PR.

[yanboliang]

Just sent triton-lang/triton#800 to fix it in triton. It includes triton supports for fp16/fp32 as well. I'll remove my workaround for fp16/bf16 and add unit tests in Inductor after that PR gets in.

[voznesenskym]

Just sent openai/triton#800 to fix it in triton. It includes triton supports for fp16/fp32 as well. I'll remove my workaround for fp16/bf16 and add unit tests in Inductor after that PR gets in.

Im not sure mapping fp16 and bf16 to fp32 is always the right thing to do here...

[ngimel]

Actual argument (in our case) will be python float, because it is obtained from tensor.item() but I'm not sure that holds for triton in general

[yanboliang]

Actual argument (in our case) will be python float, because it is obtained from tensor.item() but I'm not sure that holds for triton in general

Yes, it maps to python float type finally. The way I did to return fp32 for bf16/fp16 is a workaround bc triton is lacking support for it. When triton-lang/triton#800 gets in and we upgrade trition version, we can remove this workaround.

[ngimel]

I still think that maybe we shouldn't add fp16/fb16 mapping to triton, and instead let our arg annotations to be fp32 (as we have now), because, after all, that's a more correct annotation for a python float that we are actually sending.