[UR] Have urDeviceParition use desc for partitioning

[veselypeta]

Specify device partitioning scheme using a desc_t instead of OpenCL style raw array.

One thing I'm not quite sure of is I've chosen to create a ur_partition_desc_t which has no additional fields and will then point to a specific partitioning description. The alternative would be to type the parameter simply as ur_base_desc_t and then cast it based on stype i.e.:

ur_result_t urDevicePartition(ur_device_handle_t hDevice, const ur_base_desc_t *pDesc, ... ){
  switch(pDesc->stype){
    case UR_STRUCTURE_TYPE_DEVICE_PARTITION_EQUALLY_DESC:
      const auto *eqlDesc = static_cast<ur_device_partition_equally_desc_t *>(pDesc);
    case ...:
    default:
      return UR_RESULT_ERROR_INVALID_ENUMERATION;
  };
}

This way it would avoid the indirection - but the typename isn't really indicative as to what you should pass in I suppose. We could also just do something like typedef ur_device_partition_desc_t ur_base_desc_t?

---

Another thing to consider is how do we handle the query UR_DEVICE_INFO_PARTITION_TYPE: previously this would just return the ur_device_partition_property_t[] of properties passed into urDevicePartition, but what should be returned if a descriptor is used? The descriptors are allocated on the caller size and a pointer to the structure chain is passed. Do we copy the structure chain somewhere into the UR adapter and the query will just return a pointer to the start of the chain?

ur_device_partition_desc_t *desc = nullptr;
urDeviceGetInfo(device, UR_DEVICE_INFO_PARTITION_TYPE, sizeof(desc), &desc, nullptr);
// Use the pointer...

[veselypeta]

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.

[veselypeta]

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

[pbalcer]

:(

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.

[kbenzie]

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?

[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, ...);

[kbenzie]

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.

[fabiomestre]

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