Layout Class — pytorch Architecture

Source Code

aten/src/ATen/native/transformers/cuda/mem_eff_attention/gemm/custom_mma_base.h lines 103–118

  template <typename Element, typename OperandShape, typename OperandLayout>
  struct OperandSharedStorage {
    AlignedBuffer<Element, OperandShape::kCount> buffer;
    using TensorRef = TensorRef<Element, OperandLayout>;

    CUTLASS_DEVICE
    static OperandLayout Layout() {
      return OperandLayout::packed({OperandShape::kRow, OperandShape::kColumn});
    }

    /// Returns a TensorRef to the operand
    CUTLASS_HOST_DEVICE
    TensorRef ref() {
      return TensorRef{buffer.data(), Layout()};
    }
  };