rever internal overlapping detection to 3-value classification: Yes, No and TooHard. TooHard is okay for torch to assume no internal overlapping. (#711)

Author: iclementine

src/flag_gems/ops/scatter.py

@@ -7,7 +7,11 @@
 
 from flag_gems.utils.code_cache import code_cache_dir
 from flag_gems.utils.code_utils import IndentedBuffer, write_atomic
-from flag_gems.utils.shape_utils import has_internal_overlapping, restride_dim
+from flag_gems.utils.shape_utils import (
+    MemOverlap,
+    has_internal_overlapping,
+    restride_dim,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -332,7 +336,7 @@ def scatter(inp, dim, index, src, reduce=None):
             torch.bfloat16,
         ), "Unsupported operation: reduce scatter bfloat tensors."
 
-    if has_internal_overlapping(out):
+    if has_internal_overlapping(out) == MemOverlap.Yes:
         out = out.contiguous()
 
     src_strided = src.as_strided(index.shape, src.stride())
@@ -367,8 +371,8 @@ def scatter_(inp, dim, index, src, reduce=None):
             torch.bfloat16,
         ), "Unsupported operation: reduce scatter bfloat tensors."
 
-    assert not has_internal_overlapping(
-        out
+    assert (
+        has_internal_overlapping(out) != MemOverlap.Yes
     ), "Unsupported operation: trying to inplace write to an internally overlapping tensor."
 
     src_restrided = src.as_strided(index.shape, src.stride())

src/flag_gems/ops/select_scatter.py

@@ -3,7 +3,7 @@
 import torch
 
 from ..ops.copy import copy
-from ..utils.shape_utils import has_internal_overlapping
+from ..utils.shape_utils import MemOverlap, has_internal_overlapping
 
 logger = logging.getLogger(__name__)
 
@@ -21,7 +21,7 @@ def select_scatter(inp, src, dim, index):
         list(src.shape) == valid_shape
     ), "Expected src to have a size equal to the slice of self"
 
-    if has_internal_overlapping(inp):
+    if has_internal_overlapping(inp) == MemOverlap.Yes:
         out = torch.empty(inp.size(), dtype=inp.dtype, device=inp.device)
     else:
         out = torch.empty_strided(

src/flag_gems/ops/slice_scatter.py

@@ -4,7 +4,7 @@
 import triton
 
 from ..ops.copy import copy
-from ..utils.shape_utils import has_internal_overlapping
+from ..utils.shape_utils import MemOverlap, has_internal_overlapping
 
 logger = logging.getLogger(__name__)
 
@@ -27,7 +27,7 @@ def slice_scatter(inp, src, dim=0, start=None, end=None, step=1):
         list(src.shape) == valid_shape
     ), "Expected src to have a size equal to the slice of self"
 
-    if has_internal_overlapping(inp):
+    if has_internal_overlapping(inp) == MemOverlap.Yes:
         out = torch.empty(inp.size(), dtype=inp.dtype, device=inp.device)
     else:
         out = torch.empty_strided(

src/flag_gems/runtime/backend/_kunlunxin/ops/scatter.py

@@ -7,7 +7,11 @@
 
 from flag_gems.utils.code_cache import code_cache_dir
 from flag_gems.utils.code_utils import IndentedBuffer
-from flag_gems.utils.shape_utils import has_internal_overlapping, restride_dim
+from flag_gems.utils.shape_utils import (
+    MemOverlap,
+    has_internal_overlapping,
+    restride_dim,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -327,7 +331,7 @@ def scatter(inp, dim, index, src, reduce=None):
             torch.bfloat16,
         ), "Unsupported operation: reduce scatter bfloat tensors."
 
-    if has_internal_overlapping(out):
+    if has_internal_overlapping(out) == MemOverlap.Yes:
         out = out.contiguous()
 
     src_strided = src.as_strided(index.shape, src.stride())
@@ -362,8 +366,8 @@ def scatter_(inp, dim, index, src, reduce=None):
             torch.bfloat16,
         ), "Unsupported operation: reduce scatter bfloat tensors."
 
-    assert not has_internal_overlapping(
-        out
+    assert (
+        has_internal_overlapping(out) != MemOverlap.Yes
     ), "Unsupported operation: trying to inplace write to an internally overlapping tensor."
 
     src_restrided = src.as_strided(index.shape, src.stride())

src/flag_gems/runtime/backend/_kunlunxin/ops/select_scatter.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from flag_gems.utils.shape_utils import has_internal_overlapping
+from flag_gems.utils.shape_utils import MemOverlap, has_internal_overlapping
 
 from ..ops.copy import copy
 
@@ -22,7 +22,7 @@ def select_scatter(inp, src, dim, index):
         list(src.shape) == valid_shape
     ), "Expected src to have a size equal to the slice of self"
 
-    if has_internal_overlapping(inp):
+    if has_internal_overlapping(inp) == MemOverlap.Yes:
         out = torch.empty(inp.size(), dtype=inp.dtype, device=inp.device)
     else:
         out = torch.empty_strided(

src/flag_gems/runtime/backend/_kunlunxin/ops/slice_scatter.py

@@ -4,7 +4,7 @@
 import triton
 from _kunlunxin.ops.copy import copy_slice
 
-from flag_gems.utils.shape_utils import has_internal_overlapping
+from flag_gems.utils.shape_utils import MemOverlap, has_internal_overlapping
 
 logger = logging.getLogger(__name__)
 
@@ -27,7 +27,7 @@ def slice_scatter(inp, src, dim=0, start=None, end=None, step=1):
         list(src.shape) == valid_shape
     ), "Expected src to have a size equal to the slice of self"
 
-    if has_internal_overlapping(inp):
+    if has_internal_overlapping(inp) == MemOverlap.Yes:
         out = torch.empty(inp.size(), dtype=inp.dtype, device=inp.device)
     else:
         out = torch.empty_strided(

src/flag_gems/runtime/backend/_kunlunxin/utils/pointwise_dynamic.py

@@ -10,6 +10,7 @@
 from flag_gems.utils.code_cache import code_cache_dir
 from flag_gems.utils.code_utils import IndentedBuffer
 from flag_gems.utils.shape_utils import (
+    MemOverlap,
     all_c_contiguous,
     all_the_same_shape,
     all_the_same_stride,
@@ -1215,15 +1216,13 @@ def prepare_args(self, *args, **kwargs):
             if out_tensors:
                 for index, item in enumerate(out_tensors):
                     if list(item.shape) != list(task_shape):
-                        raise ValueError(
+                        raise RuntimeError(
                             f"out tensor at index {index} shape is invalid, should be {task_shape} but is {item.shape}!"
                         )
-                    # output arguments must be dense and no overlapping for pointwise operation
-                    if has_internal_overlapping(item) and any(
-                        item is t for t in in_tensors
-                    ):
-                        raise ValueError(
-                            "Pointwise Input arguments must be dense and no overlapping."
+                    # output arguments must not have internal overlapping for pointwise operation
+                    if has_internal_overlapping(item) == MemOverlap.Yes:
+                        raise RuntimeError(
+                            "Pointwise Input arguments should not have internal overlapping."
                         )
 
             ndim = len(task_shape)

src/flag_gems/utils/pointwise_dynamic.py

@@ -9,6 +9,7 @@
 from flag_gems.utils.code_cache import code_cache_dir
 from flag_gems.utils.code_utils import IndentedBuffer, write_atomic
 from flag_gems.utils.shape_utils import (
+    MemOverlap,
     all_c_contiguous,
     all_the_same_shape,
     all_the_same_stride,
@@ -1169,15 +1170,13 @@ def prepare_args(self, *args, **kwargs):
             if out_tensors:
                 for index, item in enumerate(out_tensors):
                     if list(item.shape) != list(task_shape):
-                        raise ValueError(
+                        raise RuntimeError(
                             f"out tensor at index {index} shape is invalid, should be {task_shape} but is {item.shape}!"
                         )
-                    # output arguments must be dense and no overlapping for pointwise operation
-                    if has_internal_overlapping(item) and any(
-                        item is t for t in in_tensors
-                    ):
-                        raise ValueError(
-                            "Pointwise Input arguments must be dense and no overlapping."
+                    # output arguments must not have internal overlapping for pointwise operation
+                    if has_internal_overlapping(item) == MemOverlap.Yes:
+                        raise RuntimeError(
+                            "Pointwise Input arguments should not have internal overlapping."
                         )
 
             ndim = len(task_shape)

src/flag_gems/utils/shape_utils.py

@@ -1,3 +1,4 @@
+import enum
 import functools
 import operator
 from typing import Iterable, Sequence, Tuple
@@ -221,15 +222,21 @@ def can_use_int32_index(a):
     return True
 
 
+class MemOverlap(enum.Enum):
+    No = 0
+    Yes = 1
+    TooHard = 2
+
+
 def has_internal_overlapping(x: torch.Tensor):
     if x.is_contiguous():
-        return False
+        return MemOverlap.No
     if torch.ops.aten.is_non_overlapping_and_dense(x):
-        return False
+        return MemOverlap.No
     for size, stride in zip(x.size(), x.stride()):
         if size > 1 and stride == 0:
-            return True
-    return True
+            return MemOverlap.Yes
+    return MemOverlap.TooHard
 
 
 def restride_dim(src, dim, shape, step=0, storage_offset=None):
@@ -277,6 +284,37 @@ def add_on_kernel(
     tl.store(add_on + offsets, res, mask=block_mask)
 
 
+def check_tensor_attributes(data_list, is_tensor_list):
+    """
+    Checks if each element in data_list is a tensor and validates whether the corresponding
+    boolean value in is_tensor_list is correct.
+    Parameters:
+    - data_list: A list containing tensor and non-tensor objects.
+    - is_tensor_list: A list of boolean values indicating whether the corresponding element in data_list is a tensor.
+    Returns:
+    - True if all elements' types match their corresponding boolean values in is_tensor_list.
+    - Raise Error otherwise, and prints the index and element that do not match.
+    """
+    # Check if both lists have the same length
+    if len(data_list) != len(is_tensor_list):
+        raise ValueError(
+            "Error: The lists of inputs and is_tensor must have the same length."
+        )
+
+    for i, (data, is_tensor) in enumerate(zip(data_list, is_tensor_list)):
+        actual_is_tensor = isinstance(data, torch.Tensor)
+
+        if actual_is_tensor != is_tensor:
+            raise ValueError(
+                f"Element at index {i} is incorrect. Expected {is_tensor}, but got {actual_is_tensor}."
+            )
+
+    return True
+
+
+_initial_missing = object()
+
+
 def offset_calculator(inp, idx, strides, dim, isInp):
     """
     Calculate the flat index(a.k.a offset) for a given ravel index in a multi-dimensional array.
@@ -341,39 +379,6 @@ def offset_calculator(inp, idx, strides, dim, isInp):
     return offsets if not isInp else (offsets - idx_dim)
 
 
-def check_tensor_attributes(data_list, is_tensor_list):
-    """
-    Checks if each element in data_list is a tensor and validates whether the corresponding
-    boolean value in is_tensor_list is correct.
-
-    Parameters:
-    - data_list: A list containing tensor and non-tensor objects.
-    - is_tensor_list: A list of boolean values indicating whether the corresponding element in data_list is a tensor.
-
-    Returns:
-    - True if all elements' types match their corresponding boolean values in is_tensor_list.
-    - Raise Error otherwise, and prints the index and element that do not match.
-    """
-    # Check if both lists have the same length
-    if len(data_list) != len(is_tensor_list):
-        raise ValueError(
-            "Error: The lists of inputs and is_tensor must have the same length."
-        )
-
-    for i, (data, is_tensor) in enumerate(zip(data_list, is_tensor_list)):
-        actual_is_tensor = isinstance(data, torch.Tensor)
-
-        if actual_is_tensor != is_tensor:
-            raise ValueError(
-                f"Element at index {i} is incorrect. Expected {is_tensor}, but got {actual_is_tensor}."
-            )
-
-    return True
-
-
-_initial_missing = object()
-
-
 def offsetCalculator(inp, idx, strides, dim, isInp):
     ndim = inp.ndim
     shape = list(inp.shape)