Preserve input memory location / dtype for NN Descent

[jinsolp]

Closes #1901

Previous Code

• We almost always allocate device side fp16 arrays. This was for...
  • allowing wmma usage
  • allowing data modification for CosineExpanded preprocessing

Current PR Changes

• No logical changes apart from removing dispatching fp32 input to use fp32 vs fp16 distance computation. This is removed now and will default to using the input type (e.g. keep fp32 as fp32). One exception is when compress_to_fp16=True and input type is fp32. In this case we conver to fp16 to exploit wmma.
• Reducing redundant memory:
  • We only allocate device side arrays corresponding to input dtype if input is not device-accessible (allocate half types for fp32 if compress_to_fp16=True).
  • Remove preprocessing for CosineExpanded metric (because we don't want to allocate additional device side data arrays) and do the computation inside the calculate_metric function.

Peak memory usage Changes

• food data (5M x 384) = 7.25GiB

• sports data (13M x 284) = 18.55GiB

• notice how for FP32->FP16 Device (meaning data is already on device), previous code allocates a new fp16 array, resulting in more gpu mem usage. This PR ensures that we convert to fp16 on-th-fly (resulting in the overhead in time) instead of allocating new fp16 memory for that.

Performance Changes

• Conversion Overhead: On-the-fly conversion introduces negligible overhead.
• Cosine Metric: Now reads l2 norms inside the calculate_metric function, aligning with access pattern used by the L2 distance metric. Adds minimal overhead (e.g. previously 18.2937s VS 18.7598s for 5Mx384 data)

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.