Fix SIGABRT on zero-extent slice copy (SliceCopyStartGreaterThanDimSize)

lib/Conversion/TorchToTMTensor/TorchToTMTensor.cpp

@@ -617,16 +617,10 @@ getBroadcastShape(Location loc, llvm::ArrayRef<Value> indices, OpBuilder b) {
     indicesRank = std::max(rank, indicesRank);
   }
 
-  auto maxDim = [](int64_t dim0, int64_t dim1) {
-    if (dim0 == Torch::kUnknownSize || dim1 == Torch::kUnknownSize)
-      return Torch::kUnknownSize;
-    return std::max(dim0, dim1);
-  };
-
   Value torchCstOne =
       Torch::ConstantIntOp::create(b, loc, b.getI64IntegerAttr(1));
   llvm::SmallVector<Value> broadcastSizes(indicesRank, torchCstOne);
-  llvm::SmallVector<int64_t> broadcastShape(indicesRank, 0);
+  llvm::SmallVector<int64_t> broadcastShape(indicesRank, 1);
   for (auto index : indices) {
     auto indexTy = cast<Torch::ValueTensorType>(index.getType());
     auto shape = indexTy.getSizes();
@@ -638,9 +632,21 @@ getBroadcastShape(Location loc, llvm::ArrayRef<Value> indices, OpBuilder b) {
       auto size = shape[j];
 
       int32_t idx = broadcastShape.size() - rank + j;
-      broadcastSizes[idx] =
-          Torch::PrimMaxIntOp::create(b, loc, sizeOp, broadcastSizes[idx]);
-      broadcastShape[idx] = maxDim(size, broadcastShape[idx]);
+      // Numpy broadcasting fold: acc' = (size == 1) ? acc : size.
+      // `PrimMaxIntOp` is wrong here because it maps a size-0 dimension
+      // (from an empty index/slice) to 1 instead of preserving 0.
+      Value acc = broadcastSizes[idx];
+      Value isOne = Torch::AtenEqIntOp::create(b, loc, sizeOp, torchCstOne);
+      Value isOneInt = Torch::AtenIntBoolOp::create(b, loc, isOne);
+      Value diff = Torch::AtenSubIntOp::create(b, loc, acc, sizeOp);
+      Value scaled = Torch::AtenMulIntOp::create(b, loc, isOneInt, diff);
+      broadcastSizes[idx] = Torch::AtenAddIntOp::create(b, loc, sizeOp, scaled);
+
+      if (size == Torch::kUnknownSize ||
+          broadcastShape[idx] == Torch::kUnknownSize)
+        broadcastShape[idx] = Torch::kUnknownSize;
+      else if (size != 1)
+        broadcastShape[idx] = size;
     }
   }
   return std::make_pair(broadcastSizes, broadcastShape);

projects/pt1/e2e_testing/xfail_sets.py

@@ -76,12 +76,6 @@
     "GridSamplerBasic2_basic",
     "GridSamplerBasic3_basic",
     "GridSamplerBasic4_basic",
-    # Runtime op verification: stride mismatch in memref.cast
-    "ReduceAllDimEmpty_basic",
-    "TraceUnsignedIntModule_empty",
-    "TraceModule_empty",
-    # Crashes due to copy to a smaller destination buffer than the source buffer.
-    "SliceCopyStartGreaterThanDimSize_Module_basic",
     # unimplemented: for conversion to byte or char type dstOriginalDtype has to be passed to convertScalarToDtype
     "AtenMmInt8Types_basic",
     "AtenMmInt8ZeroK_basic",
@@ -3109,6 +3103,7 @@
     "SignAndLogarithmOfDeterminantModule_F32",
     "SignAndLogarithmOfDeterminantBatchedModule_F32",
     "SignAndLogarithmOfDeterminantDynamicModule_F32",
+    "SliceCopyStartGreaterThanDimSize_Module_basic",
     "SliceEndSleStartModule_basic",
     "SliceOutOfUpperBoundIndexModule_basic",
     "SliceStartEqEndModule_basic",
@@ -3138,8 +3133,6 @@
     "TorchPrimLoopWhileLikeModule_basic",
     # RuntimeError: Detected that you are using FX to torch.jit.trace a dynamo-optimized function
     "TorchPrimLoopWhileLikeHOPModule_basic",
-    "TraceModule_empty",
-    "TraceUnsignedIntModule_empty",
     "TransposedConv1dNegativePaddingUnitStrideDyn_basic",
     "UniformModule_basic",
     "UniformNoCorrelationModule_basic",

projects/pt1/python/test/refbackend/zero_extent_memref.py

@@ -0,0 +1,39 @@
+# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+# Also available under a BSD-style license. See LICENSE.
+
+# RUN: %PYTHON %s | FileCheck %s
+
+import numpy as np
+
+from torch_mlir_e2e_test.linalg_on_tensors_backends.refbackend import (
+    _canonicalize_zero_element_strides,
+)
+
+
+def print_element_strides(shape):
+    array = np.empty(shape, dtype=np.float32)
+    array = np.lib.stride_tricks.as_strided(
+        array, shape=shape, strides=(0,) * len(shape)
+    )
+    canonical = _canonicalize_zero_element_strides(array)
+    strides = tuple(stride // canonical.itemsize for stride in canonical.strides)
+    print(f"{shape}: {strides}")
+
+
+print_element_strides((0,))
+# CHECK: (0,): (1,)
+
+print_element_strides((0, 10, 3))
+# CHECK: (0, 10, 3): (30, 3, 1)
+
+print_element_strides((5, 0, 0))
+# CHECK: (5, 0, 0): (0, 0, 1)
+
+print_element_strides((1, 0, 3))
+# CHECK: (1, 0, 3): (0, 3, 1)
+
+nonempty = np.empty((2, 3), dtype=np.float32)
+print(_canonicalize_zero_element_strides(nonempty) is nonempty)
+# CHECK: True

projects/pt1/python/torch_mlir_e2e_test/linalg_on_tensors_backends/refbackend.py

@@ -62,6 +62,24 @@ def assert_arg_type_is_supported(ty):
 CONSUME_RETURN_FUNC_PREFIX = "refbackend_consume_func_return_"
 
 
+def _canonicalize_zero_element_strides(array):
+    if array.size != 0:
+        return array
+
+    element_strides = [1] * array.ndim
+    for dim in range(array.ndim - 2, -1, -1):
+        element_strides[dim] = element_strides[dim + 1] * array.shape[dim + 1]
+
+    # NumPy reports zero strides for empty arrays converted from PyTorch.
+    # Empty arrays have no addressable elements, so presenting the canonical
+    # contiguous strides at the memref ABI boundary is semantics-preserving.
+    return np.lib.stride_tricks.as_strided(
+        array,
+        shape=array.shape,
+        strides=tuple(stride * array.itemsize for stride in element_strides),
+    )
+
+
 def get_return_funcs(module):
     return_prefix_len = len(CONSUME_RETURN_FUNC_PREFIX)
     return_funcs = []
@@ -123,6 +141,7 @@ def invoke(*args):
             ffi_args = []
             for arg in args:
                 assert_arg_type_is_supported(arg.dtype)
+                arg = _canonicalize_zero_element_strides(arg)
                 ffi_args.append(
                     ctypes.pointer(ctypes.pointer(get_unranked_memref_descriptor(arg)))
                 )