fixed GEMM dimensions based on operand transpose flags in TP overlap

Signed-off-by: Alp Dener <[email protected]>

Author: denera

transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp

@@ -232,9 +232,9 @@ void CommOverlapBase::atomic_gemm_overlap_rs(TensorWrapper &A, bool transa, Tens
   _ub_comm->sms = _num_comm_sm;
   _ub_comm->cga_size = _cga_size;
   // Get GEMM dimensions
-  size_t m = A.size(0);
-  size_t k = A.size(1);
-  size_t n = B.size(0);
+  size_t m = (transa) ? A.size(0) : A.size(1);
+  size_t k = (transa) ? A.size(1) : A.size(0);
+  size_t n = (transb) ? B.size(1) : B.size(0);
   size_t m_chunk = m / _num_splits;
   size_t workspace_size_chunk = workspace.numel() / _stream_compute.size();
 
@@ -332,9 +332,9 @@ void CommOverlapBase::split_overlap_rs(TensorWrapper &A, bool transa, TensorWrap
   _ub_comm->use_ce = _use_ce;
   _ub_comm->sms = _num_comm_sm;
   _ub_comm->cga_size = _cga_size;
-  size_t m = A.size(0);
-  size_t k = A.size(1);
-  size_t n = B.size(0);
+  size_t m = (transa) ? A.size(0) : A.size(1);
+  size_t k = (transa) ? A.size(1) : A.size(0);
+  size_t n = (transb) ? B.size(1) : B.size(0);
   size_t m_chunk = m / _num_splits;
   size_t input_a_chunk_size = m_chunk * k;
   size_t output_chunk_size = n * m_chunk;
@@ -930,8 +930,8 @@ void CommOverlapP2PBase::split_overlap_rs(TensorWrapper &A, bool transa, TensorW
   _ub_comm->use_ce = _use_ce;
   _ub_comm->sms = _num_comm_sm;
   _ub_comm->cga_size = _cga_size;
-  size_t k = A.size(1);
-  size_t n = B.size(0);
+  size_t k = (transa) ? A.size(1) : A.size(0);
+  size_t n = (transb) ? B.size(1) : B.size(0);
 
   // Get communication and GEMM input chunk sizes
   size_t n_chunk = n / _tp_size;