[QST] Adding new parameter to Conv2dFprop in Python

What is your question?
Hi, I have implemented a modification to the conv2d kernel in Cutlass. I have added three new attributes (two tensors, mask0s and mask1s, with the same datatype as the filters) and a flag (integer) that can be used in the convolutional kernel: https://github.com/NVIDIA/cutlass/blob/24f991e87930e1159f1f5a47e329d43bcfbd76b9/include/cutlass/conv/kernel/implicit_gemm_convolution.h#L281.

To do so, I modified the cutlass code into three parts:

1. Inside Host Tensot:

  /// Host-side memory allocation
  std::vector<StorageUnit> host_;
  std::vector<StorageUnit> host_mask1s;
  std::vector<StorageUnit> host_mask0s;
  int* host_flag;

  /// Device-side memory
  device_memory::allocation<StorageUnit> device_;
  device_memory::allocation<StorageUnit> device_masks1s;
  device_memory::allocation<StorageUnit> device_masks0s;
  device_memory::allocation<int> device_flag;

 /// Accesses the tensor reference pointing to data
  TensorRef host_ref_complete(LongIndex ptr_element_offset=0) { return TensorRef(host_data_ptr_offset(ptr_element_offset), 
                                                                        layout_, 
                                                                        host_data_ptr_offset_mask0s(ptr_element_offset),
                                                                        host_data_ptr_offset_mask1s(ptr_element_offset),
                                                                        host_flag); }

  /// Accesses the tensor reference pointing to data
  ConstTensorRef device_ref_complete(LongIndex ptr_element_offset=0) const {
    return TensorRef(device_data_ptr_offset(ptr_element_offset), 
    layout_, device_data_ptr_offset_mask0s(ptr_element_offset), 
    device_data_ptr_offset_mask1s(ptr_element_offset), 
    device_data_flag());
  }


  /// Copies data from device to host
  void sync_host() {
    if (device_backed()) {
      device_memory::copy_to_host(
          host_.data(), device_.get(), device_.size());

      device_memory::copy_to_host(
          host_mask1s.data(), device_masks1s.get(), device_masks1s.size());
      
      device_memory::copy_to_host(
          host_mask0s.data(), device_masks0s.get(), device_masks0s.size());

      device_memory::copy_to_host(
          host_flag, device_flag.get(), device_flag.size());
      
    }
  }

  /// Copies data from host to device
  void sync_device() {
    if (device_backed()) {
      device_memory::copy_to_device(
          device_.get(), host_.data(), host_.size());
      
      device_memory::copy_to_device(
          device_masks1s.get(), host_mask1s.data(), host_mask1s.size());
      
      device_memory::copy_to_device(
          device_masks0s.get(), host_mask0s.data(), host_mask0s.size());

      device_memory::copy_to_device(
          device_flag.get(), host_flag, 1);
          
    }
  }

I have shown only the most critical functions globally, but inside the class, I modified constructors and reset methods to allocate memory from the host and later with copy_to_devidce() in the GPU.

2. In Tensor_ref:

/// Pointer
Element* ptr_;
Element* ptr_mask0s;
Element* ptr_mask1s;
int *check;

CUTLASS_HOST_DEVICE
TensorRef(
  Element *ptr,                   ///< pointer to start of 
  Layout const &layout,            ///< layout object containing stride and mapping function
  Element *ptr_mask0s,            ///< pointer to start of tensor mask0s
  Element *ptr_mask1s,            ///< pointer to start of  tensor mask1s
  int *device_flag              ///< pointer to check flag
):
  ptr_(ptr), layout_(layout), check(device_flag), ptr_mask0s(ptr_mask0s), ptr_mask1s(ptr_mask1s){ 
  }

/// Returns the pointer to referenced data
CUTLASS_HOST_DEVICE
Element * data() const { return ptr_; }

/// Returns the pointer to referenced data
CUTLASS_HOST_DEVICE
Element * data_mask0s() const { return ptr_mask0s; }

/// Returns the pointer to referenced data
CUTLASS_HOST_DEVICE
Element * data_mask1s() const { return ptr_mask1s; }

int * check_flag() const { return check; }

3. In the same conv kernel:

/// Parameters structure
  struct Params {
    ConvProblemSize problem_size;
    cutlass::gemm::GemmCoord grid_tiled_shape;
    gemm::GemmCoord implicit_gemm_problem_size;
    int swizzle_log_tile;
    int gemm_k_iterations;
    int gemm_k_iterations_per_channel;
    int *first_call;
    typename Mma::IteratorA::Params iterator_A;
    typename Mma::IteratorA::Element const *ptr_A;
    typename Mma::IteratorB::Params iterator_B;
    typename Mma::IteratorB::Element *ptr_B;
    typename Mma::IteratorB::Element *ptr_B_mask0s;
    typename Mma::IteratorB::Element *ptr_B_mask1s;
    typename Epilogue::OutputTileIterator::Params iterator_C;
    typename Epilogue::OutputTileIterator::Element *ptr_C;
    typename Epilogue::OutputTileIterator::Params iterator_D;
    typename Epilogue::OutputTileIterator::Element *ptr_D;
    typename EpilogueOutputOp::Params output_op;
    int *semaphore;
    SplitKMode split_k_mode;

    //
    // Methods
    //

    CUTLASS_HOST_DEVICE
    Params(): swizzle_log_tile(0), gemm_k_iterations(0) { }

    /// 
    CUTLASS_HOST_DEVICE
    Params(
      Arguments const &args,
      int *semaphore = nullptr
    ):
      problem_size(args.problem_size),
      implicit_gemm_problem_size(cutlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, args.problem_size)),
      iterator_A(Mma::IteratorA::getParams(args.problem_size, args.ref_A.layout())),
      ptr_A(args.ref_A.data()),
      iterator_B(args.problem_size, args.ref_B.layout()),
      ptr_B(args.ref_B.data()),
      iterator_C(ConvOutputIteratorParameter::layout(args.ref_C), implicit_gemm_tensor_c_extent(kConvolutionalOperator, args.problem_size)),
      ptr_C(args.ref_C.data()),
      iterator_D(ConvOutputIteratorParameter::layout(args.ref_D), implicit_gemm_tensor_c_extent(kConvolutionalOperator, args.problem_size)),
      ptr_D(args.ref_D.data()),
      output_op(args.output_op),
      first_call(args.ref_B.check_flag()),
      ptr_B_mask0s(args.ref_B.data_mask0s()),
      ptr_B_mask1s(args.ref_B.data_mask1s()),
      semaphore(semaphore),
      split_k_mode(args.split_k_mode)
    {
      gemm_k_iterations = implicit_gemm_k_iterations(
        kConvolutionalOperator,
        ThreadblockShape::kK,
        args.problem_size,
        kIteratorAlgorithm,
        kGroupMode,
        ThreadblockShape::kN);

      gemm_k_iterations_per_channel = implicit_gemm_k_iterations_per_channel(
          kConvolutionalOperator, args.problem_size, kIteratorAlgorithm);

      ThreadblockSwizzle threadblock_swizzle;
      printf("construct params on kernel/implicit Gemem lin256...\n");

      grid_tiled_shape = threadblock_swizzle.get_tiled_shape(
        implicit_gemm_problem_size,
        {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
        args.problem_size.split_k_slices);

      swizzle_log_tile = threadblock_swizzle.get_log_tile(grid_tiled_shape);
    }
  };

All of this code is currently working, and it does ok without problems or memory errors while executing example 16.

My question now is how to adapt the host code to Python. The device code works well and takes the same arguments as before; the only difference is that now the tensor_ref class can have three more attributes. However, when defining the convolution on python there is a part in which it is specified the arguments in https://github.com/NVIDIA/cutlass/blob/375e284e6aef68b81d58b116dff9e0970f64c5cd/python/cutlass/backend/conv2d_operation.py#L214

I tried to modify this part of the code by adding two ElementB pointers and an integer and using the new constructor of tenser_ref, which I have previously added:

struct ${operation_name}_TemporaryArgs {
    int conv_kind;
    cutlass::conv::Conv2dProblemSize problem_size;
    ElementA* ptr_A;
    ElementB* ptr_B;
    ElementB* ptr_B_0;
    ElementB* ptr_B_1;
    ElementC* ptr_C;
    ElementC* ptr_D;
    int tensor_c_numel;
    typename EpilogueOutputOp::Params epilogue_params;
    int split_k_mode;
  };

  typename ${operation_name}${operation_suffix}::Arguments
  construct_arguments(${operation_name}_TemporaryArgs args) {
    cutlass::conv::Operator conv_operator = static_cast<cutlass::conv::Operator>(args.conv_kind);
    auto tc_A = cutlass::conv::implicit_gemm_tensor_a_extent(conv_operator, args.problem_size);
    auto tc_B = cutlass::conv::implicit_gemm_tensor_b_extent(conv_operator, args.problem_size);
    auto tc_C = cutlass::conv::implicit_gemm_tensor_c_extent(conv_operator, args.problem_size);
    auto tc_D = cutlass::conv::implicit_gemm_tensor_c_extent(conv_operator, args.problem_size);

    auto size_C = tc_C.at(0) * tc_C.at(1) * tc_C.at(2) * tc_C.at(3);
    if (args.tensor_c_numel >= 0 && args.tensor_c_numel == tc_C.at(3) && args.tensor_c_numel < size_C) {
      // C is interpreted as bias
      tc_C = {0, 0, 0, 0};
    }


    int device_flag = 0;
    cutlass::TensorRef<ElementA, LayoutA> tref_A(args.ptr_A, LayoutA::packed(tc_A));
    cutlass::TensorRef<ElementB, LayoutA> tref_B(args.ptr_B, LayoutB::packed(tc_B), args.ptr_B_0, args.ptr_B_1, &device_flag);
    cutlass::TensorRef<ElementC, LayoutA> tref_C(args.ptr_C, LayoutC::packed(tc_C));
    cutlass::TensorRef<ElementC, LayoutA> tref_D(args.ptr_D, LayoutC::packed(tc_D));

    return {
      args.problem_size,
      tref_A,
      tref_B,
      tref_C,
      tref_D,
      args.epilogue_params,
      static_cast<cutlass::conv::SplitKMode>(args.split_k_mode)
    };
  }

But of course, it gives an error of cuda illegal address when executing sync() to the parameters because I have added two pointers and an integer that has not correctly copied to GPU memory:

Traceback (most recent call last):
  File "python/prova.py", line 81, in <module>
    plan.run(input, weight, tensor_C, output, stride, padding, dilation, alpha, beta, print_module=print_module)
  File "/mnt/beegfs/gap/izcagal@upvnet.upv.es/cutlass/python/cutlass/op/conv.py", line 939, in run
    return super().run(
  File "/mnt/beegfs/gap/izcagal@upvnet.upv.es/cutlass/python/cutlass/op/conv.py", line 890, in run
    arguments.sync()
  File "/mnt/beegfs/gap/izcagal@upvnet.upv.es/cutlass/python/cutlass/backend/conv2d_operation.py", line 191, in sync
    return super().sync()
  File "/mnt/beegfs/gap/izcagal@upvnet.upv.es/cutlass/python/cutlass/backend/arguments.py", line 107, in sync
    raise RuntimeError("CUDA Error %s" % str(err))
RuntimeError: CUDA Error cudaError_t.cudaErrorIllegalAddress

So, the question is, where do I need to add these pointers and the integer to properly initialize them on the host CPU memory and copy them to the GPU??

In the example 16 it is used host_tensor class, but in python, there is not such a class, so I am a little bit lost on this.

Any feedback would be appreciated.

Thanks.

Trace argument construction and synchronization through python/cutlass/backend/conv2d_operation.py and python/cutlass/backend/arguments.py, comparing the Python path with host_tensor.h, tensor_ref.h, and implicit_gemm_convolution.h. Start at arguments.sync() and the generated construct_arguments code; done means Conv2dFprop can construct and synchronize the added mask pointers and flag without cudaErrorIllegalAddress.