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