Merge branch 'develop' of https://github.com/ROCmSoftwarePlatform/composable_kernel into wmma_op

include/ck/tensor_operation/gpu/grid/gridwise_set_buffer_value.hpp

@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
 #include "ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp"
 
 namespace ck {

include/ck/tensor_operation/gpu/warp/xdlops_gemm.hpp

@@ -594,6 +594,7 @@ struct XdlopsGemm
     static constexpr auto I5 = Number<5>{};
 
     using CIndex   = MultiIndex<2>;
+    using CIndex4D = MultiIndex<4>;
 
     __device__ static constexpr index_t GetNumBlks() { return mfma_instr.num_output_blks; }
 
@@ -822,6 +823,16 @@ struct XdlopsGemm
         return TransposeC ? CIndex{n_offset, m_offset} : CIndex{m_offset, n_offset};
     }
 
+    __device__ static CIndex4D GetBeginOfThreadBlk4D(index_t /* xdlops_i */, index_t /* blk_i */)
+    {
+        const auto blk_idx = GetBlkIdx();
+
+        const auto blk_id = blk_idx[I0];
+        const auto blk_td = blk_idx[I1];
+
+        return TransposeC ? CIndex4D{blk_td, I0, blk_id, I0} : CIndex4D{I0, blk_id, I0, blk_td};
+    }
+
     static constexpr auto mfma = MfmaSelector<base_type, MPerXdlops, NPerXdlops>{};
 
     static constexpr auto mfma_instr = mfma.selected_mfma;

library/include/ck/library/tensor_operation_instance/device_operation_instance_factory.hpp

@@ -3,7 +3,10 @@
 
 #pragma once
 
-#include <cstdlib>
+#include "ck/utility/data_type.hpp"
+#include "ck/utility/tuple.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
 
 namespace ck {
 namespace tensor_operation {
@@ -15,6 +18,8 @@ using F64  = double;
 using F32  = float;
 using F16  = ck::half_t;
 using BF16 = ck::bhalf_t;
+using I8   = int8_t;
+using I32  = int32_t;
 
 using Empty_Tuple = ck::Tuple<>;