[API/MemRefValue] Validate strides against canonical strides for non-empty shapes

Author: jhalakpatel

mlir-tensorrt/executor/lib/Runtime/API/API.cpp

@@ -671,6 +671,38 @@ static StatusOr<int64_t> getFootprintInBytes(llvm::ArrayRef<int64_t> shape,
   return sizeBytes;
 }
 
+static std::vector<int64_t> getCanonicalStride(const llvm::ArrayRef<int64_t>& shape) {
+    if (shape.empty())
+        return {};
+
+    std::vector<int64_t> canonicalStride(shape.size(), 1);
+    int64_t cumulativeProduct = 1;
+
+    for (int64_t dimIndex = shape.size() - 1; dimIndex >= 0; --dimIndex) {
+        bool isFirstZeroDim = (shape[dimIndex] == 0 && dimIndex != static_cast<int64_t>(shape.size()) - 1);
+        if (isFirstZeroDim)
+            canonicalStride[dimIndex] = 1;
+        else if (shape[dimIndex] != 1)
+            canonicalStride[dimIndex] = cumulativeProduct;
+        cumulativeProduct *= isFirstZeroDim ? 1 : shape[dimIndex];
+    }
+
+    return canonicalStride;
+}
+
+static bool areStridesEquivalent(llvm::ArrayRef<int64_t> shape,
+                          llvm::ArrayRef<int64_t> stride,
+                          llvm::ArrayRef<int64_t> expectedStride) {
+  if (shape.size() != stride.size() || shape.size() != expectedStride.size())
+    return false;
+
+  for (size_t i = 0; i < shape.size(); ++i)
+    if (stride[i] != expectedStride[i] && shape[i] != 0 && shape[i] != 1)
+      return false;
+
+  return true;
+}
+
 StatusOr<std::unique_ptr<MemRefValue>> MemRefValue::create(
     RuntimeClient *client, mlirtrt::runtime::PointerType addressSpace,
     int64_t bitsPerElement, uintptr_t ptr, int64_t offset,
@@ -691,6 +723,17 @@ StatusOr<std::unique_ptr<MemRefValue>> MemRefValue::create(
     return getInvalidArgStatus("a specific device must be provided for MemRefs "
                                "that are device-visible");
 
+  // Check if given strides match canonical stride
+  if (!strides.empty() && !shape.empty()) {
+      std::vector<int64_t> canonicalStride = getCanonicalStride(shape);
+      if (!areStridesEquivalent(shape, strides, canonicalStride)) {      
+      std::string errorMsg = llvm::formatv(
+          "Given strides [{0}] do not match canonical strides [{1}] for shape [{2}]",
+          strides, llvm::ArrayRef(canonicalStride), shape);
+      return getInvalidArgStatus(errorMsg.c_str());
+    }
+  }
+
   return std::unique_ptr<MemRefValue>(
       new MemRefValue(client, addressSpace, bitsPerElement, ptr, offset, shape,
                       strides, device, scalarType));

mlir-tensorrt/test/python/mlir_tensorrt_runtime/test_create_memref.py

@@ -514,3 +514,59 @@ def create_dangling_memref():
 # CHECK-LABEL: Test memref maintains data's lifetime
 # CHECK-NEXT: -- Inner scope: np.from_dlpack(): [1 2]
 # CHECK-NEXT: -- Outer scope: np.from_dlpack(): [1 2]
+
+
+def check_non_canonical_stride(client):
+    try:
+        t = cp.arange(12, dtype=cp.float32).reshape(3, 4)
+        a = cp.transpose(t)
+        memref = client.create_memref_view_from_dlpack(a.__dlpack__())
+        print(
+            "Test failed: Expected exception was not thrown for CuPy non-canonical stride"
+        )
+        return False
+    except Exception as e:
+        error_message = str(e)
+        if (
+            "Given strides do not match canonical stride for the shape.\nShape: 4, 3\nExpected canonical stride: 3, 1\nGiven stride: 1, 4'"
+            in error_message
+        ):
+            print(
+                "Test passed: Correct error message received for non-canonical stride"
+            )
+            return True
+        else:
+            print(
+                f"Test failed: Unexpected error message for non-canonical stride: {error_message}"
+            )
+            return False
+
+
+def check_canonical_stride(client):
+    try:
+        t = cp.arange(12, dtype=cp.float32).reshape(3, 4)
+        memref = client.create_memref_view_from_dlpack(t.__dlpack__())
+        print("Test passed: No exception thrown for canonical stride")
+        return True
+    except Exception as e:
+        print(f"Test failed: Unexpected exception for canonical stride: {str(e)}")
+        return False
+
+
+def test_memref_strides():
+    print("Testing non-canonical stride:")
+    non_canonical_result = check_non_canonical_stride(client)
+
+    print("Testing canonical stride:")
+    canonical_result = check_canonical_stride(client)
+
+
+print("Test memref strides")
+test_memref_strides()
+
+# CHECK-LABEL: Test memref strides
+# CHECK-NEXT: Testing non-canonical stride:
+# CHECK-NEXT: Test failed: Unexpected error message for non-canonical stride: InvalidArgument:
+# CHECK-SAME: InvalidArgument: Given strides [1, 4] do not match canonical strides [3, 1] for shape [4, 3]
+# CHECK-NEXT: Testing canonical stride:
+# CHECK-NEXT: Test passed: No exception thrown for canonical stride