Add .npy support to halide_image_io

[steven-johnson]

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. see #8177

[abadams]

It's not, and for_each_value below doesn't guarantee a planar order. I think you want write_planar_payload

[steven-johnson]

Yeah. So, to be clear, if we use write_planar_payload, we can definitely assert that fortran_order=false for things we write?

[abadams]

Yes, I believe so.

[abadams]

pos + k.size()? That would make the code below robust to keys with ':' in them.

[steven-johnson]

Would needs to be k.size() + 2 to account for the enclosing quotes... but then the enclosing logic later fails because it assumes that we stop at the following key. That said, the worst case is that if someone produces a file that has extra keys (or keys with unlikely chars), we just refuse to read the file.

[abadams]

So if there's an unusual key in there, this will be wrong. Does the spec guarantee that only those three keys exist for the versions we support?

[steven-johnson]

It's not clear; the spec says The dictionary contains three keys but doesn't say it contains only these.

[steven-johnson]

(BTW, I have verified that a .npy file produced by NumPy in Python on my Mac loads fine here... sample size = 1 :-)

[abadams]

I guess another way this can fail is if any of these strings appear in the key values somewhere. Seems impossible?

[steven-johnson]

A 'real' .npy file wouldn't do that (it would be malformed), but a malicious one could. Let me think about how I can write a more robust parser here. (Sigh... if only the spec had just required a specific layout that was easier to parse in non-Python)

[steven-johnson]

Added support for float16, since NumPy supports it just fine (which makes this the only format I know of that we support that can handle it OK)

[steven-johnson]

So I think the only real blocker here is how robust we need to be when parsing the Python dict; from a quick survey of other C/C++ parsing libraries I could find via google, most of are even more naive than this code; also, based on looking at the NumPy source code that generates canonical files, it never inserts anything weird into it, so it seems likely to assume that most files will be 'well-formed'. I'm fine with failing to read a non-well-formed file.... the real question is whether it's possible to malform a file in a way that can cause malicious exploits, of course.

[abadams]

The spec looks like we can expect the header to be very consistent. If we want to be rigid about it, how about writing a regex that the header must match:

   static std::regex r_header("^\\{'descr': '[<|]([ifu])(\\d+)', "
                               "'fortran_order': False, "
                               "'shape': \\(([\\d, ]+)\\), \\}$");
    static std::regex r_num("(\\d+)");
    std::smatch m;
    if (!std::regex_match(header, m, r_header)) {
        return false;
    }

    char type_code;
    int bytes;
    std::vector<int> extents;

    type_code = m[1].str()[0];
    bytes = std::stoi(m[2].str());
    std::string shape_string = m[3].str();
    for (auto it = std::sregex_token_iterator(shape_string.begin(), shape_string.end(), r_num);
         it != std::sregex_token_iterator(); it++) {
        std::cout << it->str() << "\n";
        extents.push_back(std::stoi(it->str()));
    }

    return true;

[steven-johnson]

The spec looks like we can expect the header to be very consistent

Are we reading the same spec? I didn't perceive that guarantee in the language (though I'd be happy if it was present)

[steven-johnson]

Sadly, the spec explicitly says the keys aren't required to be in order (wtf?)...

[abadams]

huh, the document I found describing it says that they are guaranteed to be in alphabetical order. I guess it's not really a spec, but it is the most precise description I found: https://paulbourke.net/dataformats/npy/#:~:text=The%20npy%20files%20contain%20a,string%20with%20datatype%20and%20dimensions.

[steven-johnson]

That link says:

For repeatability and readability, this dictionary is formatted using pprint.pformat() so the keys are in alphabetic order.

which is indeed what the NumPy source does, but in https://numpy.org/devdocs/reference/generated/numpy.lib.format.html they say

For repeatability and readability, the dictionary keys are sorted in alphabetic order. This is for convenience only. A writer SHOULD implement this if possible. A reader MUST NOT depend on this.

[steven-johnson]

That said: we could just flout the SHOULD part and say that we're gonna require them sorted; that would be safer and would probably eliminate ~zero real-world files (and if we find counterexamples we can address it then)

[abadams]

How about this old-school version:

struct NpyHeader {
    char type_code;
    int type_bytes;
    std::vector<int> extents;

    bool parse(const std::string &header) {
        const char *ptr = &header[0];
        if (*ptr++ != '{') {
            return false;
        }
        while (true) {
            char endian;
            int consumed;
            if (std::sscanf(ptr, "'descr': '%c%c%d'%n", &endian, &type_code, &type_bytes, &consumed) == 3) {
                if (endian != '<' && endian != '|') {
                    return false;
                }
                ptr += consumed;
            } else if (std::strncmp(ptr, "'fortran_order': False", 22) == 0) {
                ptr += 22;
            } else if (std::strncmp(ptr, "'shape': (", 10) == 0) {
                ptr += 10;
                int n;
                while (std::sscanf(ptr, "%d%n", &n, &consumed) == 1) {
                    extents.push_back(n);
                    ptr += consumed;
                    if (*ptr == ',') ptr++;
                    if (*ptr == ' ') ptr++;
                }
                if (*ptr++ != ')') {
                    return false;
                }
            } else if (*ptr == '}') {
                return true;
            } else {
                return false;
            }
            if (*ptr == ',') ptr++;
            if (*ptr == ' ') ptr++;
            assert(ptr <= &header.back());
        }
    }
};

[steven-johnson]

How about this old-school version:

trying now

[steven-johnson]

Looks good and deals with pathologies reasonably, PTAL