aboutsummaryrefslogtreecommitdiffhomepage
path: root/ptx/tools/implicit_ld_src.py
blob: cb6780eba6761cab322d84e4982ca5954f98b37a (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
import os
import subprocess
import tempfile

types = ["b8", "b16", "b32", "b64", "u8", "u16", "u32", "u64", "s8", "s16", "s32", "s64", "f32", "f64"]

for input_type in types:
    for op_type in types:
        with tempfile.TemporaryDirectory() as dir:
            f_name = os.path.join(dir, 'ptx')
            out_name = os.path.join(dir, 'out')
            with open(f_name, 'w') as f:
                f.write(
                f"""
                .version 6.5
                .target sm_30
                .address_size 64
                .visible .entry VecAdd_kernel(
                    .param .{input_type} input
                )
                {{
                    .reg.{op_type}        r1;
                    ld.param.{op_type} 	r1, [input];
                    ret;
                }}
                """)
            err = subprocess.run(f"ptxas {f_name} -o {out_name}")
            if err.returncode == 0:
                print(f"{op_type} {input_type}")
            else:
                print(f"[INVALID] {op_type} {input_type}")