[Transformation] Add information about depthwise conv to created nodes when...

[Transformation] Add information about depthwise conv to created nodes when lowering convs to matmul

src/finn/transformation/lower_convs_to_matmul.py

@@ -63,12 +63,25 @@ class LowerConvsToMatMul(Transformation):
                 ifm_dim = model.get_tensor_shape(n.input[0])[-1]  # assume NCHW
                 ofm_dim = model.get_tensor_shape(n.output[0])[-1]  # assume NCHW
 
-                # if depthwise conv create sparse matrix
+                # if depthwise conv create sparse matrix and variable "dw"
+                # to store as attribute in Im2Col that indicates that the created
+                # Im2Col node belongs to a depthwise convolution
+                dw = False
                 if group == ifm_ch and ofm_ch == ifm_ch:
                     W_sparse = np.zeros((ofm_ch, ifm_ch, k, k))
                     for ch in range(ifm_ch):
                         W_sparse[ch][ch] = W_conv[ch][0]
                     W_conv = W_sparse.astype(np.float32)
+                    # we need to store information of the
+                    # sparsity of the weight matrix. For this
+                    # we use the sparsity annotation of the
+                    # weight tensor
+                    sparsity = {"dw": {"kernel_shape": k}}
+                    model.set_tensor_sparsity(weight_name, sparsity)
+                    # additionally create variable "dw" to store
+                    # as attribute in Im2Col that indicates that the created
+                    # Im2Col node belongs to a depthwise convolution
+                    dw = True
 
                 # reuse conv weights for new matmul weights
                 # conv weights are [OFM][IFM][k][k]
@@ -124,6 +137,7 @@ class LowerConvsToMatMul(Transformation):
                     kernel_size=k,
                     pad_amount=pad,
                     input_shape="(1,{},{},{})".format(ifm_dim, ifm_dim, ifm_ch),
+                    dw=dw,
                 )
                 # do matmul
                 matmul_node = helper.make_node(