Allow scalar broadcasting in VisitorRowBroadcast and VisitorColBroadcast

include/cutlass/epilogue/threadblock/fusion/visitor_load.hpp

@@ -335,7 +335,8 @@ struct VisitorAuxLoad{
 template<
   class ThreadMap,
   class Element,
-  class StrideMNL
+  class StrideMNL,
+  bool EnableNullptr = false
 >
 struct VisitorRowBroadcast {
 
@@ -403,10 +404,31 @@ struct VisitorRowBroadcast {
       auto src_v = filter(tC_gRow);
       auto coord_v = filter(tC_cRow);
       auto dst_v = filter(tC_rRow);
+
+      if constexpr (EnableNullptr) {
+        if (params_ptr->ptr_row == nullptr) {
+          // In this case we are loading from a scalar and broadcasting
+          VecType filled_vec;
+          CUTLASS_PRAGMA_UNROLL
+          for (int i = 0; i < VecLength; i++) {
+            reinterpret_cast<Element *>(&filled_vec)[i] = params_ptr->null_default;
+          }
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int i = 0; i < size(src_v); ++i) {
+            if (get<1>(coord_v(i)) < n) {
+              dst_v(i) = filled_vec;
+            }
+          }
+          return;
+        }
+      }
+
+      // In this case we are loading from a row vector and broadcasting
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < size(src_v); ++i) {
         bool guard = get<1>(coord_v(i)) < n;
-        cutlass::arch::global_load<VecType, sizeof(VecType)>(dst_v(i), (void const*)&src_v(i), guard);
+        cutlass::arch::global_load<VecType, sizeof(VecType)>(dst_v(i), (void const *)&src_v(i), guard);
       }
     }
 
@@ -464,7 +486,8 @@ struct VisitorRowBroadcast {
 template<
   class ThreadMap,
   class Element,
-  class StrideMNL = Stride<_1,_0,_0>
+  class StrideMNL = Stride<_1,_0,_0>,
+  bool EnableNullptr = false
 >
 struct VisitorColBroadcast {
 
@@ -524,11 +547,29 @@ struct VisitorColBroadcast {
     CUTLASS_DEVICE void
     begin_epilogue() {
       clear(tC_rCol);
+
       Tensor pred = make_tensor<bool>(shape(tC_gCol));
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < size(pred); ++i) {
         pred(i) = get<0>(tC_cCol(i)) < m;
       }
+
+      if constexpr (EnableNullptr) {
+        if (params_ptr->ptr_col == nullptr) {
+          // In this case we are loading from a scalar and broadcasting
+          auto dst_v = filter(tC_rCol);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int i = 0; i < size(dst_v); ++i) {
+            if (pred(i)) {
+              dst_v(i) = params_ptr->null_default;
+            }
+          }
+          return;
+        }
+      }
+
+      // In this case we are loading from a column vector and broadcasting
       copy_if(pred, tC_gCol, tC_rCol);
     }