composable_kernel/dispatcher/examples/gemm/cpp/05_json_export.cpp at develop · ROCm/composable_kernel · GitHub

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
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

/**
 * Example 05: JSON Export
 *
 * Demonstrates exporting registry information to JSON format.
 *
 * Build: cd dispatcher/build && cmake .. && make gemm_05_json_export
 */

#include <hip/hip_runtime.h>
#include <iostream>
#include <fstream>

#include "ck_tile/dispatcher.hpp"
#include "ck_tile/dispatcher/kernel_decl.hpp"
#include "ck_tile/dispatcher/example_args.hpp"

using namespace ck_tile::dispatcher;
using namespace ck_tile::dispatcher::utils;
using Signature = decl::Signature;
using Algorithm = decl::Algorithm;

// =============================================================================
// KERNEL SET: Multiple kernels for JSON export demo
// =============================================================================

DECL_KERNEL_SET(json_export_kernels,
                .add(Signature().dtype("fp16").layout("rcr"),
                     Algorithm()
                         .tile(64, 64, 32)
                         .wave(2, 2, 1)
                         .warp(32, 32, 16)
                         .pipeline("compv3")
                         .scheduler("intrawave")
                         .epilogue("cshuffle"),
                     "gfx942")
                    .add(Signature().dtype("fp16").layout("rcr"),
                         Algorithm()
                             .tile(128, 128, 64)
                             .wave(2, 2, 1)
                             .warp(32, 32, 16)
                             .pipeline("compv3")
                             .scheduler("intrawave")
                             .epilogue("cshuffle"),
                         "gfx942"));

// =============================================================================
// MAIN
// =============================================================================

int main(int argc, char* argv[])
{
    ExampleArgs args("Example 05: JSON Export", "Export registry information to JSON format");
    args.add_option("--output", "registry.json", "Output JSON file path");
    args.add_option("--arch", "gfx942", "GPU architecture");
    args.add_flag("--list", "List all kernel sets");

    if(!args.parse(argc, argv))
        return 0;

    print_header("Example 05: JSON Export");

    std::string gfx_arch = args.get("--arch", "gfx942");

    if(args.has("--list"))
    {
        std::cout << "\nDeclared Kernel Sets:\n";
        KernelSetRegistry::instance().print();
        return 0;
    }

    std::string output_file = args.get("--output", "registry.json");

    // =========================================================================
    // Setup Registry
    // =========================================================================
    std::cout << "\nSetting up registry...\n";
    Registry registry;
    registry.set_name("json_export_registry");

    REGISTER_GENERATED_KERNELS(registry, gfx_arch);

    std::cout << "  Registry: " << registry.get_name() << "\n";
    std::cout << "  Kernels:  " << registry.size() << "\n";

    // =========================================================================
    // Export to JSON
    // =========================================================================
    std::cout << "\nExporting to JSON...\n";

    std::string json = registry.export_json(true);

    std::cout << "\nJSON Preview (first 500 chars):\n";
    print_separator();
    std::cout << json.substr(0, std::min(size_t(500), json.size()));
    if(json.size() > 500)
        std::cout << "\n...";
    std::cout << "\n";
    print_separator();

    // Write to file
    std::ofstream file(output_file);
    if(file.is_open())
    {
        file << json;
        file.close();
        std::cout << "\nExported to: " << output_file << "\n";
        std::cout << "File size: " << json.size() << " bytes\n";
    }
    else
    {
        std::cerr << "Failed to write to: " << output_file << "\n";
        return 1;
    }

    // =========================================================================
    // Also show kernel set declarations
    // =========================================================================
    std::cout << "\nKernel Set Declarations:\n";
    print_separator();
    KernelSetRegistry::instance().print();
    print_separator();

    return 0;
}