From d2dfb1d487387dcf64988f9912a8f5ae60c39cbe Mon Sep 17 00:00:00 2001
From: Yvonne Chen <yche@synopsys.com>
Date: Thu, 16 Dec 2021 11:36:57 +0100
Subject: [PATCH] refine the code structure to support all axis>=3 case

---
 src/caffe/layers/conv_layer.cpp | 27 +++++----------------------
 1 file changed, 5 insertions(+), 22 deletions(-)

diff --git a/src/caffe/layers/conv_layer.cpp b/src/caffe/layers/conv_layer.cpp
index cd761535..28485526 100644
--- a/src/caffe/layers/conv_layer.cpp
+++ b/src/caffe/layers/conv_layer.cpp
@@ -236,28 +236,11 @@ void ConvolutionLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
 
   if(this->submanifold_sparse_)
   {
-    if(bottom[0]->num_axes()==4)
-    {
-      CHECK_EQ(bottom[0]->height(), top[0]->height())<<
-          "Input and output blob height not equal! Submanifold sparse computation is invalid!";
-      CHECK_EQ(bottom[0]->width(), top[0]->width())<<
-          "Input and output blob width not equal! Submanifold sparse computation is invalid!";
-    }
-    else if(bottom[0]->num_axes()==5)
-    {
-      CHECK_EQ(bottom[0]->shape(2), top[0]->shape(2))<<
-          "Input and output blob depth not equal! Submanifold sparse computation is invalid!";
-      CHECK_EQ(bottom[0]->shape(3), top[0]->shape(3))<<
-          "Input and output blob height not equal! Submanifold sparse computation is invalid!";
-      CHECK_EQ(bottom[0]->shape(4), top[0]->shape(4))<<
-          "Input and output blob width not equal! Submanifold sparse computation is invalid!";
-    }
-    else
-    {
-      CHECK_EQ(bottom[0]->num_axes(), 3)<<"Not support Submanifold sparse computation for such blob dimension yet!";
-      CHECK_EQ(bottom[0]->shape(2), top[0]->shape(2))<<
-           "Input and output blob length not equal! Submanifold sparse computation is invalid!";
-    }
+    CHECK_GE(bottom[0]->num_axes(), 3)<<"Input blob dimension must >=3!";
+    for(int i=2; i<bottom[0]->num_axes();i++)
+      CHECK_EQ(bottom[0]->shape(i), top[0]->shape(i))<<
+          "Input and output blob shape does not match! Submanifold sparse computation is invalid!";
+
     LOG(INFO)<<"Starts submanifold sparse computation.";
 
     for(int index=0; index<bottom[0]->count(2); index++)