Skip to content

Compute typeIdWithDefault and denseTypeIdWithDefault without reference to Backend #53544

@ezyang

Description

@ezyang

Currently they are implemented as

// NB: device_idx here is NOT a DeviceIndex, but index into PythonArgs
c10::DispatchKey typeIdWithDefault(PythonArgs& r, int64_t device_idx, c10::DispatchKey dispatch_key) {
  auto device_type = r.isNone(device_idx) ? computeDeviceType(dispatch_key) : r.device(device_idx).type();
  return backendToDispatchKey(backendToBackendOfDeviceType(dispatchKeyToBackend(dispatch_key), device_type));
}

// NB: device_idx here is NOT a DeviceIndex, but index into PythonArgs
c10::DispatchKey denseTypeIdWithDefault(PythonArgs& r, int64_t device_idx, c10::DispatchKey dispatch_key) {
  auto device_type = r.isNone(device_idx) ? computeDeviceType(dispatch_key) : r.device(device_idx).type();
  return backendToDispatchKey(toDense(backendToBackendOfDeviceType(dispatchKeyToBackend(dispatch_key), device_type)));
}     

This forces people who add new device types to also add a corresponding Backend for them. But the backend is just an incidental part of the computation, and we eventually turn it back to dispatch key. Do the calculation to dispatch key directly!

Metadata

Metadata

Assignees

Labels

module: bootcampWe plan to do a full writeup on the issue, and then get someone to do it for onboardingtriagedThis issue has been looked at a team member, and triaged and prioritized into an appropriate module

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions