xref: /llvm-project/llvm/test/CodeGen/NVPTX/ld-st-addrrspace.py (revision b279f6b098d3849f7f1c1f539b108307d5f8ae2d)
1# This test generates all variants of load/store instructions and verifies that
2# LLVM generates correct PTX for them.
3
4# RUN: %python %s > %t.ll
5# RUN: llc < %t.ll -mtriple=nvptx -mcpu=sm_30 | FileCheck -check-prefixes=CHECK,CHECK_P32 %t.ll
6# RUN: llc < %t.ll -mtriple=nvptx64 -mcpu=sm_30 | FileCheck -check-prefixes=CHECK,CHECK_P64 %t.ll
7# RUN: %if ptxas && !ptxas-12.0 %{ llc < %t.ll -mtriple=nvptx -mcpu=sm_30 | %ptxas-verify %}
8# RUN: %if ptxas %{ llc < %t.ll -mtriple=nvptx64 -mcpu=sm_30 | %ptxas-verify %}
9
10from __future__ import print_function
11
12from itertools import product
13from string import Template
14
15llvm_type_to_ptx_type = {
16    "i8": "u8",
17    "i16": "u16",
18    "i32": "u32",
19    "i64": "u64",
20    "half": "b16",
21    "<2 x half>": "b32",
22    "float": "f32",
23    "double": "f64",
24}
25
26llvm_type_to_ptx_reg = {
27    "i8": "r",
28    "i16": "r",
29    "i32": "r",
30    "i64": "rd",
31    "half": "rs",
32    "<2 x half>": "r",
33    "float": "f",
34    "double": "fd",
35}
36
37addrspace_id = {
38    "": 0,
39    ".global": 1,
40    ".shared": 3,
41    ".const": 4,
42    ".local": 5,
43    ".param": 101,
44}
45
46
47def gen_load_tests():
48    load_template = """
49define ${type} @${testname}(${type} addrspace(${asid})* %ptr) {
50; CHECK: ${testname}
51; CHECK_P32: ld${_volatile}${_volatile_as}.${ptx_type} %${ptx_reg}{{[0-9]+}}, [%r{{[0-9]+}}]
52; CHECK_P64: ld${_volatile}${_volatile_as}.${ptx_type} %${ptx_reg}{{[0-9]+}}, [%rd{{[0-9]+}}]
53; CHECK: ret
54  %p = ${generic_ptr}
55  %a = load ${volatile} ${type}, ${type}* %p
56  ret ${type} %a
57}
58"""
59    for op_type, volatile, space in product(
60        ["i8", "i16", "i32", "i64", "half", "float", "double", "<2 x half>"],
61        [True, False],  # volatile
62        ["", ".shared", ".global", ".const", ".local", ".param"],
63    ):
64
65        # Volatile is only supported for global, shared and generic.
66        if volatile and not space in ["", ".global", ".shared"]:
67            continue
68
69        # Volatile is only supported for global, shared and generic.
70        # All other volatile accesses are done in generic AS.
71        if volatile and not space in ["", ".global", ".shared"]:
72            volatile_as = ""
73        else:
74            volatile_as = space
75
76        params = {
77            "type": op_type,
78            "volatile": "volatile" if volatile else "",
79            "_volatile": ".volatile" if volatile else "",
80            "_volatile_as": volatile_as,
81            "_space": space,
82            "ptx_reg": llvm_type_to_ptx_reg[op_type],
83            "ptx_type": llvm_type_to_ptx_type[op_type],
84            "asid": addrspace_id[space],
85        }
86
87        testname = Template("ld_${_volatile}${_space}.${ptx_type}").substitute(params)
88        params["testname"] = testname.replace(".", "_")
89
90        # LLVM does not accept "addrspacecast Type* addrspace(0) to Type*", so we
91        # need to avoid it for generic pointer tests.
92        if space:
93            generic_ptr_template = (
94                "addrspacecast ${type} addrspace(${asid})* %ptr " "to ${type}*"
95            )
96        else:
97            generic_ptr_template = "select i1 true, ${type}* %ptr, ${type}* %ptr"
98        params["generic_ptr"] = Template(generic_ptr_template).substitute(params)
99
100        print(Template(load_template).substitute(params))
101
102
103def main():
104    gen_load_tests()
105
106
107main()
108