Update debug_to_file API to remove type_code (#8183)

* Add .npy support to halide_image_io

The .npy format is NumPy's native format for storing multidimensional arrays (aka tensors/buffers). Being able to load/save in this format makes it (potentially) a lot easier to interchange data with the Python ecosystem, as well as providing a file format that support floating-point data more robustly than any of the others that we current support.

This adds load/save support for a useful subset:
- We support the int/uint/float types common in Halide (except for f16/bf16 for now)
- We don't support reading or writing files that are in `fortran_order`
- We don't support any object/struct/etc files, only numeric primitives
- We only support loading files that are in the host's endianness (typically little-endian)

Note that at present this doesn't support f16 / bf16 formats, but that could likely be added with minimal difficulty.

The tricky bit of this is that the reading code has to parse a (limited) Python dict in text form. Please review that part carefully.

TODO: we could probably add this as an option for `debug_to_file()` without too much pain in a followup PR.

* clang-tidy

* clang-tidy

* Address review comments

* Allow for "keys" as well as 'keys'

* Add .npy support to debug_to_file()

Built on top of #8175, this adds .npy as an option. This is actually pretty great because it's easy to do something like

```
ss = numpy.load("my_file.npy")
print(ss)
```

in Python and get nicely-formatted output, which can sometimes be a lot easier for debugging that inserting lots of print() statements (see #8176)

Did a drive-by change to the correctness test to use this format instead of .mat.

* Add float16 support

* Add support for Float16 images in npy

* Assume little-endian

* Remove redundant halide_error()

* naming convention

* naming convention

* Test both mat and npy

* Don't call halide_error()

* Use old-school parser

* clang-tidy

* Update debug_to_file API to remove type_code

* Clean up into single table

* Update CodeGen_LLVM.cpp

* Fix tmp codes

* Update InjectHostDevBufferCopies.cpp

* Update InjectHostDevBufferCopies.cpp

* trigger buildbots

src/CodeGen_LLVM.cpp

@@ -2637,7 +2637,7 @@ void CodeGen_LLVM::visit(const Call *op) {
     // handled in the standard library, but ones with e.g. varying
     // types are handled here.
     if (op->is_intrinsic(Call::debug_to_file)) {
-        internal_assert(op->args.size() == 3);
+        internal_assert(op->args.size() == 2);
         const StringImm *filename = op->args[0].as<StringImm>();
         internal_assert(filename) << "Malformed debug_to_file node\n";
         // Grab the function from the initial module
@@ -2647,10 +2647,10 @@ void CodeGen_LLVM::visit(const Call *op) {
         // Make the filename a global string constant
         Value *user_context = get_user_context();
         Value *char_ptr = codegen(Expr(filename));
-        vector<Value *> args = {user_context, char_ptr, codegen(op->args[1])};
+        vector<Value *> args = {user_context, char_ptr};
 
-        Value *buffer = codegen(op->args[2]);
-        buffer = builder->CreatePointerCast(buffer, debug_to_file->getFunctionType()->getParamType(3));
+        Value *buffer = codegen(op->args[1]);
+        buffer = builder->CreatePointerCast(buffer, debug_to_file->getFunctionType()->getParamType(2));
         args.push_back(buffer);
 
         value = builder->CreateCall(debug_to_file, args);

src/DebugToFile.cpp

@@ -42,37 +42,6 @@ class DebugToFile : public IRMutator {
                 num_elements *= bound.extent;
             }
 
-            // TODO: why do we bother with this? halide_debug_to_file()
-            // can infer the type-and-size it needs from the buffer's type field.
-            int type_code = 0;
-            Type t = op->types[0];
-            if (t == Float(32)) {
-                type_code = 0;
-            } else if (t == Float(64)) {
-                type_code = 1;
-            } else if (t == UInt(8) || t == UInt(1)) {
-                type_code = 2;
-            } else if (t == Int(8)) {
-                type_code = 3;
-            } else if (t == UInt(16)) {
-                type_code = 4;
-            } else if (t == Int(16)) {
-                type_code = 5;
-            } else if (t == UInt(32)) {
-                type_code = 6;
-            } else if (t == Int(32)) {
-                type_code = 7;
-            } else if (t == UInt(64)) {
-                type_code = 8;
-            } else if (t == Int(64)) {
-                type_code = 9;
-            } else if (t == Float(16)) {
-                type_code = 10;
-            } else {
-                user_error << "Type " << t << " not supported for debug_to_file\n";
-            }
-            args.emplace_back(type_code);
-
             Expr buf = Variable::make(Handle(), f.name() + ".buffer");
             args.push_back(buf);
 

src/InjectHostDevBufferCopies.cpp

@@ -76,8 +76,8 @@ class FindBufferUsage : public IRVisitor {
                 op->args[i].accept(this);
             }
         } else if (op->is_intrinsic(Call::debug_to_file)) {
-            internal_assert(op->args.size() == 3);
-            if (is_buffer_var(op->args[2])) {
+            internal_assert(op->args.size() == 2);
+            if (is_buffer_var(op->args[1])) {
                 devices_touched.insert(current_device_api);
                 devices_writing.insert(current_device_api);
             }

src/runtime/HalideRuntime.h

@@ -463,7 +463,6 @@ extern halide_get_library_symbol_t halide_set_custom_get_library_symbol(halide_g
  * Cannot be replaced in JITted code at present.
  */
 extern int32_t halide_debug_to_file(void *user_context, const char *filename,
-                                    int32_t type_code,
                                     struct halide_buffer_t *buf);
 
 /** Types in the halide type system. They can be ints, unsigned ints,

src/runtime/write_debug_image.cpp

@@ -26,26 +26,6 @@ namespace Halide {
 namespace Runtime {
 namespace Internal {
 
-// Mappings from the type_code passed in to the type codes of the
-// formats. See "type_code" in DebugToFile.cpp
-
-constexpr int kNumTypeCodes = 11;
-
-// TIFF sample type values are:
-//     1 => Unsigned int
-//     2 => Signed int
-//     3 => Floating-point
-WEAK int16_t pixel_type_to_tiff_sample_type[kNumTypeCodes] = {
-    // float, double, uint8, int8, ... uint64, int64
-    3, 3, 1, 2, 1, 2, 1, 2, 1, 2, 0};
-
-// See the .mat level 5 documentation for matlab class codes.
-WEAK uint8_t pixel_type_to_matlab_class_code[kNumTypeCodes] = {
-    7, 6, 9, 8, 11, 10, 13, 12, 15, 14, 0};
-
-WEAK uint8_t pixel_type_to_matlab_type_code[kNumTypeCodes] = {
-    7, 9, 2, 1, 4, 3, 6, 5, 13, 12, 0};
-
 #pragma pack(push)
 #pragma pack(2)
 
@@ -138,37 +118,56 @@ constexpr char big_endian_char = '>';
 constexpr char no_endian_char = '|';
 constexpr char host_endian_char = (host_is_big_endian ? big_endian_char : little_endian_char);
 
-struct npy_dtype_info_t {
+struct npy_type_info_t {
     char byte_order;
     char kind;
-    size_t item_size;
+    uint16_t item_size;
 };
 
-struct htype_to_dtype {
+struct mat_type_info_t {
+    uint8_t class_code, type_code;
+};
+
+struct tmp_type_info_t {
+    int8_t type_code;
+};
+
+struct tiff_type_info_t {
+    int8_t type_code;
+};
+
+struct halide_type_to_dst_type_t {
     halide_type_t htype;
-    npy_dtype_info_t dtype;
+    npy_type_info_t npy;
+    mat_type_info_t mat;
+    tmp_type_info_t tmp;
+    tiff_type_info_t tiff;
 };
 
-WEAK htype_to_dtype npy_dtypes[] = {
-    {halide_type_t(halide_type_float, 16), {host_endian_char, 'f', 2}},
-    {halide_type_of<float>(), {host_endian_char, 'f', sizeof(float)}},
-    {halide_type_of<double>(), {host_endian_char, 'f', sizeof(double)}},
-    {halide_type_of<int8_t>(), {no_endian_char, 'i', sizeof(int8_t)}},
-    {halide_type_of<int16_t>(), {host_endian_char, 'i', sizeof(int16_t)}},
-    {halide_type_of<int32_t>(), {host_endian_char, 'i', sizeof(int32_t)}},
-    {halide_type_of<int64_t>(), {host_endian_char, 'i', sizeof(int64_t)}},
-    {halide_type_of<uint8_t>(), {no_endian_char, 'u', sizeof(uint8_t)}},
-    {halide_type_of<uint16_t>(), {host_endian_char, 'u', sizeof(uint16_t)}},
-    {halide_type_of<uint32_t>(), {host_endian_char, 'u', sizeof(uint32_t)}},
-    {halide_type_of<uint64_t>(), {host_endian_char, 'u', sizeof(uint64_t)}},
+// See the .mat level 5 documentation for matlab class codes.
+
+// clang-format off
+WEAK halide_type_to_dst_type_t debug_to_file_type_map[] = {  //                        mat       tmp   tiff
+    { halide_type_of<float>(),              {host_endian_char, 'f', sizeof(float)},    {7, 7},   {0},  {3} },
+    { halide_type_of<double>(),             {host_endian_char, 'f', sizeof(double)},   {6, 9},   {1},  {3} },
+    { halide_type_of<uint8_t>(),            {no_endian_char,   'u', sizeof(uint8_t)},  {9, 2},   {2},  {1} },
+    { halide_type_of<bool>(),               {no_endian_char,   'u', sizeof(uint8_t)},  {9, 2},   {2},  {1} },
+    { halide_type_of<int8_t>(),             {no_endian_char,   'i', sizeof(int8_t)},   {8, 1},   {3},  {2} },
+    { halide_type_of<uint16_t>(),           {host_endian_char, 'u', sizeof(uint16_t)}, {11, 4},  {4},  {1} },
+    { halide_type_of<int16_t>(),            {host_endian_char, 'i', sizeof(int16_t)},  {10, 3},  {5},  {2} },
+    { halide_type_of<uint32_t>(),           {host_endian_char, 'u', sizeof(uint32_t)}, {13, 6},  {6},  {1} },
+    { halide_type_of<int32_t>(),            {host_endian_char, 'i', sizeof(int32_t)},  {12, 5},  {7},  {2} },
+    { halide_type_of<uint64_t>(),           {host_endian_char, 'u', sizeof(uint64_t)}, {15, 13}, {8},  {1} },
+    { halide_type_of<int64_t>(),            {host_endian_char, 'i', sizeof(int64_t)},  {14, 12}, {9},  {2} },
+    { halide_type_t(halide_type_float, 16), {host_endian_char, 'f', 2},                {0, 0},   {-1}, {0} },
 };
+// clang-format on
 
 }  // namespace Internal
 }  // namespace Runtime
 }  // namespace Halide
 
-WEAK extern "C" int halide_debug_to_file(void *user_context, const char *filename,
-                                         int32_t type_code, struct halide_buffer_t *buf) {
+WEAK extern "C" int halide_debug_to_file(void *user_context, const char *filename, struct halide_buffer_t *buf) {
 
     if (buf->is_bounds_query()) {
         halide_error(user_context, "Bounds query buffer passed to halide_debug_to_file");
@@ -209,15 +208,19 @@ WEAK extern "C" int halide_debug_to_file(void *user_context, const char *filenam
 
     uint32_t final_padding_bytes = 0;
 
-    if (ends_with(filename, ".npy")) {
-        npy_dtype_info_t di = {0, 0, 0};
-        for (const auto &d : npy_dtypes) {
-            if (d.htype == buf->type) {
-                di = d.dtype;
-                break;
-            }
+    const halide_type_to_dst_type_t *type_found = nullptr;
+    for (const auto &d : debug_to_file_type_map) {
+        if (d.htype == buf->type) {
+            type_found = &d;
+            break;
         }
-        if (di.byte_order == 0) {
+    }
+    if (!type_found) {
+        return halide_error_code_debug_to_file_failed;
+    }
+
+    if (ends_with(filename, ".npy")) {
+        if (type_found->npy.item_size == 0) {
             return halide_error_code_debug_to_file_failed;
         }
 
@@ -227,9 +230,9 @@ WEAK extern "C" int halide_debug_to_file(void *user_context, const char *filenam
         char *end = dict_string_buf + max_dict_string_size - 1;
 
         dst = halide_string_to_string(dst, end, "{'descr': '");
-        *dst++ = di.byte_order;
-        *dst++ = di.kind;
-        dst = halide_int64_to_string(dst, end, di.item_size, 1);
+        *dst++ = type_found->npy.byte_order;
+        *dst++ = type_found->npy.kind;
+        dst = halide_int64_to_string(dst, end, type_found->npy.item_size, 1);
         dst = halide_string_to_string(dst, end, "', 'fortran_order': False, 'shape': (");
         for (int d = 0; d < buf->dimensions; ++d) {
             if (d > 0) {
@@ -272,7 +275,7 @@ WEAK extern "C" int halide_debug_to_file(void *user_context, const char *filenam
             return halide_error_code_debug_to_file_failed;
         }
     } else if (ends_with(filename, ".tiff") || ends_with(filename, ".tif")) {
-        if (type_code == 10) {
+        if (type_found->tiff.type_code == 0) {
             return halide_error_code_debug_to_file_failed;
         }
 
@@ -320,9 +323,8 @@ WEAK extern "C" int halide_debug_to_file(void *user_context, const char *filenam
                         __builtin_offsetof(halide_tiff_header, height_resolution));  // Height resolution
         tag++->assign16(284, 1, 2);                                                  // Planar configuration -- planar
         tag++->assign16(296, 1, 1);                                                  // Resolution Unit -- none
-        tag++->assign16(339, 1,
-                        pixel_type_to_tiff_sample_type[type_code]);  // Sample type
-        tag++->assign32(32997, 1, depth);                            // Image depth
+        tag++->assign16(339, 1, type_found->tiff.type_code);                         // Sample type
+        tag++->assign32(32997, 1, depth);                                            // Image depth
 
         header.ifd0_end = 0;
         header.width_resolution[0] = 1;
@@ -351,7 +353,7 @@ WEAK extern "C" int halide_debug_to_file(void *user_context, const char *filenam
             }
         }
     } else if (ends_with(filename, ".mat")) {
-        if (type_code == 10) {
+        if (type_found->mat.type_code == 0) {
             return halide_error_code_debug_to_file_failed;
         }
 
@@ -406,7 +408,7 @@ WEAK extern "C" int halide_debug_to_file(void *user_context, const char *filenam
             // This is a matrix
             14, 40 + padded_dimensions * 4 + padded_name_size + (uint32_t)payload_bytes + final_padding_bytes,
             // The element type
-            6, 8, pixel_type_to_matlab_class_code[type_code], 1,
+            6, 8, type_found->mat.class_code, 1,
             // The shape
             5, (uint32_t)(dims * 4)};
 
@@ -430,21 +432,20 @@ WEAK extern "C" int halide_debug_to_file(void *user_context, const char *filenam
         }
 
         // Payload header
-        uint32_t payload_header[2] = {
-            pixel_type_to_matlab_type_code[type_code], (uint32_t)payload_bytes};
+        uint32_t payload_header[2] = {type_found->mat.type_code, (uint32_t)payload_bytes};
         if (!f.write(payload_header, sizeof(payload_header))) {
             return halide_error_code_debug_to_file_failed;
         }
     } else {
-        if (type_code == 10) {
+        if (type_found->tmp.type_code < 0) {
             return halide_error_code_debug_to_file_failed;
         }
 
         int32_t header[] = {shape[0].extent,
                             shape[1].extent,
                             shape[2].extent,
                             shape[3].extent,
-                            type_code};
+                            type_found->tmp.type_code};
         if (!f.write((void *)(&header[0]), sizeof(header))) {
             return halide_error_code_debug_to_file_failed;
         }