Remove Tensor.h includes from spdiags cpu kernel (#84500)

This file uses `Tensor::operator[]` in the middle of a `cpu_kernel`
which is not allowed because it relies on the thread-local dispatcher
state. Instead, we should just do the stride calculations.
Pull Request resolved: #84500
Approved by: https://github.com/ezyang

Author: peterbell10

aten/src/ATen/native/cpu/SparseFactories.cpp

@@ -1,35 +1,25 @@
+#define TORCH_ASSERT_NO_OPERATORS
+#include <ATen/native/sparse/SparseFactories.h>
+
 #include <ATen/Dispatch.h>
-#include <ATen/SparseTensorImpl.h>
-#include <ATen/SparseTensorUtils.h>
-#include <ATen/TensorIndexing.h>
 #include <ATen/TensorIterator.h>
-#include <ATen/core/ATen_fwd.h>
-#include <ATen/core/Tensor.h>
+#include <ATen/core/TensorBase.h>
 #include <ATen/native/cpu/Loops.h>
-#include <ATen/native/sparse/SparseFactories.h>
-#include <c10/core/Scalar.h>
-#include <c10/util/ArrayRef.h>
+#include <c10/core/ScalarType.h>
 #include <c10/util/Exception.h>
 
-#ifndef AT_PER_OPERATOR_HEADERS
-#include <ATen/Functions.h>
-#include <ATen/NativeFunctions.h>
-#else
-#include <ATen/ops/sparse_coo_tensor.h>
-#endif
-
 namespace at {
 namespace native {
-using namespace at::sparse;
 
 namespace {
 void _spdiags_kernel_cpu(
     TensorIterator& iter,
-    const Tensor& diagonals,
-    Tensor& values,
-    Tensor& indices) {
-  auto* row_index_write_ptr = indices[0].data_ptr<int64_t>();
-  auto* col_index_write_ptr = indices[1].data_ptr<int64_t>();
+    const TensorBase& diagonals,
+    TensorBase& values,
+    TensorBase& indices) {
+  auto* row_index_write_ptr = indices.data_ptr<int64_t>();
+  auto* col_index_write_ptr = row_index_write_ptr + indices.stride(0);
+  const int64_t diagonals_index_stride = diagonals.stride(0);
   const int64_t diagonals_read_stride = diagonals.stride(1);
   AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND4(
       at::ScalarType::BFloat16,
@@ -39,7 +29,9 @@ void _spdiags_kernel_cpu(
       diagonals.scalar_type(),
       "spdiags_cpu",
       [&] {
-        auto* values_write_ptr = values.data_ptr<scalar_t>();
+        auto* const values_write_ptr = values.data_ptr<scalar_t>();
+        const auto* const diagonals_ptr = diagonals.data_ptr<scalar_t>();
+
         cpu_kernel(
             iter,
             [&](int64_t diag_index,
@@ -52,8 +44,9 @@ void _spdiags_kernel_cpu(
                 auto* vals_start = values_write_ptr + out_offset;
                 const int64_t first_col = std::max<int64_t>(diag_offset, 0);
                 const int64_t first_row = first_col - diag_offset;
-                auto* data_read = diagonals[diag_index].data_ptr<scalar_t>() +
-                    first_col * diagonals_read_stride;
+                auto* data_read = (diagonals_ptr +
+                                   diagonals_index_stride * diag_index +
+                                   first_col * diagonals_read_stride);
                 for (int64_t i = 0; i < n_out; ++i) {
                   rows_start[i] = first_row + i;
                   cols_start[i] = first_col + i;

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

@@ -1,4 +1,5 @@
 #include <ATen/Dispatch.h>
+#include <ATen/TensorIterator.h>
 #include <ATen/native/sparse/SparseFactories.h>
 
 #ifndef AT_PER_OPERATOR_HEADERS

aten/src/ATen/native/sparse/SparseFactories.h

@@ -1,14 +1,14 @@
 #pragma once
-#include <ATen/TensorIterator.h>
-#include <ATen/core/ATen_fwd.h>
-#include <ATen/core/Tensor.h>
 #include <ATen/native/DispatchStub.h>
 
 namespace at {
+struct TensorIterator;
+class TensorBase;
+
 namespace native {
 
 using spdiags_kernel_fn_t =
-    void (*)(TensorIterator&, const Tensor&, Tensor&, Tensor&);
+    void (*)(TensorIterator&, const TensorBase&, TensorBase&, TensorBase&);
 
 DECLARE_DISPATCH(spdiags_kernel_fn_t, spdiags_kernel_stub);
 } // namespace native