From 0568a7a03c598b76d12cd8d72c30d88e779e9e89 Mon Sep 17 00:00:00 2001
From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com>
Date: Fri, 16 Jan 2026 15:08:41 -0600
Subject: [PATCH] Add static_for_indexed for reduced template instantiation

Provide static_for_indexed as an alternative to static_for that takes
a struct with templated operator() instead of a lambda. This reduces
instantiation bloat because struct types are named and reused, unlike
lambdas which create unique types at each call site.

Usage:
  struct MyLoop {
      Array& arr;
      template <index_t I>
      void operator()() const { arr[I] = I * 2; }
  };
  static_for_indexed<0, 8, 1>{}(MyLoop{arr});

Note: This is opt-in for high-impact code paths. The example_grouped_conv
target has minimal static_for usage (26 instantiations), but larger targets
like blockwise_gemm_pipeline (90+ usages per file) would benefit more.
---
 include/ck/utility/common_header.hpp      |  1 +
 include/ck/utility/static_for_indexed.hpp | 97 +++++++++++++++++++++++
 2 files changed, 98 insertions(+)
 create mode 100644 include/ck/utility/static_for_indexed.hpp

diff --git a/include/ck/utility/common_header.hpp b/include/ck/utility/common_header.hpp
index 78c3b78de1..39bf305a28 100644
--- a/include/ck/utility/common_header.hpp
+++ b/include/ck/utility/common_header.hpp
@@ -12,6 +12,7 @@
 #include "ck/utility/data_type.hpp"
 #include "ck/utility/functional.hpp"
 #include "ck/utility/functional2.hpp"
+#include "ck/utility/static_for_indexed.hpp"
 #include "ck/utility/functional3.hpp"
 #include "ck/utility/functional4.hpp"
 #include "ck/utility/enable_if.hpp"
diff --git a/include/ck/utility/static_for_indexed.hpp b/include/ck/utility/static_for_indexed.hpp
new file mode 100644
index 0000000000..15d5a2c9c1
--- /dev/null
+++ b/include/ck/utility/static_for_indexed.hpp
@@ -0,0 +1,97 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include "ck/utility/sequence.hpp"
+#include "ck/utility/number.hpp"
+
+namespace ck {
+
+namespace detail {
+
+// Helper to call functor with compile-time index as template parameter
+// This avoids per-lambda instantiation by using a struct with templated operator()
+template <typename F, index_t I>
+__host__ __device__ __forceinline__ constexpr void call_at_index(F& f)
+{
+    f.template operator()<I>();
+}
+
+template <typename Seq>
+struct static_for_indexed_impl;
+
+template <index_t... Is>
+struct static_for_indexed_impl<Sequence<Is...>>
+{
+    template <typename F>
+    __host__ __device__ constexpr void operator()(F& f) const
+    {
+        (call_at_index<F, Is>(f), ...);
+    }
+
+    template <typename F>
+    __host__ __device__ constexpr void operator()(const F& f) const
+    {
+        (call_at_index<const F, Is>(const_cast<const F&>(f)), ...);
+    }
+};
+
+} // namespace detail
+
+// static_for_indexed: Optimized compile-time loop that reduces template instantiation bloat
+//
+// Unlike static_for which takes a lambda (creating unique instantiation per lambda type),
+// static_for_indexed expects a functor with a templated operator():
+//
+//   struct MyLoop {
+//       SomeType& data;
+//       template <index_t I>
+//       __host__ __device__ void operator()() const { data[Number<I>{}] = I * 2; }
+//   };
+//   static_for_indexed<0, N, 1>{}(MyLoop{data});
+//
+// This reduces instantiations because MyLoop is a named type reused across call sites,
+// rather than unique lambda types at each usage.
+//
+template <index_t NBegin, index_t NEnd, index_t Increment = 1>
+struct static_for_indexed
+{
+    template <class F>
+    __host__ __device__ constexpr void operator()(F& f) const
+    {
+        detail::static_for_indexed_impl<
+            typename arithmetic_sequence_gen<NBegin, NEnd, Increment>::type>{}(f);
+    }
+
+    template <class F>
+    __host__ __device__ constexpr void operator()(const F& f) const
+    {
+        detail::static_for_indexed_impl<
+            typename arithmetic_sequence_gen<NBegin, NEnd, Increment>::type>{}(f);
+    }
+};
+
+// Convenience macro for inline loop body definition
+// Usage:
+//   CK_STATIC_FOR_INDEXED(0, 8, 1, I, {
+//       arr[Number<I>{}] = I * 2;
+//   });
+//
+// Note: Captures must be handled via the struct members, not lambda-style captures
+#define CK_STATIC_FOR_INDEXED_IMPL2(begin, end, incr, idx, captures, body) \
+    do                                                                     \
+    {                                                                      \
+        struct CK_LocalLoop                                                \
+        {                                                                  \
+            captures template <ck::index_t idx>                            \
+            __host__ __device__ void operator()() const body               \
+        } ck_loop_body;                                                    \
+        ck::static_for_indexed<begin, end, incr>{}(ck_loop_body);          \
+    } while(0)
+
+// Simpler macro without explicit captures (for self-contained loop bodies)
+#define CK_STATIC_FOR_INDEXED(begin, end, incr, idx, body) \
+    CK_STATIC_FOR_INDEXED_IMPL2(begin, end, incr, idx, , body)
+
+} // namespace ck