Revision a0aa20afbc521b6215e4c34deb8070b71a8aa386 authored by Andrew Adams on 14 March 2024, 01:00:04 UTC, committed by Andrew Adams on 14 March 2024, 01:00:04 UTC
1 parent 83616f2
HalideTraceDump.cpp
#include "HalideBuffer.h"
#include "HalideTraceUtils.h"
#include "halide_image_io.h"
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fcntl.h>
#include <map>
#include <string>
#include <vector>
/** \file
*
* A tool which can read a binary Halide trace file, and dump files
* containing the final pixel values recorded for each traced Func.
*
* Currently dumps into supported Halide image formats.
*/
using namespace Halide;
using namespace Internal;
using Halide::Runtime::Buffer;
using std::map;
using std::string;
using std::vector;
struct BufferOutputOpts {
enum OutputType {
PNG = 0,
JPG,
PGM,
TMP,
MAT
};
enum OutputType type;
};
struct FuncInfo {
int min_coords[16];
int max_coords[16];
int dimensions;
halide_type_t type;
Buffer<> values;
FuncInfo() = default;
FuncInfo(Packet *p) {
int real_dims = p->dimensions / p->type.lanes;
if (real_dims > 16) {
fprintf(stderr, "Error: found trace packet with dimensionality > 16. Aborting.\n");
exit(1);
}
for (int i = 0; i < real_dims; i++) {
min_coords[i] = INT32_MAX;
max_coords[i] = INT32_MIN;
}
dimensions = real_dims;
type = p->type;
type.lanes = 1;
}
void add_preprocess(Packet *p) {
int real_dims = p->dimensions / p->type.lanes;
int lanes = p->type.lanes;
halide_type_t scalar_type = p->type;
scalar_type.lanes = 1;
if (scalar_type != type) {
fprintf(stderr, "Error: packet type doesn't match previous packets of same Func. Aborting.\n");
exit(1);
}
if (real_dims != dimensions) {
fprintf(stderr, "Error: packet dimensionality doesn't match previous packets of same Func. Aborting.\n");
exit(1);
}
for (int lane = 0; lane < lanes; lane++) {
for (int i = 0; i < real_dims; i++) {
if (p->coordinates()[lanes * i + lane] < min_coords[i]) {
min_coords[i] = p->coordinates()[lanes * i + lane];
}
if (p->coordinates()[lanes * i + lane] > max_coords[i]) {
max_coords[i] = p->coordinates()[lanes * i + lane];
}
}
}
}
void allocate() {
std::vector<int> extents;
for (int i = 0; i < dimensions; i++) {
extents.push_back(max_coords[i] - min_coords[i] + 1);
}
values = Buffer<>(type, extents);
if (values.data() == nullptr) {
fprintf(stderr, "Memory allocation failure. Aborting.\n");
exit(1);
}
}
void add(Packet *p) {
halide_type_t scalar_type = p->type;
scalar_type.lanes = 1;
if (scalar_type == halide_type_of<float>()) {
add_typed<float>(p);
} else if (scalar_type == halide_type_of<double>()) {
add_typed<double>(p);
} else if (scalar_type == halide_type_of<uint8_t>()) {
add_typed<uint8_t>(p);
} else if (scalar_type == halide_type_of<uint16_t>()) {
add_typed<uint16_t>(p);
} else if (scalar_type == halide_type_of<uint32_t>()) {
add_typed<uint32_t>(p);
} else if (scalar_type == halide_type_of<uint64_t>()) {
add_typed<uint64_t>(p);
} else if (scalar_type == halide_type_of<int8_t>()) {
add_typed<int8_t>(p);
} else if (scalar_type == halide_type_of<int16_t>()) {
add_typed<int16_t>(p);
} else if (scalar_type == halide_type_of<int32_t>()) {
add_typed<int32_t>(p);
} else if (scalar_type == halide_type_of<int64_t>()) {
add_typed<int64_t>(p);
} else if (scalar_type == halide_type_of<bool>()) {
add_typed<bool>(p);
} else {
printf("Packet with unknown type\n");
exit(1);
}
}
template<typename T>
void add_typed(Packet *p) {
Buffer<T> &buf = values.as<T>();
int lanes = p->type.lanes;
if (!allocated()) {
fprintf(stderr, "Packet storage not allocated. Aborting.\n");
exit(1);
}
if (p->dimensions < 0) {
fprintf(stderr, "Negative dimensionality found. Aborting.\n");
exit(1);
}
for (int lane = 0; lane < lanes; lane++) {
int coord[16];
for (int i = 0; i < dimensions; i++) {
coord[i] = p->coordinates()[lanes * i + lane] - min_coords[i];
}
buf(coord) = p->get_value_as<T>(lane);
}
}
bool allocated() const {
return (values.data() != nullptr);
}
};
bool check_and_continue(bool condition, const char *msg) {
if (!condition) {
fprintf(stderr, "Failed to dump func: %s\n", msg);
}
return condition;
}
void dump_func(string name, FuncInfo &func, BufferOutputOpts output_opts) {
// Remove special characters
for (char &c : name) {
if (!std::isalnum(c)) {
c = '_';
}
}
string filename;
switch (output_opts.type) {
case BufferOutputOpts::PNG:
filename = name + ".png";
break;
case BufferOutputOpts::JPG:
filename = name + ".jpg";
break;
case BufferOutputOpts::PGM:
filename = name + ".pgm";
break;
case BufferOutputOpts::TMP:
filename = name + ".tmp";
break;
case BufferOutputOpts::MAT:
filename = name + ".mat";
break;
default:
exit(1);
}
printf("[INFO] Dumping func '%s' to file: %s\n", name.c_str(), filename.c_str());
// Rely on save_image to do type-checking
Halide::Tools::convert_and_save_image<Buffer<>, check_and_continue>(func.values, filename);
}
void finish_dump(map<string, FuncInfo> &func_info, BufferOutputOpts output_opts) {
printf("\nTrace stats:\n");
printf(" Funcs:\n");
for (auto &pair : func_info) {
const string &name = pair.first;
FuncInfo &info = pair.second;
printf(" %s:\n", name.c_str());
// Type
printf(" Type: ");
if (info.type.code == halide_type_code_t::halide_type_int) {
printf("int%d\n", info.type.bits);
} else if (info.type.code == halide_type_code_t::halide_type_uint) {
printf("uint%d\n", info.type.bits);
} else if (info.type.code == halide_type_code_t::halide_type_float) {
printf("float%d\n", info.type.bits);
} else {
fprintf(stderr, "Unsupported Func type. Aborting.\n");
exit(1);
}
// Dimensions
printf(" Dimensions: %d\n", info.dimensions);
// Size of the func
printf(" Size: ");
for (int idx = 0; idx < info.dimensions; idx++) {
if (idx > 0) {
printf("x");
}
printf("%d", (info.max_coords[idx] - info.min_coords[idx]) + 1);
}
printf("\n");
// Minima
printf(" Minimum stored to in each dim: {");
for (int idx = 0; idx < info.dimensions; idx++) {
if (idx > 0) {
printf(", ");
}
printf("%d", info.min_coords[idx]);
}
printf("}\n");
// Maxima
printf(" Maximum stored to in each dim: {");
for (int idx = 0; idx < info.dimensions; idx++) {
if (idx > 0) {
printf(", ");
}
printf("%d", info.max_coords[idx]);
}
printf("}\n");
}
for (auto &pair : func_info) {
string name = pair.first;
FuncInfo &info = pair.second;
dump_func(name, info, output_opts);
}
printf("Done.\n");
}
void usage(char *const *argv) {
const string usage =
"Usage: " + string(argv[0]) +
" -i trace_file -t {png,jpg,pgm,tmp,mat}\n"
"\n"
"This tool reads a binary trace produced by Halide, and dumps all\n"
"Funcs into individual image files in the current directory.\n"
"To generate a suitable binary trace, use Func::trace_stores(), or the\n"
"target features trace_stores and trace_realizations, and run with\n"
"HL_TRACE_FILE=<filename>.\n";
fprintf(stderr, "%s\n", usage.c_str());
exit(1);
}
int main(int argc, char *const *argv) {
char *buf_filename = nullptr;
char *buf_imagetype = nullptr;
BufferOutputOpts outputopts;
for (int i = 1; i < argc - 1; i++) {
string arg = argv[i];
if (arg == "-t") {
i++;
buf_imagetype = argv[i];
} else if (arg == "-i") {
i++;
buf_filename = argv[i];
}
}
if (buf_filename == nullptr) {
usage(argv);
}
if (buf_imagetype == nullptr) {
usage(argv);
}
string imagetype(buf_imagetype);
if (imagetype == "jpg") {
outputopts.type = BufferOutputOpts::JPG;
} else if (imagetype == "png") {
outputopts.type = BufferOutputOpts::PNG;
} else if (imagetype == "pgm") {
outputopts.type = BufferOutputOpts::PGM;
} else if (imagetype == "tmp") {
outputopts.type = BufferOutputOpts::TMP;
} else if (imagetype == "mat") {
outputopts.type = BufferOutputOpts::MAT;
} else {
usage(argv);
}
FILE *file_desc = fopen(buf_filename, "r");
if (file_desc == nullptr) {
fprintf(stderr, "[Error opening file: %s. Exiting.\n", argv[1]);
exit(1);
}
printf("[INFO] Starting parse of binary trace...\n");
int packet_count = 0;
map<string, FuncInfo> func_info;
printf("[INFO] First pass...\n");
for (;;) {
Packet p;
if (!p.read_from_filedesc(file_desc)) {
printf("[INFO] Finished pass 1 after %d packets.\n", packet_count);
break;
}
// Packet read was successful.
packet_count++;
if ((packet_count % 100000) == 0) {
printf("[INFO] Pass 1: Read %d packets so far.\n", packet_count);
}
// Check if this was a store packet.
if ((p.event == halide_trace_store) || (p.event == halide_trace_load)) {
if (func_info.find(string(p.func())) == func_info.end()) {
printf("[INFO] Found Func with tracked accesses: %s\n", p.func());
func_info[string(p.func())] = FuncInfo(&p);
}
func_info[string(p.func())].add_preprocess(&p);
}
}
packet_count = 0;
fseek(file_desc, 0, SEEK_SET);
if (ferror(file_desc)) {
fprintf(stderr, "Error: couldn't seek back to beginning of trace file. Aborting.\n");
exit(1);
}
for (auto &pair : func_info) {
pair.second.allocate();
}
for (;;) {
Packet p;
if (!p.read_from_filedesc(file_desc)) {
printf("[INFO] Finished pass 2 after %d packets.\n", packet_count);
if (file_desc != nullptr) {
fclose(file_desc);
}
finish_dump(func_info, outputopts);
exit(0);
}
// Packet read was successful.
packet_count++;
if ((packet_count % 100000) == 0) {
printf("[INFO] Pass 2: Read %d packets so far.\n", packet_count);
}
// Check if this was a store packet.
if ((p.event == halide_trace_store) || (p.event == halide_trace_load)) {
if (func_info.find(string(p.func())) == func_info.end()) {
fprintf(stderr, "Unable to find Func on 2nd pass. Aborting.\n");
exit(1);
}
func_info[string(p.func())].add(&p);
}
}
}
![swh spinner](/static/img/swh-spinner.gif)
Computing file changes ...