[CK_TILE] Add indexing to pooling operator (Lwpck 3892) (#3013)

* Add indexing support to pooling operator

- Add IndexDataType template parameter to pooling problem and kernel
definitions

- Enable pooling kernel to output indices of selected elements during
max/absmax pooling

- Add overloaded operators for Max and AbsMax that track when values
change using bool changed parameter

-  Support optional index buffer allocation and management in device
memory

- Modify BlockReduce2d classes to handle index tensors alongside value
tensors

-  Add separate shared memory allocation for index data in cross-warp
reductions

- Create validate_pool_indices function to verify index correctness

- Modify pool3d.cpp example to demonstrate index output functionality

- Add tests for index output

* fixes

* Refactor BlockReduce2D functions to get rid auxiliary private types.

* comment resolutions and some changes to block_reduce2d

- index reference implementation improved
- reduce_operator.hpp cleanedup
- updated the block_reduce2d.hpp to have index calculation for
BlockReduce2dLinearCrossWarpSync as well

* conditionally used variable declaration improvement

- the conditionally used vairbales are used only when indexing is
enabled. To inform the compiler that they may be unused and declare them
with least size possible. This may allow it to be optimized compared to
the previous declarations

* comment resolutions

* lexical ordering of the indicies

- introduced accumulate methods that handle the intermediate steps if
needed to order the indexes

* add reduce_operator_accumulate.hpp to core.hpp

---------

Co-authored-by: Adam Osewski <Adam.Osewski@amd.com>

include/ck_tile/ops/pooling/pipeline/pool_default_policy.hpp

@@ -32,7 +32,8 @@ struct PoolDefaultPolicy
     {
         using P_ = BlockReduce2dProblem<typename Problem::InDataType,
                                         typename Problem::ComputeDataType,
-                                        typename Problem::BlockShape>;
+                                        typename Problem::BlockShape,
+                                        Problem::kOutputIndex>;
         return BlockReduce2d<P_>{};
     }
 
@@ -41,7 +42,8 @@ struct PoolDefaultPolicy
     {
         using P_ = BlockReduce2dProblem<typename Problem::InDataType,
                                         typename Problem::ComputeDataType,
-                                        typename Problem::BlockShape>;
+                                        typename Problem::BlockShape,
+                                        Problem::kOutputIndex>;
         return BlockReduce2dSync<P_>{};
     }
 
@@ -50,7 +52,8 @@ struct PoolDefaultPolicy
     {
         using P_ = BlockReduce2dProblem<typename Problem::InDataType,
                                         typename Problem::ComputeDataType,
-                                        typename Problem::BlockShape>;
+                                        typename Problem::BlockShape,
+                                        Problem::kOutputIndex>;
         return BlockReduce2dCrossWarpSync<P_>{};
     }
 
@@ -61,7 +64,8 @@ struct PoolDefaultPolicy
         {
             using P_ = BlockReduce2dProblem<typename Problem::InDataType,
                                             typename Problem::ComputeDataType,
-                                            typename Problem::BlockShape>;
+                                            typename Problem::BlockShape,
+                                            Problem::kOutputIndex>;
 
             using block_reduce2d = BlockReduce2d<P_>;
             using x_block_tile =
@@ -76,5 +80,25 @@ struct PoolDefaultPolicy
             return 1; // zero size arrays are an extension
         }
     }
+
+    template <typename Problem>
+    CK_TILE_HOST_DEVICE static constexpr index_t GetIndicesSmemSize()
+    {
+
+        using P_ = BlockReduce2dProblem<typename Problem::InDataType,
+                                        typename Problem::ComputeDataType,
+                                        typename Problem::BlockShape,
+                                        Problem::kOutputIndex>;
+
+        using block_reduce2d = BlockReduce2d<P_>;
+        using x_block_tile = decltype(make_static_distributed_tensor<typename Problem::InDataType>(
+            MakeXBlockTileDistribution<Problem>()));
+        using y_index_block_tile = decltype(block_reduce2d::template MakeYIndexBlockTile<
+                                            x_block_tile,
+                                            typename Problem::IndexDataType>());
+
+        return GetBlockReduce2dCrossWarpSync<Problem>()
+            .template GetIndicesSmemSize<y_index_block_tile>();
+    }
 };
 } // namespace ck_tile

include/ck_tile/ops/pooling/pipeline/pool_problem.hpp

@@ -26,6 +26,8 @@ struct PoolProblem
     using OutputIndex     = bool_constant<OutputIndex_>;
     using PropagateNan    = bool_constant<PropagateNan_>;
 
+    static constexpr bool kOutputIndex       = OutputIndex_;
+    static constexpr bool kPropagateNan      = PropagateNan_;
     static constexpr bool kNeedCrossLaneSync = BlockShape::ThreadPerWarp_N > 1;
     static constexpr bool kNeedCrossWarpSync = BlockShape::WarpPerBlock_N > 1;
 };