fix example call C0MatrixMask(N)

example/32_batched_gemm_scale_softmax_gemm/batched_multihead_attention_train_v1.cpp

@@ -669,8 +669,9 @@ void run_attention_fwd_host(const TensorQ& q_g_m_k,
     ref_gemm0_invoker.Run(ref_gemm0_argument);
 
     // masking
+    auto M          = s_g_m_n.GetLengths()[1];
     auto N          = s_g_m_n.GetLengths()[2];
-    const auto mask = DeviceGemmInstanceFWD::C0MatrixMask(N);
+    const auto mask = DeviceGemmInstanceFWD::C0MatrixMask(M, N);
     s_g_m_n.ForEach([&](auto& self, auto idx) {
         if(mask.IsMaskedElement(idx[1], idx[2]))
             self(idx) = -ck::NumericLimits<float>::Infinity();

example/32_batched_gemm_scale_softmax_gemm/batched_multihead_attention_train_v2.cpp

@@ -257,8 +257,9 @@ void run_attention_fwd_host(const TensorQ& q_g_m_k,
     ref_gemm0_invoker.Run(ref_gemm0_argument);
 
     // masking
+    auto M          = s_g_m_n.GetLengths()[1];
     auto N          = s_g_m_n.GetLengths()[2];
-    const auto mask = DeviceGemmInstanceFWD::C0MatrixMask(N);
+    const auto mask = DeviceGemmInstanceFWD::C0MatrixMask(M, N);
     s_g_m_n.ForEach([&](auto& self, auto idx) {
         if(mask.IsMaskedElement(idx[1], idx[2]))
             self(idx) = -ck::NumericLimits<float>::Infinity();

example/32_batched_gemm_scale_softmax_gemm/grouped_multihead_attention_train_v1.cpp

@@ -668,8 +668,9 @@ void run_attention_fwd_host(const TensorQ& q_g_m_k,
     ref_gemm0_invoker.Run(ref_gemm0_argument);
 
     // masking
+    auto M          = s_g_m_n.GetLengths()[1];
     auto N          = s_g_m_n.GetLengths()[2];
-    const auto mask = DeviceGemmInstanceFWD::C0MatrixMask(N);
+    const auto mask = DeviceGemmInstanceFWD::C0MatrixMask(M, N);
     s_g_m_n.ForEach([&](auto& self, auto idx) {
         if(mask.IsMaskedElement(idx[1], idx[2]))
             self(idx) = -ck::NumericLimits<float>::Infinity();

example/32_batched_gemm_scale_softmax_gemm/grouped_multihead_attention_train_v2.cpp

@@ -256,8 +256,9 @@ void run_attention_fwd_host(const TensorQ& q_g_m_k,
     ref_gemm0_invoker.Run(ref_gemm0_argument);
 
     // masking
+    auto M          = s_g_m_n.GetLengths()[1];
     auto N          = s_g_m_n.GetLengths()[2];
-    const auto mask = DeviceGemmInstanceFWD::C0MatrixMask(N);
+    const auto mask = DeviceGemmInstanceFWD::C0MatrixMask(M, N);
     s_g_m_n.ForEach([&](auto& self, auto idx) {
         if(mask.IsMaskedElement(idx[1], idx[2]))
             self(idx) = -ck::NumericLimits<float>::Infinity();

example/47_gemm_bias_softmax_gemm_permute/gemm_bias_softmax_gemm_permute.cpp

@@ -361,7 +361,7 @@ int main(int argc, char* argv[])
             c0de_element_op(acc0_g_m_n(idx), acc0_g_m_n(idx), d0_g_m_n(idx));
         });
         // masking
-        const auto mask = DeviceOpInstance::C0MatrixMask(N);
+        const auto mask = DeviceOpInstance::C0MatrixMask(M, N);
         acc0_g_m_n.ForEach([&](auto& self, auto idx) {
             if(mask.IsMaskedElement(idx[1], idx[2]))
                 self(idx) = -ck::NumericLimits<float>::Infinity();