diff --git a/ci/baseimage.cuda.Dockerfile b/ci/baseimage.cuda.Dockerfile
index 288d682..a022c37 100644
--- a/ci/baseimage.cuda.Dockerfile
+++ b/ci/baseimage.cuda.Dockerfile
@@ -63,7 +63,7 @@ RUN ldconfig
 # create environments for several configurations and install dependencies
 RUN spack env create -d /cosma-env-cuda && \
     spack -e /cosma-env-cuda add "cosma@master %gcc +cuda +tests +scalapack +shared ^mpich" && \
-    spack -e /cosma-env-cuda add "tiled-mm@master" && \
+    spack -e /cosma-env-cuda add "tiled-mm@2.3.1" && \
     spack -e /cosma-env-cuda develop -p /src cosma@master && \
     spack -e /cosma-env-cuda install --only=dependencies --fail-fast
 
@@ -75,7 +75,7 @@ RUN spack env create -d /cosma-env-cuda-gpu-direct && \
 
 RUN spack env create -d /cosma-env-cuda-nccl && \
     spack -e /cosma-env-cuda-nccl add "cosma@master %gcc +cuda +tests +scalapack +shared +nccl ^mpich " && \
-    spack -e /cosma-env-cuda-nccl add "tiled-mm@master" && \
+    spack -e /cosma-env-cuda-nccl add "tiled-mm@2.3.1" && \
     spack -e /cosma-env-cuda-nccl develop -p /src cosma@master && \
     spack -e /cosma-env-cuda-nccl install --only=dependencies --fail-fast
 
diff --git a/libs/Tiled-MM b/libs/Tiled-MM
index fce5dc1..85331eb 160000
--- a/libs/Tiled-MM
+++ b/libs/Tiled-MM
@@ -1 +1 @@
-Subproject commit fce5dc1b1021febd2b70cf9a506cb75ac23b0684
+Subproject commit 85331eb36ec45e644761e9d44604090a92e6c4d2
diff --git a/spack/packages/tiled-mm/package.py b/spack/packages/tiled-mm/package.py
new file mode 100644
index 0000000..c057a35
--- /dev/null
+++ b/spack/packages/tiled-mm/package.py
@@ -0,0 +1,52 @@
+# Copyright 2013-2024 Lawrence Livermore National Security, LLC and other
+# Spack Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+from spack.package import *
+
+
+class TiledMm(CMakePackage, CudaPackage, ROCmPackage):
+    """Matrix multiplication on GPUs for matrices stored on a CPU. Similar to cublasXt,
+    but ported to both NVIDIA and AMD GPUs."""
+
+    homepage = "https://github.com/eth-cscs/Tiled-MM/"
+    url = "https://github.com/eth-cscs/Tiled-MM/archive/refs/tags/v2.0.tar.gz"
+    git = "https://github.com/eth-cscs/Tiled-MM.git"
+
+    maintainers("mtaillefumier", "simonpintarelli", "RMeli")
+
+    license("BSD-3-Clause")
+
+    version("master", branch="master")
+
+    version("2.3.1", sha256="68914a483e62f796b790ea428210b1d5ef5943d6289e53d1aa62f56a20fbccc8")
+    version("2.3", sha256="504c6201f5a9be9741c55036bf8e2656ae3f4bc19996295b264ee5e303c9253c")
+    version("2.2", sha256="6d0b49c9588ece744166822fd44a7bc5bec3dc666b836de8bf4bf1a7bb675aac")
+    version("2.0", sha256="ea554aea8c53d7c8e40044e6d478c0e8137d7e8b09d7cb9650703430d92cf32e")
+
+    variant("shared", default=True, description="Build shared libraries")
+    variant("examples", default=False, description="Enable examples")
+    variant("tests", default=False, description="Enable tests")
+
+    depends_on("rocblas", when="+rocm")
+    depends_on("cxxopts", when="+tests")
+    depends_on("cxxopts", when="+examples")
+
+    conflicts("~cuda~rocm")
+    conflicts("+cuda", when="+rocm")
+
+    def cmake_args(self):
+        args = [
+            self.define_from_variant("BUILD_SHARED_LIBS", "shared"),
+            self.define_from_variant("TILEDMM_WITH_EXAMPLES", "examples"),
+            self.define_from_variant("TILEDMM_WITH_TESTS", "tests"),
+        ]
+
+        if "+rocm" in self.spec:
+            args.extend([self.define("TILEDMM_GPU_BACKEND", "ROCM")])
+
+        if "+cuda" in self.spec:
+            args.extend([self.define("TILEDMM_GPU_BACKEND", "CUDA")])
+
+        return args
diff --git a/tests/pdgemm.cpp b/tests/pdgemm.cpp
index facf477..ebdbb78 100644
--- a/tests/pdgemm.cpp
+++ b/tests/pdgemm.cpp
@@ -685,7 +685,7 @@ INSTANTIATE_TEST_CASE_P(
             // matrix dimensions
             43176, 217, // matrix A
             43176, 217, // matrix B
-            217, 217, // matrix C
+            2176, 217, // matrix C
 
             // block sizes
             1696, 108, // matrix A