From 5ca1ed01016530056c4507661c24d6c21efc983d Mon Sep 17 00:00:00 2001 From: Yanbo Liang Date: Mon, 24 Oct 2022 19:41:25 -0700 Subject: [PATCH] Add bf16/fp16/fp64 support for ty_to_cpp (#800) In ```torch._inductor```, we [convert 0d CPU tensor to scalar during triton codegen](https://github.com/pytorch/pytorch/pull/87329), so need add missing triton support for bf16/fp16/fp64. --- python/triton/compiler.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/python/triton/compiler.py b/python/triton/compiler.py index 1332f2c76..ab7733b60 100644 --- a/python/triton/compiler.py +++ b/python/triton/compiler.py @@ -913,7 +913,10 @@ def ty_to_cpp(ty): "i64": "int64_t", "u32": "uint32_t", "u64": "uint64_t", + "fp16": "float", + "bf16": "float", "fp32": "float", + "fp64": "double", }[ty] @@ -943,6 +946,8 @@ def generate_launcher(identifier, constants, signature): 'i64': 'int64_t', 'u32': 'uint32_t', 'u64': 'uint64_t', + 'fp16': 'float', + 'bf16': 'float', 'fp32': 'float', 'fp64': 'double', }[ty]