Minor refactor.

codegen/CMakeLists.txt

@@ -21,14 +21,14 @@ list(APPEND CMAKE_MODULE_PATH ${CK_ROOT}/cmake)
 include(Embed)
 file(GLOB_RECURSE KERNEL_FILES CONFIGURE_DEPENDS
 	${CK_ROOT}/include/ck/*.hpp)
-#printouts fot debug purposes
+# printouts for debug purposes
 #message(STATUS "KERNEL_FILES: ${KERNEL_FILES}")
 #message(STATUS "RELATIVE: ${CK_ROOT}/include")
 add_embed_library(ck_headers ${KERNEL_FILES} RELATIVE ${CK_ROOT}/include)
 
 file(GLOB SOURCES CONFIGURE_DEPENDS src/*.cpp)
 
-##message(STATUS "SOURCE_FILES: ${SOURCES}")
+#message(STATUS "SOURCE_FILES: ${SOURCES}")
 # TODO: Use object library
 add_library(ck_host STATIC ${SOURCES})
 add_library(composable_kernel::ck_host ALIAS ck_host)
@@ -39,7 +39,6 @@ set_target_properties(ck_host PROPERTIES
 
 target_include_directories(ck_host SYSTEM PRIVATE
     $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
-    # $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}/library/src/jit_library/solution_instances>
     $<BUILD_INTERFACE:${CMAKE_CURRENT_BINARY_DIR}/solution_instances>
     $<BUILD_INTERFACE:${CMAKE_CURRENT_BINARY_DIR}/embed/ck_headers/include>
 )

codegen/include/ck/host/device_batched_gemm_softmax_gemm/operation.hpp

@@ -27,6 +27,8 @@ struct Operation_Xdl_CShuffle
     TensorDesc B{};
     TensorDesc B1{};
     TensorDesc C{};
+    DataType acc                    = DataType::Float;
+    DataType cs_type                = DataType::Half;
     std::string a_elem_op           = PassThrough;
     std::string b_elem_op           = PassThrough;
     std::string b1_elem_op          = PassThrough;

codegen/src/device_batched_gemm_softmax_gemm_operation_xdl_cshuffle.cpp

@@ -41,8 +41,6 @@ void Operation_Xdl_CShuffle::update_prologue(const std::string& pro)
     if(!prologue.empty())
     {
         this->prologue = pro;
-        // TODO
-        // this->cde_elem_op = "CDEElementOp";
     }
     else
     {
@@ -55,8 +53,6 @@ void Operation_Xdl_CShuffle::update_epilogue(const std::string& epi)
     if(!epilogue.empty())
     {
         this->epilogue = epi;
-        // TODO
-        // this->cde_elem_op = "CDEElementOp";
     }
     else
     {
@@ -336,8 +332,8 @@ Solution Operation_Xdl_CShuffle::ToSolution() const
         {"B0DataType", ToString(this->B.element)},
         {"B1DataType", ToString(this->B1.element)},
         {"CDataType", ToString(this->C.element)},
-        {"AccDataType", ToString(DataType::Float)},
-        {"CShuffleDataType", ToString(DataType::Half)},
+        {"AccDataType", ToString(this->acc)},
+        {"CShuffleDataType", ToString(this->cs_type)},
         {"AElementwiseOperation", this->a_elem_op},
         {"B0ElementwiseOperation", this->b_elem_op},
         {"Acc0ElementwiseOperation", this->acc_elem_op},