fix(compiler): fix the bug that activation function is lost when conv…

…erting convolution without bias but with activation function. GitOrigin-RevId: 15d4378ae5759158b63945d41a4fda038ee2c480
MegEngine · May 10, 2024 · 6708597 · 6708597
1 parent 4bb59d2
commit 6708597
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 22 deletions.
diff --git a/compiler/lib/Conversion/MGBToKernel/MGBToKernel.cpp b/compiler/lib/Conversion/MGBToKernel/MGBToKernel.cpp
@@ -248,10 +248,12 @@ class ConvertConvLike final : public ConversionPattern {
                 return failure();
             operand_segment_sizes = {1, 1, 0, 0, 1};
         } else if (isa<MGB::ConvBias>(op)) {
-            // FIXME: only conv_bias(input, weight, bias) is supported now
-            if (operands.size() != 3)
+            if (operands.size() == 3)
+                operand_segment_sizes = {1, 1, 1, 0, 1};
+            else if (operands.size() == 2)
+                operand_segment_sizes = {1, 1, 0, 0, 1};
+            else
                 return failure();
-            operand_segment_sizes = {1, 1, 1, 0, 1};
         } else if (isa<MGB::ConvolutionBackwardData>(op)) {
             if (operands.size() != 2)
                 return failure();

diff --git a/compiler/lib/Target/MGB/importer.cpp b/compiler/lib/Target/MGB/importer.cpp
@@ -639,24 +639,12 @@ class Importer {
         } else if (auto conv = opr->try_cast_final<opr::ConvBiasForward>()) {
             auto&& p = conv->param();
             auto&& out = opr->output(0);
-            if (opr->input().size() == 2) {
-                CC_ASSERT(
-                        !is_dynamic_value(m_var2value.at(opr->input(0))) &&
-                        !is_dynamic_value(m_var2value.at(opr->input(1))));
-                mlir::Value value = m_builder.create<mlir::MGB::Convolution>(
-                        m_builder.getUnknownLoc(), var_to_shaped_type(out),
-                        m_var2value.at(opr->input(0)), m_var2value.at(opr->input(1)),
-                        p.mode, p.pad_h, p.pad_w, p.stride_h, p.stride_w, p.dilate_h,
-                        p.dilate_w, p.sparse, p.format, p.compute_mode);
-                m_var2value.emplace(out, value);
-            } else {
-                mlir::Value value = m_builder.create<mlir::MGB::ConvBias>(
-                        m_builder.getUnknownLoc(), var_to_shaped_type(out),
-                        var_array_to_value_array(opr->input(), true), p.nonlineMode,
-                        p.mode, p.sparse, p.format, p.pad_h, p.pad_w, p.stride_h,
-                        p.stride_w, p.dilate_h, p.dilate_w, p.compute_mode);
-                m_var2value.emplace(out, value);
-            }
+            mlir::Value value = m_builder.create<mlir::MGB::ConvBias>(
+                    m_builder.getUnknownLoc(), var_to_shaped_type(out),
+                    var_array_to_value_array(opr->input(), true), p.nonlineMode, p.mode,
+                    p.sparse, p.format, p.pad_h, p.pad_w, p.stride_h, p.stride_w,
+                    p.dilate_h, p.dilate_w, p.compute_mode);
+            m_var2value.emplace(out, value);
         } else if (auto resize_opr = opr->try_cast_final<opr::ResizeForward>()) {
             auto&& p = resize_opr->param();
             auto&& out = opr->output(0);

diff --git a/script/build_and_test_not_standard_os.sh b/script/build_and_test_not_standard_os.sh
@@ -28,7 +28,7 @@ cmake --build "$MEGCC_BUILD_DIR" -j$(nproc) --target mgb-to-tinynn --target mgb-
 
 function check_key_words() {
     #elf self mangle words, we do not care!!
-    white_list="@MEGW mgb1 5Mbg6 MGBi O:MgBnWk <mbG =MEG>Yr]< 4emUi0B >HMgE kMEG RmEg MbGV4 MEgIy @MEg mGe#S BMgb MGB( mBg: MBgr8C A&mGB mEg; mGb>/ mEg= .strtab .shstrtab A=MgE= mgb=g MGe= g=MgE <mgE= =Mgb> MGE< 8<MGE= =Mge =Mgb=K <MBG <MGE =MGB="
+    white_list="@MEGW mgb1 5Mbg6 MGBi O:MgBnWk <mbG =MEG>Yr]< 4emUi0B >HMgE kMEG RmEg MbGV4 MEgIy @MEg mGe#S BMgb MGB( mBg: MBgr8C A&mGB mEg; mGb>/ mEg= .strtab .shstrtab A=MgE= mgb=g MGe= g=MgE <mgE= =Mgb> MGE< 8<MGE= =Mge =Mgb=K <MBG <MGE =MGB= y?=Mgb="
     elf_file=$1
     if [ ! -f ${elf_file} ];then
         echo "ERR: can not find ${elf_file}"