Make sparse empty constructors specialize instead of fail on symbolic inputs (#129983)

Exercised in pytorch/pytorch#128327

Signed-off-by: Edward Z. Yang <[email protected]>
Pull Request resolved: pytorch/pytorch#129983
Approved by: https://github.com/anijain2305

Author: ezyang

aten/src/ATen/native/native_functions.yaml

@@ -2370,8 +2370,10 @@
     MPS: empty_mps
     Meta: empty_meta_symint
     MkldnnCPU: empty_mkldnn
-    SparseCPU, SparseCUDA, SparseMeta: empty_sparse
-    SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: empty_sparse_compressed
+    SparseCPU, SparseCUDA: empty_sparse
+    SparseMeta: empty_sparse_symint
+    SparseCsrCPU, SparseCsrCUDA: empty_sparse_compressed
+    SparseCsrMeta: empty_sparse_compressed_symint
     QuantizedCPU, QuantizedCUDA, QuantizedMeta: empty_unknown_quantized
   tags: core
 

aten/src/ATen/native/sparse/SparseCsrTensor.cpp

@@ -644,6 +644,17 @@ SPARSE_COMPRESSED_TENSOR(csc, kSparseCsc)
 SPARSE_COMPRESSED_TENSOR(bsr, kSparseBsr)
 SPARSE_COMPRESSED_TENSOR(bsc, kSparseBsc)
 
+Tensor empty_sparse_compressed_symint(
+    SymIntArrayRef size,
+    std::optional<ScalarType> dtype,
+    std::optional<Layout> layout,
+    std::optional<Device> device,
+    std::optional<bool> pin_memory,
+    std::optional<MemoryFormat> optional_memory_format) {
+  // TODO: Don't specialize
+  return empty_sparse_compressed(C10_AS_INTARRAYREF_SLOW_ALLOC(size), dtype, layout, device, pin_memory, optional_memory_format);
+}
+
 // Warning: ideally, torch.empty(..., layout=<sparse compressed
 // format>) ought to be unsupported because it does not return a valid
 // sparse compressed tensor without initialization of compressed

aten/src/ATen/native/sparse/SparseTensor.cpp

@@ -226,6 +226,17 @@ SparseTensor new_with_dims_and_tensor_sparse_symint(
 /** Public creation API that dispatch to methods above **/
 
 /** Empty init **/
+Tensor empty_sparse_symint(
+    SymIntArrayRef size,
+    std::optional<ScalarType> dtype,
+    std::optional<Layout> layout,
+    std::optional<Device> device,
+    std::optional<bool> pin_memory,
+    std::optional<MemoryFormat> optional_memory_format) {
+  // TODO: Don't specialize
+  return empty_sparse(C10_AS_INTARRAYREF_SLOW_ALLOC(size), dtype, layout, device, pin_memory, optional_memory_format);
+}
+
 Tensor empty_sparse(
     IntArrayRef size,
     std::optional<ScalarType> dtype,

c10/core/SymIntArrayRef.h

@@ -2,8 +2,10 @@
 
 #include <c10/core/SymInt.h>
 #include <c10/util/ArrayRef.h>
+#include <c10/util/DimVector.h>
 #include <c10/util/Exception.h>
 #include <c10/util/Optional.h>
+#include <c10/util/irange.h>
 #include <cstdint>
 
 namespace c10 {
@@ -45,7 +47,22 @@ inline at::IntArrayRef asIntArrayRefSlow(
   return asIntArrayRefUnchecked(ar);
 }
 
+// Even slower than asIntArrayRefSlow, as it forces an allocation for a
+// destination int, BUT it is able to force specialization (it never errors)
+inline c10::DimVector asIntArrayRefSlowAlloc(
+    c10::SymIntArrayRef ar,
+    const char* file,
+    int64_t line) {
+  c10::DimVector res(ar.size(), 0);
+  for (const auto i : c10::irange(ar.size())) {
+    res[i] = ar[i].guard_int(file, line);
+  }
+  return res;
+}
+
 #define C10_AS_INTARRAYREF_SLOW(a) c10::asIntArrayRefSlow(a, __FILE__, __LINE__)
+#define C10_AS_INTARRAYREF_SLOW_ALLOC(a) \
+  c10::asIntArrayRefSlowAlloc(a, __FILE__, __LINE__)
 
 // Prefer using a more semantic constructor, like
 // fromIntArrayRefKnownNonNegative