update ConvKernel isAvailable checking

GitOrigin-RevId: 69e316b8ca9f826d8134ebd7a932ccabc328e401
MegEngine · Nov 20, 2023 · feee3f2 · feee3f2
1 parent 252a8d9
commit feee3f2
Show file tree

Hide file tree

Showing 9 changed files with 13 additions and 12 deletions.
diff --git a/compiler/lib/KernelGen/Arm/Arm64/KernelPack.cpp b/compiler/lib/KernelGen/Arm/Arm64/KernelPack.cpp
@@ -3,10 +3,10 @@
 #include "BatchedMatmul/BatchedMatmul.h"
 #include "ConvKernel.h"
 #include "Elemwise/Elemwise.h"
-#include "Rotate.h"
 #include "InternalKernel/InternalKernel.h"
 #include "KernelPack.h"
 #include "MatMulKernel/MatMul.h"
+#include "Rotate.h"
 
 using namespace megcc;
 using namespace KernelGen;
@@ -46,7 +46,7 @@ struct AllA64Kernel {
 
         inner_map[KernelPack::KernType::BatchMatmulKernel] = {
                 std::make_shared<Arm64::Fp32BatchedMatmul>()};
-        
+
         inner_map[KernelPack::KernType::RotateKernel] = {
                 std::make_shared<Arm64::RotateKernel>()};
     }

diff --git a/compiler/lib/KernelGen/Arm/Arm64/Rotate.h b/compiler/lib/KernelGen/Arm/Arm64/Rotate.h
@@ -1,14 +1,14 @@
 #pragma once
-#include <string>
 #include <sstream>
-#include "compiler/KernelGen/KernelGen.h"
+#include <string>
 #include "Utils/SymbolHelper.h"
 #include "Utils/Utils.h"
+#include "compiler/KernelGen/KernelGen.h"
 namespace megcc {
 namespace KernelGen {
 namespace Arm64 {
 
-class RotateKernel : public KernelFunc  {
+class RotateKernel : public KernelFunc {
 public:
     bool IsAvailable(TContext* context) const override { return false; };
     std::string GetKernelBody(TContext* context) const override { return ""; };

diff --git a/compiler/lib/KernelGen/Common/ConvKernel.h b/compiler/lib/KernelGen/Common/ConvKernel.h
@@ -14,7 +14,8 @@ class ConvImpl : public KernelFunc {
     static bool is_channel_broadcast_bias(TContext* ctx) {
         if (is_bias(ctx)) {
             CCOperand bias = ctx->getAttrOprand("operand:2");
-            return bias.shape[0] == 1 && bias.shape[2] == 1 && bias.shape[3] == 1;
+            return (bias.shape[0] == 1 && bias.shape[2] == 1 && bias.shape[3] == 1) ||
+                   bias.shape.size() == 1;
         }
         return false;
     }

diff --git a/compiler/lib/Target/CMakeLists.txt b/compiler/lib/Target/CMakeLists.txt
@@ -1,4 +1,4 @@
 add_subdirectory(MGB)
 add_subdirectory(TinyNN)
 add_subdirectory(Hako)
-add_subdirectory(onnx)
+add_subdirectory(onnx)
diff --git a/compiler/lib/Target/onnx/CMakeLists.txt b/compiler/lib/Target/onnx/CMakeLists.txt
@@ -19,4 +19,4 @@ target_include_directories(
                                ${PROTOBUF_INCLUDE_DIR})
 # add onnx-imported
 target_link_libraries(MLIRONNXImporter PUBLIC $<BUILD_INTERFACE:onnx_imported>)
-# target_compile_options(MLIRONNXImporter PUBLIC -fexceptions)
+# target_compile_options(MLIRONNXImporter PUBLIC -fexceptions)
diff --git a/compiler/test/kernel/opr/arm/cv.cpp b/compiler/test/kernel/opr/arm/cv.cpp
@@ -164,6 +164,6 @@ TEST(AARCH64, CVrotateFp16) {
         checker.exec({{1, 19, 19, 1}, {}});
 
         checker.exec({{1, 19, 19, 3}, {}});
-    } 
+    }
 }
 #endif
diff --git a/compiler/tools/CMakeLists.txt b/compiler/tools/CMakeLists.txt
@@ -7,4 +7,4 @@ add_subdirectory(dump-kernel)
 add_subdirectory(megcc-translate)
 add_subdirectory(kernel_exporter)
 add_subdirectory(onnx-importer)
-add_subdirectory(onnx-to-tinynn)
+add_subdirectory(onnx-to-tinynn)
diff --git a/compiler/tools/onnx-importer/CMakeLists.txt b/compiler/tools/onnx-importer/CMakeLists.txt
@@ -4,4 +4,4 @@ add_llvm_executable(onnx-importer onnx-importer.cpp NO_INSTALL_RPATH)
 llvm_update_compile_flags(onnx-importer)
 target_link_libraries(onnx-importer PRIVATE ${dialect_libs} MLIRONNXImporter Common
                                             ${ONNX_LIBS} ${PROTOBUF_LIBS})
-mlir_check_all_link_libraries(onnx-importer)
+mlir_check_all_link_libraries(onnx-importer)
diff --git a/compiler/tools/onnx-to-tinynn/CMakeLists.txt b/compiler/tools/onnx-to-tinynn/CMakeLists.txt
@@ -17,4 +17,4 @@ target_link_libraries(
           HakoParse
           ${ONNX_LIBS}
           ${PROTOBUF_LIBS})
-mlir_check_all_link_libraries(onnx-to-tinynn)
+mlir_check_all_link_libraries(onnx-to-tinynn)
-Original file line number
+Diff line change
@@ Expand Up / @@ -164,6 +164,6 @@ TEST(AARCH64, CVrotateFp16) { @@
             checker.exec({{1, 19, 19, 1}, {}});
             checker.exec({{1, 19, 19, 3}, {}});
-        }
+        }
     }
     #endif