[core][sparse][pruning] Add (i8i8)-> fp16 support to cuSPARSELt matmul (#109214)
Summary:
This PR adds in support for sparse matmul using cuSPASRELt with int8
inputs and fp16 outputs.
It does so by adding a out_dtype flag to `torch_cslt_sparse_mm`.
Because the only mixed_dtype support present in cuSPARSELt is for int8
input and fp16 output, we error out if:
* out_dtype is set and the input tensors are not int8.
* out_dtype is set to any value other than fp16
Test Plan:
python test/test_sparse_semi_structured -k int8_in_fp16_out
Reviewers:
@cphursh
Subscribers:
Tasks:
Tags:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/109214
Approved by: https://github.com/cpuhrsch
