[API/MemRef] Implement canonical stride validation for MemRefValue

Add stride validation in MemRefValue::create to compute canonical stride and compare against given strides. We need to handle special cases for zero-sized and unit-sized dimensions since frameworks deal with them arbitrarily while converting to the corresponding DLPack tensor. Add Python tests to verify both canonical and non-canonical stride validation.

Author: jhalakpatel

mlir-tensorrt/executor/include/mlir-executor-c/Runtime/Runtime.h

@@ -130,7 +130,8 @@ MLIR_CAPI_EXPORTED MTRT_Status
 mtrtMemRefCreate(MTRT_RuntimeClient client, MTRT_PointerType pointerKind,
                  int64_t bitsPerElement, int64_t rank, const int64_t *shape,
                  const int64_t *strides, MTRT_Device device, MTRT_Stream stream,
-                 MTRT_ScalarTypeCode scalarType, MTRT_MemRefValue *result);
+                 MTRT_ScalarTypeCode scalarType, MTRT_MemRefValue *result,
+                 bool assertCanonicalStrides = false);
 
 /// Creates an externally managed MemRef value. The caller provides all the
 /// metadata for the MemRef including the shape, strides (in elements), pointer,
@@ -142,7 +143,8 @@ MLIR_CAPI_EXPORTED MTRT_Status mtrtMemRefCreateExternal(
     MTRT_RuntimeClient client, MTRT_PointerType pointerKind,
     int64_t bitsPerElement, uintptr_t ptr, int64_t offset, int64_t rank,
     const int64_t *shape, const int64_t *strides, MTRT_Device device,
-    MTRT_ScalarTypeCode scalarType, MTRT_MemRefValue *result);
+    MTRT_ScalarTypeCode scalarType, MTRT_MemRefValue *result,
+    bool assertCanonicalStrides = false);
 
 /// Destroys `MTRT_MemRefValue` in a potentially asynchronous manner.
 /// If `buffer` is a device buffer, device memory is freed in the stream

mlir-tensorrt/executor/include/mlir-executor/Runtime/API/API.h

@@ -647,7 +647,8 @@ class MemRefValue : public RuntimeValue {
          int64_t bitsPerElement, uintptr_t ptr, int64_t offset,
          llvm::ArrayRef<int64_t> shape, llvm::ArrayRef<int64_t> strides,
          std::optional<const Device *> device,
-         std::optional<ScalarType> scalarType);
+         std::optional<ScalarType> scalarType,
+         std::optional<bool> assertCanonicalStrides = false);
 
   mlirtrt::runtime::PointerType getBufferKind() { return addressSpace; }
   int64_t getElementBitWidth() const { return bitsPerElement; }
@@ -917,15 +918,17 @@ class RuntimeClient {
                  llvm::ArrayRef<int64_t> shape, llvm::ArrayRef<int64_t> strides,
                  std::optional<const Device *> device = {},
                  std::optional<CudaStream> stream = {},
-                 std::optional<ScalarType> scalarType = {});
+                 std::optional<ScalarType> scalarType = {},
+                 std::optional<bool> assertCanonicalStrides = {});
 
   StatusOr<std::unique_ptr<MemRefValue>>
   createExternalMemRef(PointerType addressSpace, int64_t bitsPerElement,
                        uintptr_t ptr, int64_t offset,
                        llvm::ArrayRef<int64_t> shape,
                        llvm::ArrayRef<int64_t> strides,
                        std::optional<const Device *> device = {},
-                       std::optional<ScalarType> scalarType = {});
+                       std::optional<ScalarType> scalarType = {},
+                       std::optional<bool> assertCanonicalStrides = {});
 
   /// Frees the memory in `value`. The `stream` may optionally be provided
   /// for resources that can be deallocated asynchronously.

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

@@ -231,7 +231,8 @@ MTRT_Status
 mtrtMemRefCreate(MTRT_RuntimeClient client, MTRT_PointerType pointerKind,
                  int64_t bitsPerElement, int64_t rank, const int64_t *shape,
                  const int64_t *strides, MTRT_Device device, MTRT_Stream stream,
-                 MTRT_ScalarTypeCode scalarType, MTRT_MemRefValue *result) {
+                 MTRT_ScalarTypeCode scalarType, MTRT_MemRefValue *result,
+                 bool assertCanonicalStrides) {
   StatusOr<std::unique_ptr<MemRefValue>> bufferImpl =
       unwrap(client)->allocateMemRef(
           unwrap(pointerKind), bitsPerElement,
@@ -244,7 +245,8 @@ mtrtMemRefCreate(MTRT_RuntimeClient client, MTRT_PointerType pointerKind,
               : std::optional(unwrap(stream)->getRawStream()),
           scalarType != MTRT_ScalarTypeCode::MTRT_ScalarTypeCode_unknown
               ? std::optional(ScalarType(unwrap(scalarType)))
-              : std::nullopt);
+              : std::nullopt,
+          std::optional(assertCanonicalStrides));
 
   if (bufferImpl.isError())
     return wrap(bufferImpl.getStatus());
@@ -257,7 +259,8 @@ MTRT_Status mtrtMemRefCreateExternal(
     MTRT_RuntimeClient client, MTRT_PointerType pointerKind,
     int64_t bitsPerElement, uintptr_t ptr, int64_t offset, int64_t rank,
     const int64_t *shape, const int64_t *strides, MTRT_Device device,
-    MTRT_ScalarTypeCode scalarType, MTRT_MemRefValue *result) {
+    MTRT_ScalarTypeCode scalarType, MTRT_MemRefValue *result,
+    bool assertCanonicalStrides) {
   StatusOr<std::unique_ptr<MemRefValue>> bufferImpl =
       unwrap(client)->createExternalMemRef(
           unwrap(pointerKind), bitsPerElement, ptr, offset,
@@ -267,7 +270,8 @@ MTRT_Status mtrtMemRefCreateExternal(
                                    : std::optional(unwrap(device)),
           scalarType == MTRT_ScalarTypeCode_unknown
               ? std::nullopt
-              : std::optional(ScalarType(unwrap(scalarType))));
+              : std::optional(ScalarType(unwrap(scalarType))),
+          std::optional(assertCanonicalStrides));
 
   if (bufferImpl.isError())
     return wrap(bufferImpl.getStatus());

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

@@ -671,12 +671,50 @@ 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);
+        // For dimensions with size 0 or 1, the stride can be arbitrary.
+        // We set it to 1 here, but other values would also be valid.
+        if (isFirstZeroDim || shape[dimIndex] == 1)
+            canonicalStride[dimIndex] = 1;
+        else
+            canonicalStride[dimIndex] = cumulativeProduct;
+        // For zero-sized dimensions (except the last one), we don't update the cumulative product
+        // This allows for consistent handling of zero-sized dimensions across different frameworks
+        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)
+    // Allow arbitrary strides for dimensions with size 0 or 1
+    // This accounts for discrepancies in how different frameworks handle these cases
+    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,
     llvm::ArrayRef<int64_t> shape, llvm::ArrayRef<int64_t> strides,
-    std::optional<const Device *> device,
-    std::optional<ScalarType> scalarType) {
+    std::optional<const Device *> device, std::optional<ScalarType> scalarType,
+    std::optional<bool> assertCanonicalStrides) {
   if (!client)
     return getInvalidArgStatus("a valid RuntimeClient must be provided to "
                                "create a tracked MemRef object");
@@ -691,6 +729,19 @@ 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 (assertCanonicalStrides && *assertCanonicalStrides) {
+    std::vector<int64_t> canonicalStride = getCanonicalStride(shape);
+    if (!strides.empty() &&
+        !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));
@@ -777,7 +828,7 @@ StatusOr<std::unique_ptr<MemRefValue>> RuntimeClient::allocateMemRef(
     PointerType addressSpace, int64_t bitsPerElement,
     llvm::ArrayRef<int64_t> shape, llvm::ArrayRef<int64_t> strides,
     std::optional<const Device *> device, std::optional<CudaStream> stream,
-    std::optional<ScalarType> scalarType) {
+    std::optional<ScalarType> scalarType, std::optional<bool> assertCanonicalStrides) {
   if (addressSpace == PointerType::device ||
       addressSpace == PointerType::unified) {
     if (!device || !*device)
@@ -800,7 +851,7 @@ StatusOr<std::unique_ptr<MemRefValue>> RuntimeClient::allocateMemRef(
   // Create the descriptor.
   StatusOr<std::unique_ptr<MemRefValue>> bufferImpl =
       MemRefValue::create(this, addressSpace, bitsPerElement, allocation->ptr,
-                          0, shape, strides, device, scalarType);
+                          0, shape, strides, device, scalarType, assertCanonicalStrides);
   if (bufferImpl.isError())
     return bufferImpl.getStatus();
 
@@ -811,11 +862,11 @@ StatusOr<std::unique_ptr<MemRefValue>> RuntimeClient::createExternalMemRef(
     PointerType addressSpace, int64_t bitsPerElement, uintptr_t ptr,
     int64_t offset, llvm::ArrayRef<int64_t> shape,
     llvm::ArrayRef<int64_t> strides, std::optional<const Device *> device,
-    std::optional<ScalarType> scalarType) {
+    std::optional<ScalarType> scalarType, std::optional<bool> assertCanonicalStrides) {
   // Create the descriptor.
   StatusOr<std::unique_ptr<MemRefValue>> memref =
       MemRefValue::create(this, addressSpace, bitsPerElement, ptr, offset,
-                          shape, strides, device, scalarType);
+                          shape, strides, device, scalarType, assertCanonicalStrides);
   if (!memref.isOk())
     return memref.getStatus();
 

mlir-tensorrt/python/bindings/Runtime/RuntimePyBind.cpp

@@ -371,11 +371,11 @@ createMemRefViewFromDLPack(PyRuntimeClient &client, py::capsule capsule) {
     s = mtrtMemRefCreateExternal(client, addressSpace, bytesPerElement * 8,
                                  reinterpret_cast<uintptr_t>(data), offset,
                                  rank, shape, strides, device, elementType,
-                                 &result);
+                                 &result, true /*assertCanonicalStrides*/);
   } else {
     s = mtrtMemRefCreate(client, addressSpace, bytesPerElement * 8, rank, shape,
                          strides, device, mtrtStreamGetNull(), elementType,
-                         &result);
+                         &result, true /*assertCanonicalStrides*/);
   }
 
   THROW_IF_MTRT_ERROR(s);

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