[xla:gpu] Update cutlass version and include util headers

PiperOrigin-RevId: 591825221

Author: ezhulenev

third_party/cutlass.BUILD

@@ -17,8 +17,21 @@ filegroup(
     ]),
 )
 
+filegroup(
+    name = "cutlass_util_header_files",
+    srcs = glob([
+        "tools/util/include/**",
+    ]),
+)
+
 cc_library(
     name = "cutlass",
-    hdrs = [":cutlass_header_files"],
-    strip_include_prefix = "/include",
+    hdrs = [
+        ":cutlass_header_files",
+        ":cutlass_util_header_files",
+    ],
+    includes = [
+        "include",
+        "tools/util/include",
+    ],
 )

workspace2.bzl

@@ -42,9 +42,9 @@ def _tf_repositories():
     tf_http_archive(
         name = "cutlass_archive",
         build_file = "//third_party:cutlass.BUILD",
-        sha256 = "ea1b7f96919460a5d80b09c1b246652539a8605600b2be4cccc02c254bccbe50",
-        strip_prefix = "cutlass-5783d6dbd0c34032371cce2bd999fc76007520d7",
-        urls = tf_mirror_urls("https://github.com/chsigg/cutlass/archive/5783d6dbd0c34032371cce2bd999fc76007520d7.tar.gz"),
+        sha256 = "84cf3fcc47c440a8dde016eb458f8d6b93b3335d9c3a7a16f388333823f1eae0",
+        strip_prefix = "cutlass-afa7b7241aabe598b725c65480bd9fa71121732c",
+        urls = tf_mirror_urls("https://github.com/chsigg/cutlass/archive/afa7b7241aabe598b725c65480bd9fa71121732c.tar.gz"),
     )
 
     tf_http_archive(

xla/service/gpu/kernels/cutlass_gemm_adaptor.cu.h

@@ -20,10 +20,12 @@ limitations under the License.
 #include <cstdint>
 #include <memory>
 
+#include "cute/layout.hpp"  // IWYU pragma: keep
 #include "cutlass/cutlass.h"
 #include "cutlass/gemm/gemm_enumerated_types.h"
 #include "cutlass/gemm_coord.h"
 #include "cutlass/layout/matrix.h"
+#include "cutlass/util/packed_stride.hpp"  // IWYU pragma: keep
 #include "xla/service/gpu/kernels/cutlass_gemm.h"
 
 namespace xla::gpu::kernel::gemm_universal {