I'm closing this because I don't think there is any way to satisfy the query UR_DEVICE_INFO_PARTITION_TYPE by returning a struct as it will inevitably include pointers and will be impossible to safely give data back to the caller especially for info about partitioning by counts.

struct ur_partition_info_t {
  ... 
  uint32_t *counts;
  size_t counts_len;
  ...
};

Even with such a struct the actual counts data is not copied to the caller before the entry-point returns.

@fabiomestre @kbenzie I think in this situation we have to stick with the OpenCL style array.

:(

As a simpler alternative, would it be possible to make the key-value pair an explicit struct with proper types so that the values aren't just all cast to int? Something like:

struct ur_partition_property_t {
  enum ur_device_partition_t partition_type;
  union {
      struct ur_device_partition_equally_desc_t equally;
      ...
  } value;
};

You can still use it in a simple array but would be easier to use with proper types in place.

Edit: Looking at the current docs, I've also noticed that the properties are in a "null-terminated array". This is redundant when retrieving the values because we already have the size, and could easily be changed in the urDevicePartition call to accept a length. Bounds checking would make the API safer to use.

That could be nicer @pbalcer.

I still think we can do better than the OpenCL approach, is the concearn about having pointers in structs to do with allocation lifetimes of the urDeviceGetInfo query @veselypeta?

I like your idea @pbalcer, we could perhaps do something like this:

 struct ur_partition_desc {
        ur_device_partition_t type;
        union {
            uint32_t equally; // Used to specify PARTITION_EQUALLY
            uint32_t count;   // Used to specify BY_COUNTS
            ur_device_affinity_domain_flags_t
                affinity_domain; // Used to specify AFFINITY_DOMAIN
        } value;
    };

    // Equally
    ur_partition_desc eq_desc;
    eq_desc.type = UR_DEVICE_PARTITION_EQUALLY;
    eq_desc.value.equally = 3u;

    urDevicePartition(..., eq_desc, /*count*/ 1, ...);


    // By Counts
    std::vector<ur_partition_desc> counts_desc;
    for(size_t i = 3; i < 8; ++i){
        ur_partition_desc desc;
        desc.type = UR_DEVICE_PARTITION_BY_COUNTS;
        desc.value.count = i; // some count
    }
    urDevicePartition(..., counts_desc.data(), /* count */ counts_desc.size());


    // then to query UR_DEVICE_INFO_PARTITION_TYPE

    size_t query_size;
    // maybe we return length of array instead of size in bytes
    urDeviceGetInfo(..., UR_DEVICE_INFO_PARTITION_TYPE, &query_size);  
    size_t len = query_size / sizeof(ur_partition_desc);
    std::vector<ur_partition_desc> query_desc(len);
    urDeviceGetInfo(..., query_desc.data, query_size, ...);

Ideally we would also enable pNext chains, I believe that should be simple enough by moving these array arguments into a struct which is built on the base structs.

struct ur_partition_desc {
        ur_device_partition_t type;
        union {
            uint32_t equally; // Used to specify PARTITION_EQUALLY
            uint32_t count;   // Used to specify BY_COUNTS
            ur_device_affinity_domain_flags_t
                affinity_domain; // Used to specify AFFINITY_DOMAIN
        } value;
    };

    // Equally
    ur_partition_desc eq_desc;
    eq_desc.type = UR_DEVICE_PARTITION_EQUALLY;
    eq_desc.value.equally = 3u;

    urDevicePartition(..., eq_desc, /*count*/ 1, ...);


    // By Counts
    std::vector<ur_partition_desc> counts_desc;
    for(size_t i = 3; i < 8; ++i){
        ur_partition_desc desc;
        desc.type = UR_DEVICE_PARTITION_BY_COUNTS;
        desc.value.count = i; // some count
    }
    urDevicePartition(..., counts_desc.data(), /* count */ counts_desc.size());


    // then to query UR_DEVICE_INFO_PARTITION_TYPE

    size_t query_size;
    // maybe we return length of array instead of size in bytes
    urDeviceGetInfo(..., UR_DEVICE_INFO_PARTITION_TYPE, &query_size);  
    size_t len = query_size / sizeof(ur_partition_desc);
    std::vector<ur_partition_desc> query_desc(len);
    urDeviceGetInfo(..., query_desc.data, query_size, ...);

I could see this solution working +1