CUTLASS 3.5.1 (#1623)

* CUTLASS 3.5.1

* updates, optimizations, fixes

tools/library/src/conv_operation_3x.hpp

@@ -595,12 +595,10 @@ private:
 
       const TensorStride stride_A = vector_to_array_strides(config.stride_a, the_stride_size);
       const TensorStride stride_B = vector_to_array_strides(config.stride_b, the_stride_size);
-      const TensorStride stride_C = vector_to_array_strides(config.stride_c, the_stride_size);
 
       // cutlass::library::Conv2dConfiguration has no member stride_d.
       // The code below imitates the testbed,
       // which just sets D's strides to C's strides.
-      const TensorStride stride_D = stride_C;
 
       const int num_groups = config.problem_size.groups;
       if (num_groups != 1) {
@@ -773,9 +771,7 @@ private:
 
       const TensorStride stride_A = coord_to_array_strides(input_stride_a);
       const TensorStride stride_B = coord_to_array_strides(input_stride_b);
-      const TensorStride stride_C = coord_to_array_strides(input_stride_c);
 
-      const TensorStride stride_D = stride_C;
       const int num_groups = config.problem_size.groups;
       if (num_groups != 1) {
         CUTLASS_TRACE_HOST("CUTLASS 3 kernels currently only support groups = 1.");
@@ -852,6 +848,12 @@ private:
     std::cerr << "ConvOperation3x::update_operator_arguments_from_arguments\n";
 #endif
 
+    auto status = UpdateFusionArgs<decltype(out_args.epilogue.thread)>::update_(
+      out_args.epilogue.thread, in_args);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
     out_args.mainloop.ptr_A = reinterpret_cast<ElementA const*>(in_args.A);
     out_args.mainloop.ptr_B = reinterpret_cast<ElementB const*>(in_args.B);
 

tools/library/src/gemm_operation_3x.hpp

@@ -250,6 +250,10 @@ protected:
     /* Query device SM count to pass onto the kernel as an argument, where needed */
     operator_args.hw_info.sm_count = arguments->sm_count;
 
+    if constexpr (!std::is_const_v<decltype(operator_args.scheduler.max_swizzle_size)>) {
+      operator_args.scheduler.max_swizzle_size = arguments->swizzle_size;
+    }
+
     if constexpr (!std::is_const_v<decltype(operator_args.scheduler.raster_order)>) {
       using Enum_t = decltype(operator_args.scheduler.raster_order);
       switch (arguments->raster_order) {

tools/library/src/reference/conv_reference_operation.h

@@ -489,7 +489,7 @@ template <
   typename InnerProductOp_ = multiply_add<ElementAccumulator_>
 >
 void make_conv_fprop(Manifest &manifest) {
-  
+#if !defined(CUTLASS_PROFILER_DISABLE_REFERENCE)
   manifest.append(new ConvReferenceOperation<
     Provider::kReferenceHost,
     cutlass::conv::Operator::kFprop,
@@ -515,6 +515,7 @@ void make_conv_fprop(Manifest &manifest) {
     ConvertOp_,
     InnerProductOp_
   >);
+#endif // !defined(CUTLASS_PROFILER_DISABLE_REFERENCE)
 }
 
 /// Constructs Dgrad and Wgrad reference operators.
@@ -532,7 +533,7 @@ template <
   typename InnerProductOp_ = multiply_add<ElementAccumulator_>
 >
 void make_conv_backwards(Manifest &manifest) {
-  
+#if !defined(CUTLASS_PROFILER_DISABLE_REFERENCE)
   manifest.append(new ConvReferenceOperation<
     Provider::kReferenceHost,
     cutlass::conv::Operator::kDgrad,
@@ -584,6 +585,7 @@ void make_conv_backwards(Manifest &manifest) {
     ConvertOp_,
     InnerProductOp_
   >);
+#endif // !defined(CUTLASS_PROFILER_DISABLE_REFERENCE)
 }
 
 /// Six operators for the price of one.

tools/library/src/reference/gemm_reference_operation.h

@@ -293,7 +293,7 @@ template <
   typename InnerProductOp_ = multiply_add<ElementAccumulator_>
 >
 void make_gemm(Manifest &manifest) {
-  
+#if !defined(CUTLASS_PROFILER_DISABLE_REFERENCE)
   manifest.append(new GemmReferenceOperation<
     Provider::kReferenceHost,
     ElementA_, LayoutA_, TransformA,
@@ -317,6 +317,7 @@ void make_gemm(Manifest &manifest) {
     ConvertOp_,
     InnerProductOp_
   >);
+#endif
 }
 
 /// Helper to create NN, NT, TN, and TT GEMM layouts.