gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add bf16/fp16/fp64 support for ty_to_cpp (#800)

Browse Source 
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.
This commit is contained in:
Yanbo Liang
2022-10-24 19:41:25 -07:00
committed by GitHub
parent db3aa1d1fb
commit 5ca1ed0101
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/triton/compiler.py
View File

@@ -913,7 +913,10 @@ def ty_to_cpp(ty):
"i64": "int64_t", "i64": "int64_t",
"u32": "uint32_t", "u32": "uint32_t",
"u64": "uint64_t", "u64": "uint64_t",
"fp16": "float",
"bf16": "float",
"fp32": "float", "fp32": "float",
"fp64": "double",
}[ty] }[ty]
@@ -943,6 +946,8 @@ def generate_launcher(identifier, constants, signature):
'i64': 'int64_t', 'i64': 'int64_t',
'u32': 'uint32_t', 'u32': 'uint32_t',
'u64': 'uint64_t', 'u64': 'uint64_t',
'fp16': 'float',
'bf16': 'float',
'fp32': 'float', 'fp32': 'float',
'fp64': 'double', 'fp64': 'double',
}[ty] }[ty]