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

[llvmlistbot at llvm.org]

Author: Adam Paszke
Date: 2023-12-19T08:40:26+01:00
New Revision: 12e4332501bca3bc1e29be94e134eea231578985

URL: https://github.com/llvm/llvm-project/commit/12e4332501bca3bc1e29be94e134eea231578985
DIFF: https://github.com/llvm/llvm-project/commit/12e4332501bca3bc1e29be94e134eea231578985.diff

LOG: [mlir][nvgpu] Fix the TMA stride setup (#75838)

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

Added: 
    

Modified: 
    mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

Removed: 
    


################################################################################
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,