refactor(dnn): refactor lowbit tensor format

GitOrigin-RevId: b646dc085b

dnn/include/megdnn/basic_types.h

@@ -170,6 +170,7 @@ struct TensorLayout : public TensorShape {
 
 #if MEGDNN_CC_HOST
         Format();
+        Format(DType dtype);
 
         const ImplBase* impl() const { return m_impl; }
 
@@ -198,6 +199,9 @@ struct TensorLayout : public TensorShape {
         //! whether this is the default tensor format
         bool is_default() const;
 
+        //! whether this is the lowbit aligned to bytes tensor format
+        bool is_lowbit_aligned() const;
+
         bool operator==(Format rhs) const { return m_impl == rhs.m_impl; }
         bool operator!=(Format rhs) const { return m_impl != rhs.m_impl; }
 #endif

dnn/include/megdnn/tensor_format.h

@@ -20,7 +20,7 @@ namespace megdnn {
 enum class TensorFormat::Type {
     DEFAULT = 0,        //!< see DefaultTensorFormat
     IMAGE2D_PACK4 = 1,  //!< see Image2DPack4TensorFormat
-    FOURBITS_ALIGNED_TO_BYTE = 2, //!< 
+    LOWBITS_ALIGNED_TO_BYTE = 2, //!< 
 };
 
 class TensorFormat::ImplBase {
@@ -205,21 +205,23 @@ using Image2DPack4TensorFormatBase = Image2DPackedTensorFormatBase<4>;
 /*!
  * \brief used for tensors storing lowbit data 
  *
- * \p SIZE_NBITS is the size in bits of element of the tensor.
- *
+ * \param m_size_nbits size in bits of elements in the tensor
+ * \param m_align_size_in_bits aligned size in bits
+ * \param m_align_size_in_elements aligned size in elements
  */
-template <size_t SIZE_NBITS_>
-class LowbitsTensorFormatBase : public TensorFormat::ImplBase {
-    static constexpr size_t SIZE_NBITS = SIZE_NBITS_;
-    size_t m_align_size_in_bits, m_align_size_in_elements;
+class LowbitsAlignedTensorFormatBase : public TensorFormat::ImplBase {
+    size_t m_size_nbits, m_align_size_in_bits, m_align_size_in_elements;
 
 protected:  //?
-    LowbitsTensorFormatBase(Type type, size_t align_size_in_bits);
+    LowbitsAlignedTensorFormatBase(Type type, size_t size_nbits,
+                                   size_t align_size_in_bits);
 
-    virtual ~LowbitsTensorFormatBase() = default;
+    virtual ~LowbitsAlignedTensorFormatBase() = default;
 
 public:
     size_t align_size_in_bits() const { return m_align_size_in_bits; }
+    
+    size_t size_nbits() const { return m_size_nbits; }
 
     std::string to_string() const override;
 
@@ -240,10 +242,10 @@ public:
             const TensorLayout& layout) const override;
 protected:
     struct SerializePack {
+        uint8_t size_nbits;
         uint8_t align_size_in_bits;
     };
 };
-using FourBitsAlignedToBytesTensorFormatBase = LowbitsTensorFormatBase<4>;
 }  // namespace detail
 
 /*!
@@ -296,19 +298,20 @@ private:
  * \brief Tensor for storing 4bit data that requires stride corresponding to
  * non-innermost dimension to be aligned to bytes, and pack 2 elems into a byte
  */
-class FourBitsAlignedToBytesTensorFormat final
-        : public detail::FourBitsAlignedToBytesTensorFormatBase {
+class LowbitsAlignedToBytesTensorFormat final
+        : public detail::LowbitsAlignedTensorFormatBase {
 public:
-    static constexpr Type TYPE = Type::FOURBITS_ALIGNED_TO_BYTE;
+    static constexpr Type TYPE = Type::LOWBITS_ALIGNED_TO_BYTE;
+    static constexpr size_t BYTE_IN_BITS = 8;
 
-    static TensorFormat make(size_t align_size_in_bits);
+    static TensorFormat make(size_t size_nbits);
 
     static TensorFormat deserialize(const Handle* handle, const void* buf,
                                     size_t size);
 
     static bool is_valid_layout(const TensorLayout& layout) {
         if (layout.format.type() == TYPE) {
-            layout.format.as_impl<FourBitsAlignedToBytesTensorFormat>()
+            layout.format.as_impl<LowbitsAlignedToBytesTensorFormat>()
                     .assert_valid(layout);
             return true;
         }
@@ -316,9 +319,9 @@ public:
     }
 
 private:
-    FourBitsAlignedToBytesTensorFormat(size_t align_size_in_bits)
-            : detail::FourBitsAlignedToBytesTensorFormatBase(
-                      TYPE, align_size_in_bits) {}
+    LowbitsAlignedToBytesTensorFormat(size_t size_nbits)
+            : detail::LowbitsAlignedTensorFormatBase(TYPE, size_nbits,
+                                                     BYTE_IN_BITS) {}
 };
 }  // namespace megdnn
 

dnn/src/common/basic_types.cpp

@@ -195,21 +195,14 @@ bool TensorShape::is_empty() const {
 /* ===================== TensorLayout =====================  */
 TensorLayout::TensorLayout() = default;
 
-TensorLayout::TensorLayout(DType dtype_) : dtype{dtype_} {}
+TensorLayout::TensorLayout(DType dtype_)
+        : dtype{dtype_}, format{Format(dtype)} {}
 
 TensorLayout::TensorLayout(DType dtype_, Format format_)
         : dtype{dtype_}, format{format_} {}
 
 TensorLayout::TensorLayout(const TensorShape& shape, DType dtype)
-        : TensorShape(shape), dtype{dtype} {
-    if (dtype.low_bit() == 4_z) {
-        format = FourBitsAlignedToBytesTensorFormat::make(8_z);
-    } else {
-        megdnn_assert(!dtype.is_low_bit(), "Unsupported data type(%s)",
-                      dtype.name());
-        format = DefaultTensorFormat::make();
-    }
-}
+        : TensorLayout(shape, dtype, Format(dtype)) {}
 
 TensorLayout::TensorLayout(const TensorShape& shape, DType dtype,
                            TensorFormat format_)

dnn/src/common/convolution.cpp

@@ -722,7 +722,7 @@ ConvolutionBase<Parameter>::deduce_layout_fwd(const TensorLayout& src,
             megdnn_assert(src.ndim == 5 &&
                                   (filter.ndim == 5 || filter.ndim == 6) &&
                                   src[src.ndim - 1] == 64 &&
-                                  filter[filter.ndim - 1] == 4,
+                                  filter[filter.ndim - 1] == 64,
                           "NCHW64 require src and filter's ndim is 5 or 6, and "
                           "last shape is 64 but got src %s, filter %s",
                           src.to_string().c_str(), filter.to_string().c_str());
@@ -754,7 +754,6 @@ ConvolutionBase<Parameter>::deduce_layout_fwd(const TensorLayout& src,
                     src[i + src_or_dst_spatial_start], cflt.dilated_spatial[i],
                     cflt.stride[i], cflt.padding[i]);
         }
-        dst.init_contiguous_stride();
     } else if (param().format == Param::Format::NCHW4) {
         megdnn_assert(src.ndim == 5,
                       "invalid src ndim for NCHW4, expected=5, got=%zu",

dnn/src/common/tensor_format.cpp

@@ -35,8 +35,8 @@ TensorFormat TensorFormat::deserialize(const std::string& bin,
         case Type::IMAGE2D_PACK4:
             return Image2DPack4TensorFormat::deserialize(
                     handle, type + 1, bin.size() - sizeof(Type));
-        case Type::FOURBITS_ALIGNED_TO_BYTE:
-            return FourBitsAlignedToBytesTensorFormat::deserialize(
+        case Type::LOWBITS_ALIGNED_TO_BYTE:
+            return LowbitsAlignedToBytesTensorFormat::deserialize(
                     handle, type + 1, bin.size() - sizeof(Type));
         default:
             megdnn_throw("invalid tensor format type in deserialize");
@@ -45,6 +45,19 @@ TensorFormat TensorFormat::deserialize(const std::string& bin,
 
 TensorFormat::Format() : m_impl{DefaultTensorFormat::make().m_impl} {}
 
+TensorFormat::Format(DType dtype) {
+    megdnn_assert(dtype.valid());
+    if (dtype.is_low_bit()) {
+        size_t size_nbits = dtype.low_bit();
+        megdnn_assert(size_nbits == 1 || size_nbits == 2 || size_nbits == 4,
+                      "unsupported lowbits data type(%s, size in bits: %zu)",
+                      dtype.name(), size_nbits);
+        m_impl = LowbitsAlignedToBytesTensorFormat::make(size_nbits).m_impl;
+    } else {
+        m_impl = DefaultTensorFormat::make().m_impl;
+    }
+}
+
 std::string TensorFormat::to_string() const {
     return m_impl->to_string();
 }
@@ -69,6 +82,10 @@ bool TensorFormat::is_default() const {
     return m_impl == default_tensor_format_obj;
 }
 
+bool TensorFormat::is_lowbit_aligned() const {
+    return type() == TensorFormat::Type::LOWBITS_ALIGNED_TO_BYTE;
+}
+
 /* ===================== DefaultFormat ===================== */
 void DefaultTensorFormat::assert_valid(const TensorLayout& layout) const {
     megdnn_assert(
@@ -440,27 +457,26 @@ template class Image2DPackedTensorFormatBase<4>;
 }  // namespace detail
 }  // namespace megdnn
 
-/* =============== FourBitsAlignedToBytesTensorFormatBase ============== */
-template <size_t SIZE_NBITS>
-LowbitsTensorFormatBase<SIZE_NBITS>::LowbitsTensorFormatBase(
-        Type type, size_t align_size_in_bits)
-        : ImplBase(type), m_align_size_in_bits(align_size_in_bits) {
-    megdnn_assert(!(m_align_size_in_bits % SIZE_NBITS),
+/* =============== LowbitsAlignedTensorFormatBase ============== */
+LowbitsAlignedTensorFormatBase::LowbitsAlignedTensorFormatBase(
+        Type type, size_t size_nbits, size_t align_size_in_bits)
+        : ImplBase(type),
+          m_size_nbits(size_nbits),
+          m_align_size_in_bits(align_size_in_bits) {
+    megdnn_assert(!(m_align_size_in_bits % m_size_nbits),
                   "align size(%zu) must be a multiple of element size(%zu)",
-                  m_align_size_in_bits, SIZE_NBITS);
-    m_align_size_in_elements = m_align_size_in_bits / SIZE_NBITS;
+                  m_align_size_in_bits, m_size_nbits);
+    m_align_size_in_elements = m_align_size_in_bits / m_size_nbits;
 }
 
-template <size_t SIZE_NBITS>
-std::string LowbitsTensorFormatBase<SIZE_NBITS>::to_string() const {
-    return ssprintf("LOWBITS{%zu,%zu}", SIZE_NBITS, m_align_size_in_bits);
+std::string LowbitsAlignedTensorFormatBase::to_string() const {
+    return ssprintf("LOWBITS{%zu,%zu}", m_size_nbits, m_align_size_in_bits);
 }
 
-template <size_t SIZE_NBITS>
-void LowbitsTensorFormatBase<SIZE_NBITS>::assert_valid(
+void LowbitsAlignedTensorFormatBase::assert_valid(
         const TensorLayout& layout) const {
     megdnn_assert(layout.dtype.valid() && layout.dtype.is_low_bit() &&
-                  layout.dtype.low_bit() == SIZE_NBITS);
+                  layout.dtype.low_bit() == m_size_nbits);
     bool has_dim_unity_stride = false;
     for (int i = layout.ndim - 1; i >= 0; --i) {
         if (!has_dim_unity_stride && layout.stride[i] == 1)
@@ -469,23 +485,28 @@ void LowbitsTensorFormatBase<SIZE_NBITS>::assert_valid(
                 layout.stride[i] >= 0 &&
                         (layout.stride[i] % m_align_size_in_elements == 0 ||
                          layout.stride[i] == 1),
-                "bad stride: %zu", layout.stride[i]);
+                "bad stride:%s, %zu", layout.to_string().c_str(),
+                layout.stride[i]);
     }
-    megdnn_assert(has_dim_unity_stride, "innermost dim not contiguous");
+    /// FIXME
+    if (layout.ndim == 0) {
+        printf("%s\n", layout.to_string().c_str());
+    }
+    megdnn_assert(layout.ndim == 0 || has_dim_unity_stride,
+                  "innermost dim not contiguous");
 }
 
-template <size_t SIZE_NBITS>
-void LowbitsTensorFormatBase<SIZE_NBITS>::serialize_append(
+void LowbitsAlignedTensorFormatBase::serialize_append(
         std::string& result) const {
     SerializePack pack;
+    pack.size_nbits = m_size_nbits;
     pack.align_size_in_bits = m_align_size_in_bits;
     megdnn_assert(pack.align_size_in_bits ==
                   m_align_size_in_bits);  // detect overflow;
     result.append(reinterpret_cast<char*>(&pack), sizeof(pack));
 }
 
-template <size_t SIZE_NBITS>
-TensorLayout::Span LowbitsTensorFormatBase<SIZE_NBITS>::span_spec(
+TensorLayout::Span LowbitsAlignedTensorFormatBase::span_spec(
         const TensorLayout& layout) const {
     assert_valid(layout);
     if (layout.ndim == 0)
@@ -507,8 +528,7 @@ TensorLayout::Span LowbitsTensorFormatBase<SIZE_NBITS>::span_spec(
     return TensorLayout::Span(0, 0, high_elem, high_byte);
 }
 
-template <size_t SIZE_NBITS>
-size_t LowbitsTensorFormatBase<SIZE_NBITS>::init_contiguous_stride(
+size_t LowbitsAlignedTensorFormatBase::init_contiguous_stride(
         TensorLayout& layout) const {
     if (!layout.ndim)
         return 0;
@@ -525,8 +545,7 @@ size_t LowbitsTensorFormatBase<SIZE_NBITS>::init_contiguous_stride(
     return accum;
 }
 
-template <size_t SIZE_NBITS>
-bool LowbitsTensorFormatBase<SIZE_NBITS>::is_contiguous_spec(
+bool LowbitsAlignedTensorFormatBase::is_contiguous_spec(
         const TensorLayout& layout) const {
     assert_valid(layout);
     ptrdiff_t expected = 1;
@@ -541,8 +560,7 @@ bool LowbitsTensorFormatBase<SIZE_NBITS>::is_contiguous_spec(
     return expected != 0;
 }
 
-template <size_t SIZE_NBITS>
-TensorLayout LowbitsTensorFormatBase<SIZE_NBITS>::collapse_contiguous_spec(
+TensorLayout LowbitsAlignedTensorFormatBase::collapse_contiguous_spec(
         const TensorLayout& layout) const {
     assert_valid(layout);
     TensorLayout res{layout};
@@ -572,12 +590,6 @@ TensorLayout LowbitsTensorFormatBase<SIZE_NBITS>::collapse_contiguous_spec(
     return res;
 }
 
-namespace megdnn {
-namespace detail {
-template class LowbitsTensorFormatBase<4>;
-}  // namespace detail
-}  // namespace megdnn
-
 /* ===================== Image2DPack4TensorFormat  ===================== */
 TensorFormat Image2DPack4TensorFormat::make_raw(
         size_t align_axis, size_t align_size_in_elements,
@@ -616,29 +628,28 @@ TensorFormat Image2DPack4TensorFormat::change_axis(size_t axis) const {
     return make_raw(axis, align_size_in_elements(), vendor());
 }
 
-/* ===================== FourBitsAlignedToBytesTensorFormat
+/* ===================== LowbitsitsAlignedToBytesTensorFormat
  * ===================== */
-TensorFormat FourBitsAlignedToBytesTensorFormat::make(
-        size_t align_size_in_bits) {
+TensorFormat LowbitsAlignedToBytesTensorFormat::make(size_t size_nbits) {
     static std::mutex mtx;
     static std::unordered_map<
-            uint32_t, std::unique_ptr<FourBitsAlignedToBytesTensorFormat>>
+            uint64_t, std::unique_ptr<LowbitsAlignedToBytesTensorFormat>>
             cache;
-    megdnn_assert(!(align_size_in_bits % 4));
+    megdnn_assert(!(8 % size_nbits));
     MEGDNN_LOCK_GUARD(mtx);
-    auto&& ptr = cache[static_cast<uint32_t>(align_size_in_bits)];
+    auto&& ptr = cache[static_cast<uint32_t>(size_nbits)];
     if (!ptr) {
-        ptr.reset(new FourBitsAlignedToBytesTensorFormat{align_size_in_bits});
+        ptr.reset(new LowbitsAlignedToBytesTensorFormat{size_nbits});
     }
     return impl_to_tensor_format(ptr.get());
 }
 
-TensorFormat FourBitsAlignedToBytesTensorFormat::deserialize(const Handle*,
-                                                             const void* buf,
-                                                             size_t size) {
+TensorFormat LowbitsAlignedToBytesTensorFormat::deserialize(const Handle*,
+                                                            const void* buf,
+                                                            size_t size) {
     megdnn_assert(size == sizeof(SerializePack));
     auto pack = *static_cast<const SerializePack*>(buf);
-    return make(pack.align_size_in_bits);
+    return make(pack.size_nbits);
 }
 
 // vim: syntax=cpp.doxygen

dnn/src/cuda/conv_bias/cudnn_conv_bias_activation.cpp

@@ -24,6 +24,9 @@ using namespace conv_bias;
 
 bool ConvBiasForwardImpl::AlgoCUDNNConvBiasActivation::is_available(
         const SizeArgs& args) const {
+    if (args.src_layout->dtype.enumv() == DTypeEnum::QuantizedS4 &&
+        args.filter_layout->dtype.enumv() == DTypeEnum::QuantizedS4)
+        return false;
     if (args.src_layout->dtype == args.filter_layout->dtype &&
         args.src_layout->dtype == dtype::BFloat16()) {
         return false;

dnn/src/cuda/conv_bias/fallback_nchw_qs4.cpp

@@ -103,15 +103,18 @@ void ConvBiasForwardImpl::AlgoFallbackNCHWQS4::exec(
     TensorND src_{ws_src, layouts[0]}, filter_{ws_filter, layouts[1]},
             bias_{args.bias_tensor->raw_ptr, layouts[2]}, z_{ws_z, layouts[3]},
             dst_{ws_dst, layouts[4]};
-    ExecArgs args_{args.opr,
+    auto conv_op = args.opr->handle()->create_operator<ConvBiasForward>();
+    conv_op->param() = args.opr->param();
+    using Format = param::ConvBias::Format;
+    conv_op->param().format = Format::NCHW64;
+    ExecArgs args_{dynamic_cast<ConvBiasForwardImpl*>(conv_op.get()),
                    src_,
                    filter_,
                    bias_,
                    z_,
                    dst_,
-                   ws.get_workspace(3),
-                   args.preprocessed_filter};
-    m_underlying_algo.exec(args);
+                   ws.get_workspace(3)};
+    m_underlying_algo.exec(args_);
     // reformat dst
     nchw642nchw(dst_, {args.dst_tensor->raw_ptr, args.dst_tensor->layout});
 }
@@ -134,6 +137,9 @@ ConvBiasForwardImpl::AlgoFallbackNCHWQS4::make_underlying_tensor_layout(
         rst.emplace_back(TensorLayout{});
     }
     rst.emplace_back(TensorLayout{{n, co / 64, ho, wo, 64}, dst.dtype});
+    for (auto& i : rst) {
+        i.init_contiguous_stride();
+    }
     return rst;
 }
 
@@ -145,13 +151,16 @@ WorkspaceBundle ConvBiasForwardImpl::AlgoFallbackNCHWQS4::get_workspace_bundle(
     auto layouts = make_underlying_tensor_layout(
             *(args.src_layout), *(args.filter_layout), *(args.bias_layout),
             *(args.z_layout), *(args.dst_layout));
-    SizeArgs args_{args.opr,
+    auto conv_op = args.opr->handle()->create_operator<ConvBiasForward>();
+    conv_op->param() = args.opr->param();
+    using Format = param::ConvBias::Format;
+    conv_op->param().format = Format::NCHW64;
+    SizeArgs args_{dynamic_cast<ConvBiasForwardImpl*>(conv_op.get()),
                    layouts[0],
                    layouts[1],
                    layouts[2],
                    layouts[3],
-                   layouts[4],
-                   args.preprocessed_filter};
+                   layouts[4]};
     size_t ws_size_underlying_algo =
             m_underlying_algo.get_workspace_in_bytes(args_);
     if (args.z_layout->ndim > 0) {

dnn/src/cuda/conv_bias/helper.cpp

@@ -136,6 +136,10 @@ void ConvBiasDesc::set_conv(DType data_type, const param::ConvBias& param,
 namespace conv_bias {
 
 bool is_cudnn_supported(const BiasForwardSizeArgs& args) {
+    if (args.src_layout->dtype.enumv() == DTypeEnum::QuantizedS4 &&
+        args.filter_layout->dtype.enumv() == DTypeEnum::QuantizedS4)
+        return false;
+
     if (args.src_layout->dtype == args.filter_layout->dtype &&
         args.src_layout->dtype == dtype::BFloat16()) {
         return false;

dnn/src/cuda/conv_bias/sass_implicit_gemm_int4_nchw64_imma.cpp

@@ -72,11 +72,11 @@ std::string ConvBiasForwardImpl::AlgoSASSInt4NCHW64IMMAImplicitGemm::kernel_key(
     auto&& param = args.opr->param();
     if (args.z_layout->ndim > 0) {
         kernel_key =
-                ssprintf("%s_conv_bias_int4_fuse_z_imma_ldg16_%ux%u",
+                ssprintf("%s_conv_bias_int4_fuse_z_imma8832_ldg16_%ux%u",
                          current_device_arch_name(), m_tile_nhw, m_tile_oc);
     } else {
         kernel_key =
-                ssprintf("%s_conv_bias_int4_imma_ldg16_%ux%u",
+                ssprintf("%s_conv_bias_int4_imma8832_ldg16_%ux%u",
                          current_device_arch_name(), m_tile_nhw, m_tile_oc);
     }
     if (param.nonlineMode == NonlineMode::H_SWISH) {
@@ -170,7 +170,7 @@ void ConvBiasForwardImpl::AlgoSASSInt4NCHW64IMMAImplicitGemm::exec(
         reorder_imma_filter_bias<4, 64>(
                 reinterpret_cast<int8_t*>(filter_ptr),
                 reinterpret_cast<int32_t*>(bias_ptr),
-                args.filter_tensor->compatible_ptr<int8_t>(),
+                reinterpret_cast<int8_t*>(args.filter_tensor->raw_ptr), 
                 args.bias_tensor->compatible_ptr<int32_t>(), co, ci, fh, fw,
                 stream);
     }
@@ -292,9 +292,10 @@ void ConvBiasForwardImpl::AlgoSASSInt4NCHW64IMMAImplicitGemm::exec_preprocess(
                                   param);
     auto&& stream = cuda_stream(args.opr->handle());
     reorder_imma_filter_bias<4, 64>(
-            args.preprocessed_filter->tensors[0].compatible_ptr<int8_t>(),
+            reinterpret_cast<int8_t*>(
+                    args.preprocessed_filter->tensors[0].raw_ptr),
             args.preprocessed_filter->tensors[1].compatible_ptr<int32_t>(),
-            args.filter_tensor->compatible_ptr<int8_t>(),
+            reinterpret_cast<int8_t*>(args.filter_tensor->raw_ptr),
             args.bias_tensor->compatible_ptr<int32_t>(), co, ci, fh, fw,
             stream);
 }

dnn/test/common/test_basic_types.cpp

@@ -320,7 +320,7 @@ TEST(BASIC_TYPES, TENSOR_LAYOUT_FMT_LOW_BITS) {
 
     layout = make_layout({16, 32, 7, 7}, {1792, 56, 8, 1},
                             dtype::QuantizedS4{1.3f});
-    layout.format = FourBitsAlignedToBytesTensorFormat::make(8_z);
+    layout.format = LowbitsAlignedToBytesTensorFormat::make(4_z);
     EXPECT_TRUE(layout.is_contiguous());
 
     layout = TensorLayout{{1, 32, 1, 1}, dtype::QuantizedS4{1.2f}};
@@ -339,12 +339,10 @@ TEST(BASIC_TYPES, TENSOR_LAYOUT_FMT_LOW_BITS_VALID) {
                               DefaultTensorFormat::make()),
                  MegDNNError);
     ASSERT_THROW(TensorLayout({1, 32, 1, 1}, dtype::QuantizedS32{1.2f},
-                              FourBitsAlignedToBytesTensorFormat::make(8_z))
-                         .span(),
+                              LowbitsAlignedToBytesTensorFormat::make(4_z)),
                  MegDNNError);
     ASSERT_THROW(TensorLayout({16, 32, 7, 7}, dtype::IntB2{},
-                              FourBitsAlignedToBytesTensorFormat::make(8_z))
-                         .span(),
+                              LowbitsAlignedToBytesTensorFormat::make(2_z)),
                  MegDNNError);
 }
 

src/core/impl/graph/operator_node.cpp

@@ -338,21 +338,26 @@ void OperatorNodeBase::init_output_format() {
     TensorFormat format, default_;
     for (auto i : input()) {
         auto cur = i->format();
-        if (cur != default_) {
+        if (!cur.is_default() && !cur.is_lowbit_aligned()) {
             if (format == default_) {
                 format = cur;
             } else {
                 mgb_assert(format == cur,
-                           "multiple non-default formats in inputs: %s vs %s",
+                           "multiple non-default or non-lowbits aligned "
+                           "formats in inputs: %s vs %s",
                            format.to_string().c_str(), cur.to_string().c_str());
             }
         }
     }
     for (auto i : output()) {
         if (i->contain_flag(VarNode::Flag::VOLATILE_CONTENT)) {
-            i->format(default_);
+            mgb_assert(format.is_default());
+            i->format(TensorFormat(i->dtype()));
         } else {
-            i->format(format);
+            if (!format.is_default())
+                i->format(format);
+            else
+                i->format(TensorFormat(i->dtype()));
         }
     }
 }

src/core/impl/graph/var_node_mem_mgr.cpp

@@ -1063,15 +1063,22 @@ bool VarNodeMemManager::fwd_in2out_readonly(
         return false;
     }
 
-    mgb_assert(
-            src != dest &&
-            src->comp_node().mem_node() == dest->comp_node().mem_node() &&
-            dest->m_mem_plan.valid() && src->m_mem_plan.valid() &&
-            dest->m_mem_plan.layout().eq_shape(sub.layout()) &&
-            dest->m_mem_plan.layout().dtype.size() == sub.layout().dtype.size()
-            );
-    assert_in_mem_opt_phase(
-            SeqMemOptimizer::Status::ALLOW_FWD_IN2OUT_READONLY);
+    bool cond_low_bit = dest->m_mem_plan.layout().dtype.is_low_bit() &&
+                        sub.layout().dtype.is_low_bit() &&
+                        dest->m_mem_plan.layout().dtype.low_bit() ==
+                                sub.layout().dtype.low_bit();
+    bool cond_normal =
+            !dest->m_mem_plan.layout().dtype.is_low_bit() &&
+            !sub.layout().dtype.is_low_bit() &&
+            dest->m_mem_plan.layout().dtype.size() == sub.layout().dtype.size();
+    MGB_MARK_USED_VAR(cond_low_bit);
+    MGB_MARK_USED_VAR(cond_normal);
+    mgb_assert(src != dest &&
+               src->comp_node().mem_node() == dest->comp_node().mem_node() &&
+               dest->m_mem_plan.valid() && src->m_mem_plan.valid() &&
+               dest->m_mem_plan.layout().eq_shape(sub.layout()) &&
+               (cond_normal || cond_low_bit));
+    assert_in_mem_opt_phase(SeqMemOptimizer::Status::ALLOW_FWD_IN2OUT_READONLY);
 
     if (!m_owner_graph->options().seq_opt.enable_mem_plan_opt)
         return false;

src/core/impl/tensor.cpp

@@ -443,8 +443,8 @@ TensorND<TensorStorage>::name
 
 DEF(resize, &)(const TensorShape& shape) {
     mgb_assert(m_layout.dtype.valid());
-    auto nr_elems = m_layout.init_contiguous_stride(shape);
-    m_storage.ensure_size(m_layout.dtype.size(nr_elems));
+    m_layout = TensorLayout(shape, m_layout.dtype);
+    m_storage.ensure_size(m_layout.span().dist_byte());
     return static_cast<ChainReturnType&>(*this);
 }
 
@@ -584,15 +584,19 @@ TensorND<TensorStorage>::copy_from(const TensorND<RStorage> &src) {
         m_layout.dtype.assert_is(src.dtype());
     else
         m_layout.dtype = src.dtype();
-    m_layout.format = {};
 
-    size_t size_bytes = dtype().size(
-            m_layout.init_contiguous_stride(src.shape()));
+    m_layout = TensorLayout(src.shape(), m_layout.dtype);
+    size_t size_bytes = m_layout.span().dist_byte();
     m_storage.ensure_size(size_bytes);
     if (!size_bytes) {
         return static_cast<ChainReturnType&>(*this);
     }
-    if (src.layout().is_physical_contiguous()) {
+    // requirement:
+    // default case, physical contiguous
+    // lowbit aligned, logical contiguous
+    if (src.layout().is_physical_contiguous() ||
+        (src.layout().format.is_lowbit_aligned() &&
+         src.layout().is_contiguous())) {
         if (should_check_overlap(*this, src)) {
             check_overlapped(m_storage.ptr(),
                              m_storage.ptr() + size_bytes,
@@ -635,10 +639,17 @@ TensorND<TensorStorage>::copy_from_fixlayout(
                          src.raw_ptr() + src_span.high_byte);
     }
 
-    bool self_contig = m_layout.is_physical_contiguous(),
-         src_contig = src.layout().is_physical_contiguous();
+    bool self_contig = m_layout.is_physical_contiguous() ||
+                       (m_layout.format.is_lowbit_aligned() &&
+                        m_layout.is_contiguous()),
+         src_contig = src.layout().is_physical_contiguous() ||
+                      (m_layout.format.is_lowbit_aligned() &&
+                       m_layout.is_contiguous());
     if (self_contig && src_contig) {
-        if (m_layout.format.is_default() && src.layout().format.is_default()) {
+        if ((m_layout.format.is_default() &&
+             src.layout().format.is_default()) ||
+            (m_layout.format.is_lowbit_aligned() &&
+             src.layout().format.is_lowbit_aligned())) {
             mgb_assert(src_span.low_byte == 0 && dst_span.low_byte == 0 &&
                        src_span.high_byte == dst_span.high_byte);
             m_storage.copy_from(src.storage(), src_span.high_byte);

src/core/impl/utils/persistent_cache.cpp

@@ -261,7 +261,8 @@ PersistentCache::Blob AlgoChooserProfileCache::Key::build_blob() const {
         ret.push_back(';');
         ret.append(ly.dtype.name());
         ret.push_back('|');
-        mgb_assert(ly.format.is_default(),
+        mgb_assert(ly.format.is_default() || (ly.format.is_lowbit_aligned() &&
+                                              ly.dtype.is_low_bit()),
                    "currently only default format is supported");
     }
     if (m_param_size) {

src/core/include/megbrain/tensor.h

@@ -68,7 +68,10 @@ class SubTensorSpec {
 
         //! get offset measured in bytes
         ptrdiff_t offset_byte() const {
-            return m_offset_elem * m_layout.dtype.size();
+            //! for lowbit cases, offset must aligned to bytes
+            mgb_assert(!m_layout.dtype.is_low_bit() ||
+                       !(m_offset_elem * m_layout.dtype.low_bit() % 8));
+            return m_layout.dtype.size(m_offset_elem);
         }
 
         /*!

src/gopt/impl/inference.cpp

@@ -554,14 +554,16 @@ void ParamFusePass::apply(OptState &state) const {
 
         SymbolVar new_var;
         bool is_default_format = var->format().is_default();
-        if (cg::is_static_var_value(var) && is_default_format) {
+        bool is_lowbit_aligned = var->format().is_lowbit_aligned();
+        if (cg::is_static_var_value(var) &&
+            (is_default_format || is_lowbit_aligned)) {
             // use ImmutableTensor for inferable vars
             HostTensorND hv;
             hv.copy_from(*inferred_val).sync();
             new_var = opr::ImmutableTensor::make(
                     *var->owner_graph(), hv, var_namer.name(var));
         } else {
-            if (is_default_format) {
+            if (is_default_format || is_lowbit_aligned) {
                 new_var = opr::SharedDeviceTensor::make_const(
                         *var->owner_graph(), inferred_val, var_namer.name(var));
             } else {

src/opr/impl/basic_arith.cpp

@@ -814,8 +814,13 @@ MGB_IMPL_OPR_GRAD(TypeCvt) {
 #endif
 
 void TypeCvt::mem_plan_fwd_in2out_writable() {
-    if (input(0)->dtype().size() == output(0)->dtype().size() &&
-            input(0)->layout().is_contiguous()) {
+    bool cond_low_bit =
+            input(0)->dtype().is_low_bit() && output(0)->dtype().is_low_bit() &&
+            input(0)->dtype().low_bit() == output(0)->dtype().low_bit();
+    bool cond_normal = !input(0)->dtype().is_low_bit() &&
+                       !output(0)->dtype().is_low_bit() &&
+                       input(0)->dtype().size() == output(0)->dtype().size();
+    if ((cond_low_bit || cond_normal) && input(0)->layout().is_contiguous()) {
         output(0)->set_fwd_in2out_writable(input(0));
     }
 }

src/opr/impl/io.cpp

@@ -120,12 +120,11 @@ public:
     explicit DevValueExecDep(DeviceTensorStorage val) : m_val{std::move(val)} {}
 };
 
-
 void intl::DeviceTensorHolder::init_output_format() {
     auto format = get_dev_tensor().format();
-    mgb_assert(format.is_default(), "non-default tensor format: %s",
-               format.to_string().c_str());
-    // no need to set output foramt since it is initialized as default
+    mgb_assert(format.is_default() || format.is_lowbit_aligned(),
+               "invalid tensor format: %s", format.to_string().c_str());
+    output(0)->format(format);
 }
 
 void intl::DeviceTensorHolder::init_output_mem_plan(bool dynamic) {

src/opr/impl/search_policy/algo_chooser.cpp

@@ -638,10 +638,18 @@ AlgoChooser<Opr>::AlgoChooserHelper::profile_single_algo(
     param.workspace = get_workspace_size_bytes(policy);
     for (int i = 0; i < arity; ++i) {
         auto&& src = m_layouts[i];
-        mgb_assert(src.format.is_default() &&
+        bool cond_normal = src.format.is_default() &&
                            (src.dtype.category() == DTypeCategory::FLOAT ||
                             src.dtype.category() == DTypeCategory::INT ||
-                            src.dtype.category() == DTypeCategory::QUANTIZED),
+                            src.dtype.category() == DTypeCategory::QUANTIZED);
+
+        bool cond_low_bit = src.dtype.is_low_bit() &&
+                            src.format.is_lowbit_aligned() &&
+                            (src.dtype.category() == DTypeCategory::QUANTIZED ||
+                             src.dtype.category() == DTypeCategory::LOWBIT);
+        MGB_MARK_USED_VAR(cond_normal);
+        MGB_MARK_USED_VAR(cond_low_bit);
+        mgb_assert(cond_normal || cond_low_bit,
                    "unsupported layout in profiling: %s",
                    src.to_string().c_str());
         param.dtypes[i] = src.dtype.enumv();

src/opr/impl/search_policy/profiler.cpp

@@ -175,15 +175,17 @@ typename TimedProfiler<Opr>::TResult TimedProfiler<Opr>::prof_impl(
     case DTypeTrait<_dt>::enumv: \
         return _dt(1.0f, static_cast<uint8_t>(0))
             cb(dtype::Quantized8Asymm);
+            cb(dtype::Quantized4Asymm);
 #undef cb
 
 #define cb(_dt)                  \
     case DTypeTrait<_dt>::enumv: \
         return _dt(1.0f)
-
+            
             cb(dtype::QuantizedS8);
             cb(dtype::QuantizedS16);
             cb(dtype::QuantizedS32);
+            cb(dtype::QuantizedS4);        
             default:
                 return DType::from_enum(enumv);
 #undef cb

src/opr/test/dnn/convolution.cpp

@@ -2603,4 +2603,306 @@ TEST_F(TestNoWeightPreprocess, NoPreprocess) {
 
 #endif
 
+namespace {
+// FIXME change comp node from "cpu0" to "gpu0"    
+TEST(TestOprDNN, ConvBiasInt4NCHW) {
+    auto run = [](size_t N, size_t C, size_t H, size_t W, size_t F, size_t S,
+                  size_t P) {
+        auto cn = CompNode::load("cpu0");
+        auto graph = ComputingGraph::make();
+
+        HostTensorGenerator<dtype::Int8> gen;
+        auto mkvar = [&gen](const char* name, const TensorShape& shp,
+                            const DType& dtype,
+                            std::shared_ptr<ComputingGraph> graph,
+                            const CompNode& cn) {
+            return opr::TypeCvt::make(
+                    opr::Host2DeviceCopy::make(*graph, gen(shp, cn))
+                            .rename(name),
+                    dtype);
+        };
+        auto mkcvar = [&gen](const char* name, const TensorShape& shp,
+                             const DType& dtype,
+                             std::shared_ptr<ComputingGraph> graph,
+                             const CompNode& cn) {
+            return opr::TypeCvt::make(
+                    opr::SharedDeviceTensor::make(*graph, *gen(shp, cn))
+                            .rename(name),
+                    dtype);
+        };
+
+        using Policy = opr::ConvBias::ExecutionPolicy;
+        using Strategy = Policy::Strategy;
+        auto x = mkvar("x", {N, C * 4, H, W}, dtype::QuantizedS4(1.19960327f),
+                       graph, cn),
+             w = mkcvar("w1", {C, C * 4, F, F}, dtype::QuantizedS4(1.19970327f),
+                        graph, cn),
+             b = mkcvar("b1", {1, C, 1, 1},
+                        dtype::QuantizedS32(1.19960327f * 1.19970327f), graph,
+                        cn);
+        opr::ConvBias::Param param;
+        param.format = opr::ConvBias::Param::Format::NCHW;
+        param.nonlineMode = opr::ConvBias::Param::NonlineMode::IDENTITY;
+        param.stride_h = param.stride_w = S;
+        param.pad_h = param.pad_w = P;
+        Policy policy;
+        policy.strategy = Strategy::PROFILE;
+
+        auto y = opr::ConvBias::make(
+                x, w, b, param, policy,
+                OperatorNodeConfig{dtype::QuantizedS4(11.9960501f)});
+        y = opr::TypeCvt::make(y, dtype::Float32());
+        auto x_f32 = opr::TypeCvt::make(x, dtype::Float32()),
+             w_f32 = opr::TypeCvt::make(w, dtype::Float32()),
+             b_f32 = opr::TypeCvt::make(b, dtype::Float32());
+        auto y_f32 = opr::ConvBias::make(x_f32, w_f32, b_f32, param, policy);
+        auto y_q4 = opr::TypeCvt::make(y_f32, dtype::QuantizedS4{11.9960501f});
+        y_q4 = opr::TypeCvt::make(y_q4, dtype::Float32());
+        HostTensorND host_y, host_y_q4;
+        auto func = graph->compile({make_callback_copy(y, host_y),
+                                    make_callback_copy(y_q4, host_y_q4)});
+        func->execute();
+        MGB_ASSERT_TENSOR_NEAR(host_y, host_y_q4, 1e-3);
+    };
+    run(2, 64, 14, 14, 3, 2, 1);
+    run(2, 64, 7, 7, 3, 1, 1);
+    run(2, 64, 14, 14, 1, 2, 0);
+    run(2, 64, 7, 7, 1, 1, 0);
+}
+
+TEST(TestOprDNN, ConvBiasInt4NCHW64) {
+    auto nchw2nchw64 = [](SymbolVar x) {
+        auto y = opr::RelayoutFormat::make(
+                x, opr::RelayoutFormat::Param::Mode::NCHW_NCHW64);
+        return y;
+    };
+
+    auto nchw642nchw = [](SymbolVar x) {
+        auto y = opr::RelayoutFormat::make(
+                x, opr::RelayoutFormat::Param::Mode::NCHW64_NCHW);
+        return y;
+    };
+
+    auto run = [&](size_t N, size_t C, size_t H, size_t W, size_t F, size_t S,
+                   size_t P) {
+        auto cn = CompNode::load("cpu0");
+        auto graph = ComputingGraph::make();
+
+        HostTensorGenerator<dtype::Int8> gen;
+        auto mkvar = [&gen](const char* name, const TensorShape& shp,
+                            const DType& dtype,
+                            std::shared_ptr<ComputingGraph> graph,
+                            const CompNode& cn) {
+            return opr::TypeCvt::make(
+                    opr::Host2DeviceCopy::make(*graph, gen(shp, cn))
+                            .rename(name),
+                    dtype);
+        };
+        auto mkcvar = [&gen](const char* name, const TensorShape& shp,
+                             const DType& dtype,
+                             std::shared_ptr<ComputingGraph> graph,
+                             const CompNode& cn) {
+            return opr::TypeCvt::make(
+                    opr::SharedDeviceTensor::make(*graph, *gen(shp, cn))
+                            .rename(name),
+                    dtype);
+        };
+
+        using Policy = opr::ConvBias::ExecutionPolicy;
+        using Strategy = Policy::Strategy;
+        auto x = mkvar("x", {N, C / 16, H, W, 64},
+                       dtype::QuantizedS4(1.19960327f), graph, cn),
+             w = mkcvar("w1", {C, C / 16, F, F, 64},
+                        dtype::QuantizedS4(1.19970327f), graph, cn),
+             b = mkcvar("b1", {1, C / 64, 1, 1, 64},
+                        dtype::QuantizedS32(1.19960327f * 1.19970327f), graph,
+                        cn);
+        opr::ConvBias::Param param;
+        param.format = opr::ConvBias::Param::Format::NCHW64;
+        param.nonlineMode = opr::ConvBias::Param::NonlineMode::IDENTITY;
+        param.stride_h = param.stride_w = S;
+        param.pad_h = param.pad_w = P;
+        Policy policy;
+        policy.strategy = Strategy::PROFILE;
+
+        auto y = opr::ConvBias::make(
+                x, w, b, param, policy,
+                OperatorNodeConfig{dtype::QuantizedS4(11.9960501f)});
+        y = opr::TypeCvt::make(y, dtype::Float32());
+        x = nchw642nchw(x);
+        w = nchw642nchw(w);
+        b = nchw642nchw(b);
+        auto x_f32 = opr::TypeCvt::make(x, dtype::Float32()),
+             w_f32 = opr::TypeCvt::make(w, dtype::Float32()),
+             b_f32 = opr::TypeCvt::make(b, dtype::Float32());
+        param.format = opr::ConvBias::Param::Format::NCHW;        
+        auto y_f32 = opr::ConvBias::make(x_f32, w_f32, b_f32, param, policy);
+        auto y_q4 = opr::TypeCvt::make(y_f32, dtype::QuantizedS4{11.9960501f});
+        y_q4 = opr::TypeCvt::make(y_q4, dtype::Float32());
+        y_q4 = nchw2nchw64(y_q4);
+        HostTensorND host_y, host_y_q4;
+        auto func = graph->compile({make_callback_copy(y, host_y),
+                                    make_callback_copy(y_q4, host_y_q4)});
+        func->execute();
+        MGB_ASSERT_TENSOR_NEAR(host_y, host_y_q4, 1e-3);
+    };
+    run(2, 64, 14, 14, 3, 2, 1);
+    run(2, 64, 7, 7, 3, 1, 1);
+    run(2, 64, 14, 14, 1, 2, 0);
+    run(2, 64, 7, 7, 1, 1, 0);
+}
+
+TEST(TestOprDNN, ConvBiasInt4Serialize) {
+    using namespace serialization;
+
+    float inp_scale = 1.20210327f;
+    float filt_scale = 1.20210406f;
+    float bias_scale = inp_scale * filt_scale;
+    DType output_dtype = dtype::QuantizedS4{inp_scale};
+
+    HostTensorGenerator<dtype::Int8> gen;
+    std::shared_ptr<HostTensorND> xv;
+    auto mkvar = [&gen](const char* name, const DType& dtype,
+                        std::shared_ptr<ComputingGraph> graph,
+                        std::shared_ptr<HostTensorND> val) {
+        return opr::TypeCvt::make(
+                opr::Host2DeviceCopy::make(*graph, val).rename(name), dtype);
+    };
+    auto mkcvar =
+            [&gen](const char* name, const TensorShape& shp, const DType& dtype,
+                   std::shared_ptr<ComputingGraph> graph, const CompNode& cn) {
+                return opr::TypeCvt::make(
+                        opr::SharedDeviceTensor::make(*graph, *gen(shp, cn))
+                                .rename(name),
+                        dtype);
+            };
+
+    auto fname = output_file("ConvBiasInt4Serialize");
+    HostTensorND y1, y2;
+    auto dump = [&]() {
+        opr::ConvBias::Param param;
+        param.mode = Mode::CONVOLUTION;
+
+        auto cn = CompNode::load("cpu0");
+        auto graph = ComputingGraph::make();
+        xv = gen({1, 64, 56, 56}, cn);
+        auto x = mkvar("x", dtype::QuantizedS4{inp_scale}, graph, xv);
+        auto w = mkcvar("w", {256, 64, 1, 1}, dtype::QuantizedS4{filt_scale}, graph, cn);
+        auto b = mkcvar("b", {1, 256, 1, 1}, dtype::QuantizedS32{bias_scale}, graph, cn);
+        auto y = opr::ConvBiasForward::make(x, w, b, param, {},
+                                            OperatorNodeConfig{output_dtype});
+        auto w1 = mkcvar("w1", {64, 256, 1, 1}, dtype::QuantizedS4{filt_scale},
+                         graph, cn);
+        auto b1 = mkcvar("b1", {1, 64, 1, 1}, dtype::QuantizedS32{bias_scale},
+                         graph, cn);
+        y = opr::ConvBiasForward::make(y, w1, b1, param, {},
+                                       OperatorNodeConfig{output_dtype});
+        y = opr::TypeCvt::make(y, dtype::Float32());
+        auto dumper = GraphDumper::make(OutputFile::make_fs(fname.c_str()));
+        auto func = graph->compile({make_callback_copy(y, y1)});
+        func->execute();
+        func->wait();
+        auto rst = dumper->dump({y});
+        ASSERT_EQ(rst.outputs.size(), 1u);
+    };
+
+    auto load = [&]() {
+        auto loader = GraphLoader::make(InputFile::make_fs(fname.c_str()));
+        auto rst = loader->load();
+        for (const auto& t : rst.tensor_map) {
+            t.second->copy_from(*xv).sync();
+        }
+        auto func = rst.graph->compile(
+                {make_callback_copy(rst.output_var_list[0], y2)});
+        func->execute();
+        func->wait();
+        ASSERT_EQ(rst.output_var_list.size(), 1u);
+        EXPECT_EQ(rst.output_var_list[0].dtype(), dtype::Float32());
+    };
+
+    dump();
+    load();
+    MGB_ASSERT_TENSOR_NEAR(y1, y2, 1e-3);
+}
+
+TEST(TestOprDNN, ConvBiasInt4SerializeWithParamFuse) {
+    using namespace serialization;
+
+    float inp_scale = 1.20210327f;
+    float filt_scale = 1.20210406f;
+    float bias_scale = inp_scale * filt_scale;
+    DType output_dtype = dtype::QuantizedS4{inp_scale};
+
+    HostTensorGenerator<dtype::Int8> gen;
+    std::shared_ptr<HostTensorND> xv;
+    auto mkvar = [&gen](const char* name, const DType& dtype,
+                        std::shared_ptr<ComputingGraph> graph,
+                        std::shared_ptr<HostTensorND> val) {
+        return opr::TypeCvt::make(
+                opr::Host2DeviceCopy::make(*graph, val).rename(name), dtype);
+    };
+    auto mkcvar =
+            [&gen](const char* name, const TensorShape& shp, const DType& dtype,
+                   std::shared_ptr<ComputingGraph> graph, const CompNode& cn) {
+                return opr::TypeCvt::make(
+                        opr::SharedDeviceTensor::make(*graph, *gen(shp, cn))
+                                .rename(name),
+                        dtype);
+            };
+
+    auto fname = output_file("ConvBiasInt4SerializeWithParamFuse");
+    HostTensorND y1, y2;
+    auto dump = [&]() {
+        opr::ConvBias::Param param;
+        param.mode = Mode::CONVOLUTION;
+
+        auto cn = CompNode::load("cpu0");
+        auto graph = ComputingGraph::make();
+        xv = gen({1, 64, 56, 56}, cn);
+        auto x = mkvar("x", dtype::QuantizedS4{inp_scale}, graph, xv);
+        auto w = mkcvar("w", {256, 64, 1, 1}, dtype::QuantizedS4{filt_scale}, graph, cn);
+        auto b = mkcvar("b", {1, 256, 1, 1}, dtype::QuantizedS32{bias_scale}, graph, cn);
+        auto y = opr::ConvBiasForward::make(x, w, b, param, {},
+                                            OperatorNodeConfig{output_dtype});
+        auto w1 = mkcvar("w1", {64, 256, 1, 1}, dtype::QuantizedS4{filt_scale},
+                         graph, cn);
+        auto b1 = mkcvar("b1", {1, 64, 1, 1}, dtype::QuantizedS32{bias_scale},
+                         graph, cn);
+        y = opr::ConvBiasForward::make(y, w1, b1, param, {},
+                                       OperatorNodeConfig{output_dtype});
+        y = opr::TypeCvt::make(y, dtype::Float32());
+        SymbolVar y_param_fused;
+        unpack_vector(gopt::GraphOptimizer{}
+                              .add_pass<gopt::ParamFusePass>()
+                              .apply({{y}})
+                              .endpoint_vars(),
+                      y_param_fused);
+        auto dumper = GraphDumper::make(OutputFile::make_fs(fname.c_str()));
+        auto func = graph->compile({make_callback_copy(y_param_fused, y1)});
+        func->execute();
+        func->wait();
+        auto rst = dumper->dump({y_param_fused});
+        ASSERT_EQ(rst.outputs.size(), 1u);
+    };
+
+    auto load = [&]() {
+        auto loader = GraphLoader::make(InputFile::make_fs(fname.c_str()));
+        auto rst = loader->load();
+        for (const auto& t : rst.tensor_map) {
+            t.second->copy_from(*xv).sync();
+        }
+        auto func = rst.graph->compile(
+                {make_callback_copy(rst.output_var_list[0], y2)});
+        func->execute();
+        func->wait();
+        ASSERT_EQ(rst.output_var_list.size(), 1u);
+        EXPECT_EQ(rst.output_var_list[0].dtype(), dtype::Float32());
+    };
+
+    dump();
+    load();
+    MGB_ASSERT_TENSOR_NEAR(y1, y2, 1e-3);
+}
+}  // namespace
+
 // vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}