[Mlir-commits] [mlir] [mlir][nvgpu] Fix the TMA stride setup (PR #75838)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-execution-engine

Author: Adam Paszke (apaszke)

<details>
<summary>Changes</summary>

There were two issues with the previous computation:
* it never looked at dimensions past the second one
* the definition was recursive, making each dimension have an extra `elementSize` power

---
Full diff: https://github.com/llvm/llvm-project/pull/75838.diff


1 Files Affected:

- (modified) mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp (+1-2) 


``````````diff
diff --git a/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp b/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
index 5ec87d58cc57f8..c45320a674568a 100644
--- a/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
+++ b/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
@@ -487,8 +487,7 @@ extern "C" MLIR_CUDA_WRAPPERS_EXPORT void *mgpuTensorMapEncodeTiledMemref(
 
   globalStrides[0] = globalDim[0] * elementSizeInBytes[tensorDataType];
   for (int r = 1; r < tensorRank - 1; r++)
-    globalStrides[r] = globalStrides[r - 1] * globalDim[1] *
-                       elementSizeInBytes[tensorDataType];
+    globalStrides[r] = globalStrides[r - 1] * globalDim[r];
 
   ScopedContext scopedContext;
   mgpuTensorMapEncodeTiled(&tensorMap, tensorDataType, tensorRank32,

``````````

</details>


https://github.com/llvm/llvm-project/pull/75838