Fix incomplete type conversions and address review feedback

jackulau · jackulau · commit a40897e698a4 · 2025-12-07T00:53:26.000-06:00
This commit addresses all critical bugs and review feedback from PR #1956: **Critical fixes (breaks stable ABI builds):** - Fixed 8 instances of `at::Half` → `apex_internal::Half` in type_shim.h - Fixed 4 instances of `at::BFloat16` → `apex_internal::BFloat16` in type_shim.h - Fixed 12 instances of `at::ScalarType::*` → `apex_internal::ScalarType::*` in nested switch statements - Fixed 4 instances of `AT_ERROR` → `APEX_ERROR` for consistency with dual-build pattern - Fixed 4 instances of `toString` → `apex_internal::toString` in error messages **CUDA stream handling (multi_tensor_apply.cuh):** - Implemented proper DeviceGuard using `torch::stable::accelerator::DeviceGuard` - Implemented proper stream retrieval using `aoti_torch_get_current_cuda_stream()` C API - Added `torch/csrc/inductor/aoti_torch/c/shim.h` include for stable ABI CUDA functions - This now properly preserves the current stream semantics like the traditional path **Documentation fixes:** - Fixed NCHW→NHWC comment error in stable_abi_utils.h:45 - Fixed NCDHW→NDHWC comment error in stable_abi_utils.h:64 **Completeness:** - Added MemoryFormat::Preserve case handling in is_contiguous() with explanatory comment These changes ensure the stable ABI infrastructure compiles correctly and addresses all feedback from maintainer review.
diff --git a/csrc/multi_tensor_apply.cuh b/csrc/multi_tensor_apply.cuh
@@ -1,6 +1,7 @@
 #ifdef TORCH_STABLE_ONLY
 #include <torch/csrc/stable/tensor.h>
 #include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
 #include <torch/headeronly/types.h>
 #include "stable_abi_utils.h"
 #else
@@ -91,9 +92,17 @@ void multi_tensor_apply(int64_t block_size, int64_t chunk_size, const apex_tenso
 #ifdef TORCH_STABLE_ONLY
   // Stable ABI: device guard and stream management
   auto device = tensor_lists[0][0].device();
-  // TODO: stable ABI device guard - for now assume correct device context
-  cudaStream_t stream = nullptr; // Use default stream for stable ABI
-  cudaGetLastError(); // Clear any prior errors
+  int32_t device_index = static_cast<int32_t>(device.index());
+
+  // Use stable ABI DeviceGuard for proper device context
+  torch::stable::accelerator::DeviceGuard device_guard(device_index);
+
+  // Get current CUDA stream using stable ABI C API
+  void* stream_ptr = nullptr;
+  auto err = aoti_torch_get_current_cuda_stream(device_index, &stream_ptr);
+  cudaStream_t stream = (err == AOTI_TORCH_SUCCESS)
+                        ? reinterpret_cast<cudaStream_t>(stream_ptr)
+                        : nullptr;
 #else
   const at::cuda::OptionalCUDAGuard device_guard(device_of(tensor_lists[0][0]));
   auto stream = at::cuda::getCurrentCUDAStream();
diff --git a/csrc/stable_abi_utils.h b/csrc/stable_abi_utils.h
@@ -42,7 +42,7 @@ inline bool is_contiguous(const torch::stable::Tensor& tensor, MemoryFormat form
   int64_t ndim = tensor.dim();
 
   if (format == MemoryFormat::ChannelsLast) {
-    // NCHW format requires ndim == 4
+    // NHWC format requires ndim == 4
     if (ndim != 4) return false;
 
     // For ChannelsLast (NHWC), strides should follow: C=1, W=C, H=W*W_size, N=H*H_size
@@ -61,7 +61,7 @@ inline bool is_contiguous(const torch::stable::Tensor& tensor, MemoryFormat form
   }
 
   if (format == MemoryFormat::ChannelsLast3d) {
-    // NCDHW format requires ndim == 5
+    // NDHWC format requires ndim == 5
     if (ndim != 5) return false;
 
     // For ChannelsLast3d (NDHWC), similar logic for 5D tensors
@@ -80,6 +80,11 @@ inline bool is_contiguous(const torch::stable::Tensor& tensor, MemoryFormat form
            (stride_n == D * H * W * C);
   }
 
+  if (format == MemoryFormat::Preserve) {
+    // Preserve means "keep current format" - not applicable for checking contiguity
+    return false;
+  }
+
   return false;
 }
 
