Fusion of operations and cast in mobilenet v1 conv2d causing large feature maps

[Mousius]

This is identified in the process of performing workspace calculation work, and can be seen in Depthwise Conv2D of quantized mobilenet_v1.

Conv2D of quantized mobilenet_v1 is producing the following Relay primitive function with fused elementwise operations that does not get really fused into a single loop, thus end up creating large feature maps. This is seen in the following generated relay and TIR primfuncs, which have these allocates:

allocate(PaddedInput, int16, [430592]);
allocate(DepthwiseConv2d, int32, [401408])

These two allocates create roughly 2.4 MB of data, moreover, the unusual cast operator at the end is making the inter-fused-operator tensors 16-bit wide where the model description states them to be 8-bits (this is taken from a single operator within mobilenet v1):

fn (%p0: Tensor[(1, 56, 56, 128), int16], %p1: Tensor[(3, 3, 128, 1), int16], %p2: Tensor[(1, 1, 1, 128), int32], Primitive=1) -> Tensor[(1, 56, 56, 128), int16] {
  %0 = nn.conv2d(%p0, %p1, padding=[1, 1, 1, 1], groups=128, channels=128, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWOI", out_dtype="int32") /* ty=Tensor[(1, 56, 56, 128), int32] */;
  %1 = add(%0, %p2) /* ty=Tensor[(1, 56, 56, 128), int32] */;
  %2 = fixed_point_multiply(%1, multiplier=2080045879, shift=-4) /* ty=Tensor[(1, 56, 56, 128), int32] */;
  %3 = clip(%2, a_min=0f, a_max=255f) /* ty=Tensor[(1, 56, 56, 128), int32] */;
  %4 = cast(%3, dtype="uint8") /* ty=Tensor[(1, 56, 56, 128), uint8] */;
  cast(%4, dtype="int16") /* ty=Tensor[(1, 56, 56, 128), int16] */
}

This gets translated to the following TIR primfunc :

primfn(placeholder_3: handle, placeholder_4: handle, placeholder_5: handle, T_cast_1: handle) -> ()
  attr = {"global_symbol": "fused_nn_conv2d_add_fixed_point_multiply_clip_cast_cast_21", "tir.noalias": True}
  buffers = {T_cast: Buffer(T_cast_2: Pointer(int16), int16, [1, 56, 56, 128], []),
             placeholder_2: Buffer(placeholder_6: Pointer(int32), int32, [1, 1, 1, 128], []),
             placeholder: Buffer(placeholder_7: Pointer(int16), int16, [1, 56, 56, 128], []),
             placeholder_1: Buffer(placeholder_8: Pointer(int16), int16, [3, 3, 128, 1], [])}
  buffer_map = {placeholder_3: placeholder, placeholder_4: placeholder_1, placeholder_5: placeholder_2, T_cast_1: T_cast} {
  attr [PaddedInput: Pointer(int16)] "storage_scope" = "global";
  allocate(PaddedInput, int16, [430592]);
  attr [DepthwiseConv2d: Pointer(int32)] "storage_scope" = "global";
  allocate(DepthwiseConv2d, int32, [401408]) {
    for (i1: int32, 0, 58) {
      for (i2: int32, 0, 58) {
        for (i3: int32, 0, 128) {
          PaddedInput[(((i1*7424) + (i2*128)) + i3)] = @tir.if_then_else(((((1 <= i1) && (i1 < 57)) && (1 <= i2)) && (i2 < 57)), (int16*)placeholder_7[((((i1*7168) + (i2*128)) + i3) - 7296)], 0i16, dtype=int16)
        }
      }
    }
    for (i: int32, 0, 56) {
      for (j: int32, 0, 56) {
        for (c: int32, 0, 128) {
          DepthwiseConv2d[(((i*7168) + (j*128)) + c)] = 0
          for (di: int32, 0, 3) {
            for (dj: int32, 0, 3) {
              DepthwiseConv2d[(((i*7168) + (j*128)) + c)] = ((int32*)DepthwiseConv2d[(((i*7168) + (j*128)) + c)] + (cast(int32, (int16*)PaddedInput[(((((i*7424) + (di*7424)) + (j*128)) + (dj*128)) + c)])*cast(int32, (int16*)placeholder_8[(((di*384) + (dj*128)) + c)])))
            }
          }
        }
      }
    }
    for (ax1: int32, 0, 56) {
      for (ax2: int32, 0, 56) {
        for (ax3: int32, 0, 128) {
          DepthwiseConv2d[(((ax1*7168) + (ax2*128)) + ax3)] = ((int32*)DepthwiseConv2d[(((ax1*7168) + (ax2*128)) + ax3)] + (int32*)placeholder_6[ax3])
        }
      }
    }
    for (i1_1: int32, 0, 56) {
      for (i2_1: int32, 0, 56) {
        for (i3_1: int32, 0, 128) {
          DepthwiseConv2d[(((i1_1*7168) + (i2_1*128)) + i3_1)] = @tir.q_multiply_shift((int32*)DepthwiseConv2d[(((i1_1*7168) + (i2_1*128)) + i3_1)], 2080045879, 31, -4, dtype=int32)
        }
      }
    }
    for (i1_2: int32, 0, 56) {
      for (i2_2: int32, 0, 56) {
        for (i3_2: int32, 0, 128) {
          DepthwiseConv2d[(((i1_2*7168) + (i2_2*128)) + i3_2)] = max(min((int32*)DepthwiseConv2d[(((i1_2*7168) + (i2_2*128)) + i3_2)], 255), 0)
        }
      }
    }
    for (ax1_1: int32, 0, 56) {
      for (ax2_1: int32, 0, 56) {
        for (ax3_1: int32, 0, 128) {
          PaddedInput[(((ax1_1*7168) + (ax2_1*128)) + ax3_1)] = cast(uint8, (int32*)DepthwiseConv2d[(((ax1_1*7168) + (ax2_1*128)) + ax3_1)])
        }
      }
    }
    for (ax1_2: int32, 0, 56) {
      for (ax2_2: int32, 0, 56) {
        for (ax3_2: int32, 0, 128) {
          T_cast_2[(((ax1_2*7168) + (ax2_2*128)) + ax3_2)] = cast(int16, (uint8*)PaddedInput[(((ax1_2*7168) + (ax2_2*128)) + ax3_2)])
        }
      }
    }
  }
}

[Mousius]

@manupa-arm @areusch

[sergio-grovety]

Hello @Mousius,
I would like to clarify something about the part of the problem related to casting to int16.
Please look at
/tvm/python/tvm/relay/qnn/op/legalizations.py->_qnn_conv2d_legalize_arm_cpu(attrs, inputs, types)
Here we see that
use_int8_on_arm == false for depthwise Conv2d operators,
and is_fast_int8_on_arm() also returns false if the target string doesn’t contain -mattr=+v8.2a,+dotprod.

So, if we have target string like llvm -device=arm_cpu, we will cast to int16.
Could you please point, which target string do you use and tell if my conclusion is correct or not.

PS: I may be misunderstanding what it means “the unusual cast operator at the end is making the inter-fused-operator tensors 16-bit”, if that’s true, could you please clarify this for me?

PS2: I use whole MobileNet, maybe I should cut out just one operator and check it?

Thank you!

[AriMIR]

Hello @Mousius,

Sorry for the large number of questions. I just started to deal with the mechanics of TVM and it is quite possible that some of the questions are pretty stupid.

I cut a Depthwise Conv2D with output (1, 56, 56, 128) from quantized mobilenet_v1 frozen graph (mobilenet_v1_1.0_224_quant_frozen.pb),
then transform it to tflite and check a relay and TIR.

There is my relay for Conv2D:
def @main(%MobilenetV1/MobilenetV1/Conv2d_2_pointwise/act_quant/FakeQuantWithMinMaxVars: Tensor[(1, 56, 56, 128), float32], %v_param_2: Tensor[(1, 3, 3, 128), float32], %v_param_3: Tensor[(128), float32]) {
%0 = qnn.quantize(%v_param_2, 0.0605373f, 160, out_dtype="uint8");
%1 = qnn.dequantize(%0, 0.0605373f, 160);
%2 = reshape(%1, newshape=[3, 3, 128, 1]);
%3 = nn.conv2d(%MobilenetV1/MobilenetV1/Conv2d_2_pointwise/act_quant/FakeQuantWithMinMaxVars, %2, padding=[1, 1, 1, 1], groups=128, channels=128, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWOI");
%4 = nn.bias_add(%3, %v_param_3, axis=3);
%5 = clip(%4, a_min=0f, a_max=6f);
%6 = qnn.quantize(%5, 0.0235285f, 0, out_dtype="uint8");
qnn.dequantize(%6, 0.0235285f, 0)
}

But it differs from the relay in your bug report.
Then I checked relay from the whole mobilenet: (from mobilenet_v1_1.0_224_quant_frozen.pb and then from mobilenet_v1_1.0_224_quant.tflite also),
But I haven’t found a relay like yours in the whole mobilenet relay on my side (tried to find “int16” or “fixed_point_multiply” lines).

Another difference is that we are using tflite quantizer (“qnn.quantize”, “qnn.dequantize” in the relay from our side)
I suggest you use another quantizer, because there are no “qnn” lines in your relay.

My questions on the above problem are:

1. It seems that you have a different version of tvm than mine, please indicate in which commit did you get this error?
2. Which quantizer did you use?
3. Please tell me why relays are different?

The next part of my question is regarding large allocations

1. In which place in the code do you check TIR? (TIR primfn on my side also differs, not at all, but in some parts)
2. Could you please give me an example when op fusion is correct

Thank you!

Best regards,
Arina Naumova,
Software developer, Grovety

[manupak]

Hi All,

Thanks for looking at this, for the relay please check the relay in 'Primitive' form (not the relay that goes on relay.build). This relay could be found at observing relay coming out of Optimize in build_module.cc. This is because the legalizations legalize the qnn dialect to primitive relay operators.

Hope this helps.

[Mousius]

Apologies for the slow reply, I've just checked out main fresh and checked this still happens. To recreate this I took the AOT test test_quant_mobilenet_tfl which uses mobilenet_v1_1.0_224_quant.tflite from https://storage.googleapis.com/download.tensorflow.org/models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz.

I put a breakpoint just after the code is generated (here-ish: https://github.com/apache/tvm/blob/main/tests/python/relay/aot/aot_test_utils.py#L570) , and you can see large allocations in tvmgen_default_fused_nn_conv2d_add_fixed_point_multiply_clip_cast_cast_21 within default_lib1.c in the temporary directory which is named in the variable base_path:

  void* PaddedInput = TVMBackendAllocWorkspace(1, dev_id, (uint64_t)861184, 0, 16);
  if (PaddedInput == NULL) {
    return -1;
  }
  void* DepthwiseConv2d = TVMBackendAllocWorkspace(1, dev_id, (uint64_t)1605632, 0, 32);
  if (DepthwiseConv2d == NULL) {
    return -1;
  }

Hope that helps 😸

[Alex-grovety]

Hello @Mousius,
I found that DepthwiseConv2d doesn't get fused as it uses default_schedule with auto_inline = False, so if set auto_inline = True we get (from default_lib1.c for test_quant_mobilenet_tfl):

TVM_DLL int32_t tvmgen_default_fused_nn_conv2d_add_fixed_point_multiply_clip_cast_cast_21(void* args, void* arg_type_ids, int32_t num_args, void* out_ret_value, void* out_ret_tcode, void* resource_handle) {
  void* arg0 = (((TVMValue*)args)[0].v_handle);
  int32_t arg0_code = ((int32_t*)arg_type_ids)[(0)];
  void* arg1 = (((TVMValue*)args)[1].v_handle);
  int32_t arg1_code = ((int32_t*)arg_type_ids)[(1)];
  void* arg2 = (((TVMValue*)args)[2].v_handle);
  int32_t arg2_code = ((int32_t*)arg_type_ids)[(2)];
  void* arg3 = (((TVMValue*)args)[3].v_handle);
  int32_t arg3_code = ((int32_t*)arg_type_ids)[(3)];
  void* placeholder = (((DLTensor*)arg0)[0].data);
  void* arg0_shape = (((DLTensor*)arg0)[0].shape);
  void* arg0_strides = (((DLTensor*)arg0)[0].strides);
  int32_t dev_id = (((DLTensor*)arg0)[0].device.device_id);
  void* placeholder1 = (((DLTensor*)arg1)[0].data);
  void* arg1_shape = (((DLTensor*)arg1)[0].shape);
  void* arg1_strides = (((DLTensor*)arg1)[0].strides);
  void* placeholder2 = (((DLTensor*)arg2)[0].data);
  void* arg2_shape = (((DLTensor*)arg2)[0].shape);
  void* arg2_strides = (((DLTensor*)arg2)[0].strides);
  void* T_cast = (((DLTensor*)arg3)[0].data);
  void* arg3_shape = (((DLTensor*)arg3)[0].shape);
  void* arg3_strides = (((DLTensor*)arg3)[0].strides);
  if (!(arg0_strides == NULL)) {
  }
  if (!(arg1_strides == NULL)) {
  }
  if (!(arg2_strides == NULL)) {
  }
  if (!(arg3_strides == NULL)) {
  }
  void* DepthwiseConv2d = TVMBackendAllocWorkspace(1, dev_id, (uint64_t)1605632, 0, 32);
  if (DepthwiseConv2d == NULL) {
    return -1;
  }
  for (int32_t i = 0; i < 56; ++i) {
    for (int32_t j = 0; j < 56; ++j) {
      for (int32_t c = 0; c < 128; ++c) {
        ((int32_t*)DepthwiseConv2d)[((((i * 7168) + (j * 128)) + c))] = 0;
        for (int32_t di = 0; di < 3; ++di) {
          for (int32_t dj = 0; dj < 3; ++dj) {
            ((int32_t*)DepthwiseConv2d)[((((i * 7168) + (j * 128)) + c))] = (((int32_t*)DepthwiseConv2d)[((((i * 7168) + (j * 128)) + c))] + (((int32_t)(((((1 <= (i + di)) && ((i + di) < 57)) && (1 <= (j + dj))) && ((j + dj) < 57)) ? ((int16_t*)placeholder)[(((((((i * 7168) + (di * 7168)) + (j * 128)) + (dj * 128)) + c) - 7296))] : (int16_t)0)) * ((int32_t)((int16_t*)placeholder1)[((((di * 384) + (dj * 128)) + c))])));
          }
        }
      }
    }
  }
  for (int32_t ax0_ax1_fused_ax2_fused_ax3_fused = 0; ax0_ax1_fused_ax2_fused_ax3_fused < 401408; ++ax0_ax1_fused_ax2_fused_ax3_fused) {
    int32_t _1 = (int32_t)(((((0 != 0) ? (((int64_t)(((int32_t*)DepthwiseConv2d)[(ax0_ax1_fused_ax2_fused_ax3_fused)] + ((int32_t*)placeholder2)[((ax0_ax1_fused_ax2_fused_ax3_fused & 127))])) << ((int64_t)0)) : ((int64_t)(((int32_t*)DepthwiseConv2d)[(ax0_ax1_fused_ax2_fused_ax3_fused)] + ((int32_t*)placeholder2)[((ax0_ax1_fused_ax2_fused_ax3_fused & 127))]))) * (int64_t)2080045879) + ((int64_t)1 << ((int64_t)((4 + 31) - 1)))) >> ((int64_t)(4 + 31)));
    int32_t _2 = (_1) < (255) ? (_1) : (255);
    ((int16_t*)T_cast)[(ax0_ax1_fused_ax2_fused_ax3_fused)] = ((int16_t)((uint8_t)((_2) > (0) ? (_2) : (0))));
  }
  if (TVMBackendFreeWorkspace(1, dev_id, DepthwiseConv2d) != 0) {
    return -1;
  }
  return 0;
}

we get cast to int16 if hardware hasn't support for fast int8 arithmetic operations. For hardware that supports fast int8 arithmetic operations we get (from default_lib1.c for test_quant_mobilenet_tfl):

TVM_DLL int32_t tvmgen_default_fused_nn_conv2d_subtract_add_fixed_point_multiply_clip_cast_10(void* args, void* arg_type_ids, int32_t num_args, void* out_ret_value, void* out_ret_tcode, void* resource_handle) {
  void* arg0 = (((TVMValue*)args)[0].v_handle);
  int32_t arg0_code = ((int32_t*)arg_type_ids)[(0)];
  void* arg1 = (((TVMValue*)args)[1].v_handle);
  int32_t arg1_code = ((int32_t*)arg_type_ids)[(1)];
  void* arg2 = (((TVMValue*)args)[2].v_handle);
  int32_t arg2_code = ((int32_t*)arg_type_ids)[(2)];
  void* arg3 = (((TVMValue*)args)[3].v_handle);
  int32_t arg3_code = ((int32_t*)arg_type_ids)[(3)];
  void* arg4 = (((TVMValue*)args)[4].v_handle);
  int32_t arg4_code = ((int32_t*)arg_type_ids)[(4)];
  void* placeholder = (((DLTensor*)arg0)[0].data);
  void* arg0_shape = (((DLTensor*)arg0)[0].shape);
  void* arg0_strides = (((DLTensor*)arg0)[0].strides);
  int32_t dev_id = (((DLTensor*)arg0)[0].device.device_id);
  void* placeholder1 = (((DLTensor*)arg1)[0].data);
  void* arg1_shape = (((DLTensor*)arg1)[0].shape);
  void* arg1_strides = (((DLTensor*)arg1)[0].strides);
  void* placeholder2 = (((DLTensor*)arg2)[0].data);
  void* arg2_shape = (((DLTensor*)arg2)[0].shape);
  void* arg2_strides = (((DLTensor*)arg2)[0].strides);
  void* placeholder3 = (((DLTensor*)arg3)[0].data);
  void* arg3_shape = (((DLTensor*)arg3)[0].shape);
  void* arg3_strides = (((DLTensor*)arg3)[0].strides);
  void* T_cast = (((DLTensor*)arg4)[0].data);
  void* arg4_shape = (((DLTensor*)arg4)[0].shape);
  void* arg4_strides = (((DLTensor*)arg4)[0].strides);
  if (!(arg0_strides == NULL)) {
  }
  if (!(arg1_strides == NULL)) {
  }
  if (!(arg2_strides == NULL)) {
  }
  if (!(arg3_strides == NULL)) {
  }
  if (!(arg4_strides == NULL)) {
  }
  void* PaddedInput = TVMBackendAllocWorkspace(1, dev_id, (uint64_t)430592, 1, 8);
  if (PaddedInput == NULL) {
    return -1;
  }
  void* DepthwiseConv2d = TVMBackendAllocWorkspace(1, dev_id, (uint64_t)1605632, 0, 32);
  if (DepthwiseConv2d == NULL) {
    return -1;
  }
  for (int32_t i1 = 0; i1 < 58; ++i1) {
    for (int32_t i2 = 0; i2 < 58; ++i2) {
      for (int32_t i3 = 0; i3 < 128; ++i3) {
        ((uint8_t*)PaddedInput)[((((i1 * 7424) + (i2 * 128)) + i3))] = ((uint8_t*)placeholder)[((((i1 * 7424) + (i2 * 128)) + i3))];
      }
    }
  }
  for (int32_t i = 0; i < 56; ++i) {
    for (int32_t j = 0; j < 56; ++j) {
      for (int32_t c = 0; c < 128; ++c) {
        ((int32_t*)DepthwiseConv2d)[((((i * 7168) + (j * 128)) + c))] = 0;
        for (int32_t di = 0; di < 3; ++di) {
          for (int32_t dj = 0; dj < 3; ++dj) {
            ((int32_t*)DepthwiseConv2d)[((((i * 7168) + (j * 128)) + c))] = (((int32_t*)DepthwiseConv2d)[((((i * 7168) + (j * 128)) + c))] + (((int32_t)((uint8_t*)PaddedInput)[((((((i * 7424) + (di * 7424)) + (j * 128)) + (dj * 128)) + c))]) * ((int32_t)((int8_t*)placeholder1)[((((di * 384) + (dj * 128)) + c))])));
          }
        }
      }
    }
  }
  for (int32_t ax1 = 0; ax1 < 56; ++ax1) {
    for (int32_t ax2 = 0; ax2 < 56; ++ax2) {
      for (int32_t ax3 = 0; ax3 < 128; ++ax3) {
        ((int32_t*)DepthwiseConv2d)[((((ax1 * 7168) + (ax2 * 128)) + ax3))] = (((int32_t*)DepthwiseConv2d)[((((ax1 * 7168) + (ax2 * 128)) + ax3))] - ((int32_t*)placeholder2)[((((ax1 * 7168) + (ax2 * 128)) + ax3))]);
      }
    }
  }
  for (int32_t ax11 = 0; ax11 < 56; ++ax11) {
    for (int32_t ax21 = 0; ax21 < 56; ++ax21) {
      for (int32_t ax31 = 0; ax31 < 128; ++ax31) {
        ((int32_t*)DepthwiseConv2d)[((((ax11 * 7168) + (ax21 * 128)) + ax31))] = (((int32_t*)DepthwiseConv2d)[((((ax11 * 7168) + (ax21 * 128)) + ax31))] + ((int32_t*)placeholder3)[(ax31)]);
      }
    }
  }
  for (int32_t i11 = 0; i11 < 56; ++i11) {
    for (int32_t i21 = 0; i21 < 56; ++i21) {
      for (int32_t i31 = 0; i31 < 128; ++i31) {
        ((int32_t*)DepthwiseConv2d)[((((i11 * 7168) + (i21 * 128)) + i31))] = ((int32_t)(((((0 != 0) ? (((int64_t)((int32_t*)DepthwiseConv2d)[((((i11 * 7168) + (i21 * 128)) + i31))]) << ((int64_t)0)) : ((int64_t)((int32_t*)DepthwiseConv2d)[((((i11 * 7168) + (i21 * 128)) + i31))])) * (int64_t)2080045879) + ((int64_t)1 << ((int64_t)((4 + 31) - 1)))) >> ((int64_t)(4 + 31))));
      }
    }
  }
  for (int32_t i12 = 0; i12 < 56; ++i12) {
    for (int32_t i22 = 0; i22 < 56; ++i22) {
      for (int32_t i32 = 0; i32 < 128; ++i32) {
        int32_t _1 = ((int32_t*)DepthwiseConv2d)[((((i12 * 7168) + (i22 * 128)) + i32))];
        int32_t _2 = (_1) < (255) ? (_1) : (255);
        ((int32_t*)DepthwiseConv2d)[((((i12 * 7168) + (i22 * 128)) + i32))] = ((_2) > (0) ? (_2) : (0));
      }
    }
  }
  for (int32_t ax12 = 0; ax12 < 56; ++ax12) {
    for (int32_t ax22 = 0; ax22 < 56; ++ax22) {
      for (int32_t ax32 = 0; ax32 < 128; ++ax32) {
        ((uint8_t*)T_cast)[((((ax12 * 7168) + (ax22 * 128)) + ax32))] = ((uint8_t)((int32_t*)DepthwiseConv2d)[((((ax12 * 7168) + (ax22 * 128)) + ax32))]);
      }
    }
  }
  if (TVMBackendFreeWorkspace(1, dev_id, DepthwiseConv2d) != 0) {
    return -1;
  }
  if (TVMBackendFreeWorkspace(1, dev_id, PaddedInput) != 0) {
    return -1;
  }
  return 0;
}

[manupak]

Thanks for the investigation and the suggestions seems to improve the things a bit.

1.) Would it be possible to share the TIR Primfunc that is the source for this code generation ?
2.) In the first example, it seems we are unnecessarily maintaining a feature map of int32 just to do that in another set of loops -- even with int16 cast.

  for (int32_t ax0_ax1_fused_ax2_fused_ax3_fused = 0; ax0_ax1_fused_ax2_fused_ax3_fused < 401408; ++ax0_ax1_fused_ax2_fused_ax3_fused) {
    int32_t _1 = (int32_t)(((((0 != 0) ? (((int64_t)(((int32_t*)DepthwiseConv2d)[(ax0_ax1_fused_ax2_fused_ax3_fused)] + ((int32_t*)placeholder2)[((ax0_ax1_fused_ax2_fused_ax3_fused & 127))])) << ((int64_t)0)) : ((int64_t)(((int32_t*)DepthwiseConv2d)[(ax0_ax1_fused_ax2_fused_ax3_fused)] + ((int32_t*)placeholder2)[((ax0_ax1_fused_ax2_fused_ax3_fused & 127))]))) * (int64_t)2080045879) + ((int64_t)1 << ((int64_t)((4 + 31) - 1)))) >> ((int64_t)(4 + 31)));
    int32_t _2 = (_1) < (255) ? (_1) : (255);
    ((int16_t*)T_cast)[(ax0_ax1_fused_ax2_fused_ax3_fused)] = ((int16_t)((uint8_t)((_2) > (0) ? (_2) : (0))));
  }

Do we know why HW feature dictates that whole feature map needs to be int16 (as opposed to doing a cast on the fly -- fused) -- especially in this case, int16 tensor bleeding out the fused primitive relay operator that is not used until the next operator ?
My question why we cant the consuming operator create int16 just at the accumulator without needing the feature map wide int16 tensor.

cc: @anijain2305

3.) In the second example, can we get the fusion done as well ?

[Alex-grovety]

1.) TIR Primfunc for first example:

PrimFunc([placeholder, placeholder, placeholder, T_cast]) attrs={"from_legacy_te_schedule": (bool)1, "global_symbol": "tvmgen_default_fused_nn_conv2d_add_fixed_point_multiply_clip_cast_cast_21", "tir.noalias": (bool)1} {
  allocate DepthwiseConv2d[int32 * 401408], storage_scope = global
  for (i, 0, 56) {
    for (j, 0, 56) {
      for (c, 0, 128) {
        DepthwiseConv2d[(((i*7168) + (j*128)) + c)] = 0
        for (di, 0, 3) {
          for (dj, 0, 3) {
            DepthwiseConv2d[(((i*7168) + (j*128)) + c)] = (DepthwiseConv2d[(((i*7168) + (j*128)) + c)] + (int32(tir.if_then_else(((((1 <= (i + di)) && ((i + di) < 57)) && (1 <= (j + dj))) && ((j + dj) < 57)), placeholder[((((((i*7168) + (di*7168)) + (j*128)) + (dj*128)) + c) - 7296)], (int16)0))*int32(placeholder[(((di*384) + (dj*128)) + c)])))
          }
        }
      }
    }
  }
  for (ax0.ax1.fused.ax2.fused.ax3.fused, 0, 401408) {
    T_cast[ax0.ax1.fused.ax2.fused.ax3.fused] = int16(uint8(max(min(tir.q_multiply_shift((DepthwiseConv2d[ax0.ax1.fused.ax2.fused.ax3.fused] + placeholder[floormod(ax0.ax1.fused.ax2.fused.ax3.fused, 128)]), 2080045879, 31, -4), 255), 0)))
  }
}

TIR Primfunc for second example:

PrimFunc([placeholder, placeholder, placeholder, placeholder, T_cast]) attrs={"from_legacy_te_schedule": (bool)1, "global_symbol": "tvmgen_default_fused_nn_conv2d_subtract_add_fixed_point_multiply_clip_cast_10", "tir.noalias": (bool)1} {
  allocate PaddedInput[uint8 * 430592], storage_scope = global
  allocate DepthwiseConv2d[int32 * 401408], storage_scope = global
  for (i1, 0, 58) {
    for (i2, 0, 58) {
      for (i3, 0, 128) {
        PaddedInput[(((i1*7424) + (i2*128)) + i3)] = placeholder[(((i1*7424) + (i2*128)) + i3)]
      }
    }
  }
  for (i, 0, 56) {
    for (j, 0, 56) {
      for (c, 0, 128) {
        DepthwiseConv2d[(((i*7168) + (j*128)) + c)] = 0
        for (di, 0, 3) {
          for (dj, 0, 3) {
            DepthwiseConv2d[(((i*7168) + (j*128)) + c)] = (DepthwiseConv2d[(((i*7168) + (j*128)) + c)] + (int32(PaddedInput[(((((i*7424) + (di*7424)) + (j*128)) + (dj*128)) + c)])*int32(placeholder[(((di*384) + (dj*128)) + c)])))
          }
        }
      }
    }
  }
  for (ax1, 0, 56) {
    for (ax2, 0, 56) {
      for (ax3, 0, 128) {
        DepthwiseConv2d[(((ax1*7168) + (ax2*128)) + ax3)] = (DepthwiseConv2d[(((ax1*7168) + (ax2*128)) + ax3)] - placeholder[(((ax1*7168) + (ax2*128)) + ax3)])
      }
    }
  }
  for (ax1, 0, 56) {
    for (ax2, 0, 56) {
      for (ax3, 0, 128) {
        DepthwiseConv2d[(((ax1*7168) + (ax2*128)) + ax3)] = (DepthwiseConv2d[(((ax1*7168) + (ax2*128)) + ax3)] + placeholder[ax3])
      }
    }
  }
  for (i1, 0, 56) {
    for (i2, 0, 56) {
      for (i3, 0, 128) {
        DepthwiseConv2d[(((i1*7168) + (i2*128)) + i3)] = tir.q_multiply_shift(DepthwiseConv2d[(((i1*7168) + (i2*128)) + i3)], 2080045879, 31, -4)
      }
    }
  }
  for (i1, 0, 56) {
    for (i2, 0, 56) {
      for (i3, 0, 128) {
        DepthwiseConv2d[(((i1*7168) + (i2*128)) + i3)] = max(min(DepthwiseConv2d[(((i1*7168) + (i2*128)) + i3)], 255), 0)
      }
    }
  }
  for (ax1, 0, 56) {
    for (ax2, 0, 56) {
      for (ax3, 0, 128) {
        T_cast[(((ax1*7168) + (ax2*128)) + ax3)] = uint8(DepthwiseConv2d[(((ax1*7168) + (ax2*128)) + ax3)])
      }
    }
  }
}

3.) For second example with fusion we get (from default_lib1.c for test_quant_mobilenet_tfl):

TVM_DLL int32_t tvmgen_default_fused_nn_conv2d_subtract_add_fixed_point_multiply_clip_cast_10(void* args, void* arg_type_ids, int32_t num_args, void* out_ret_value, void* out_ret_tcode, void* resource_handle) {
  void* arg0 = (((TVMValue*)args)[0].v_handle);
  int32_t arg0_code = ((int32_t*)arg_type_ids)[(0)];
  void* arg1 = (((TVMValue*)args)[1].v_handle);
  int32_t arg1_code = ((int32_t*)arg_type_ids)[(1)];
  void* arg2 = (((TVMValue*)args)[2].v_handle);
  int32_t arg2_code = ((int32_t*)arg_type_ids)[(2)];
  void* arg3 = (((TVMValue*)args)[3].v_handle);
  int32_t arg3_code = ((int32_t*)arg_type_ids)[(3)];
  void* arg4 = (((TVMValue*)args)[4].v_handle);
  int32_t arg4_code = ((int32_t*)arg_type_ids)[(4)];
  void* placeholder = (((DLTensor*)arg0)[0].data);
  void* arg0_shape = (((DLTensor*)arg0)[0].shape);
  void* arg0_strides = (((DLTensor*)arg0)[0].strides);
  int32_t dev_id = (((DLTensor*)arg0)[0].device.device_id);
  void* placeholder1 = (((DLTensor*)arg1)[0].data);
  void* arg1_shape = (((DLTensor*)arg1)[0].shape);
  void* arg1_strides = (((DLTensor*)arg1)[0].strides);
  void* placeholder2 = (((DLTensor*)arg2)[0].data);
  void* arg2_shape = (((DLTensor*)arg2)[0].shape);
  void* arg2_strides = (((DLTensor*)arg2)[0].strides);
  void* placeholder3 = (((DLTensor*)arg3)[0].data);
  void* arg3_shape = (((DLTensor*)arg3)[0].shape);
  void* arg3_strides = (((DLTensor*)arg3)[0].strides);
  void* T_cast = (((DLTensor*)arg4)[0].data);
  void* arg4_shape = (((DLTensor*)arg4)[0].shape);
  void* arg4_strides = (((DLTensor*)arg4)[0].strides);
  if (!(arg0_strides == NULL)) {
  }
  if (!(arg1_strides == NULL)) {
  }
  if (!(arg2_strides == NULL)) {
  }
  if (!(arg3_strides == NULL)) {
  }
  if (!(arg4_strides == NULL)) {
  }
  void* DepthwiseConv2d = TVMBackendAllocWorkspace(1, dev_id, (uint64_t)1605632, 0, 32);
  if (DepthwiseConv2d == NULL) {
    return -1;
  }
  for (int32_t i = 0; i < 56; ++i) {
    for (int32_t j = 0; j < 56; ++j) {
      for (int32_t c = 0; c < 128; ++c) {
        ((int32_t*)DepthwiseConv2d)[((((i * 7168) + (j * 128)) + c))] = 0;
        for (int32_t di = 0; di < 3; ++di) {
          for (int32_t dj = 0; dj < 3; ++dj) {
            ((int32_t*)DepthwiseConv2d)[((((i * 7168) + (j * 128)) + c))] = (((int32_t*)DepthwiseConv2d)[((((i * 7168) + (j * 128)) + c))] + (((int32_t)((uint8_t*)placeholder)[((((((i * 7424) + (di * 7424)) + (j * 128)) + (dj * 128)) + c))]) * ((int32_t)((int8_t*)placeholder1)[((((di * 384) + (dj * 128)) + c))])));
          }
        }
      }
    }
  }
  for (int32_t ax0_ax1_fused_ax2_fused_ax3_fused = 0; ax0_ax1_fused_ax2_fused_ax3_fused < 401408; ++ax0_ax1_fused_ax2_fused_ax3_fused) {
    int32_t _1 = (int32_t)(((((0 != 0) ? (((int64_t)((((int32_t*)DepthwiseConv2d)[(ax0_ax1_fused_ax2_fused_ax3_fused)] + ((int32_t*)placeholder3)[((ax0_ax1_fused_ax2_fused_ax3_fused & 127))]) - ((int32_t*)placeholder2)[(ax0_ax1_fused_ax2_fused_ax3_fused)])) << ((int64_t)0)) : ((int64_t)((((int32_t*)DepthwiseConv2d)[(ax0_ax1_fused_ax2_fused_ax3_fused)] + ((int32_t*)placeholder3)[((ax0_ax1_fused_ax2_fused_ax3_fused & 127))]) - ((int32_t*)placeholder2)[(ax0_ax1_fused_ax2_fused_ax3_fused)]))) * (int64_t)2080045879) + ((int64_t)1 << ((int64_t)((4 + 31) - 1)))) >> ((int64_t)(4 + 31)));
    int32_t _2 = (_1) < (255) ? (_1) : (255);
    ((uint8_t*)T_cast)[(ax0_ax1_fused_ax2_fused_ax3_fused)] = ((uint8_t)((_2) > (0) ? (_2) : (0)));
  }
  if (TVMBackendFreeWorkspace(1, dev_id, DepthwiseConv2d) != 0) {
    return -1;
  }
  return 0;
}

TIR Primfunc:

PrimFunc([placeholder, placeholder, placeholder, placeholder, T_cast]) attrs={"from_legacy_te_schedule": (bool)1, "global_symbol": "tvmgen_default_fused_nn_conv2d_subtract_add_fixed_point_multiply_clip_cast_10", "tir.noalias": (bool)1} {
  allocate DepthwiseConv2d[int32 * 401408], storage_scope = global
  for (i, 0, 56) {
    for (j, 0, 56) {
      for (c, 0, 128) {
        DepthwiseConv2d[(((i*7168) + (j*128)) + c)] = 0
        for (di, 0, 3) {
          for (dj, 0, 3) {
            DepthwiseConv2d[(((i*7168) + (j*128)) + c)] = (DepthwiseConv2d[(((i*7168) + (j*128)) + c)] + (int32(placeholder[(((((i*7424) + (di*7424)) + (j*128)) + (dj*128)) + c)])*int32(placeholder[(((di*384) + (dj*128)) + c)])))
          }
        }
      }
    }
  }
  for (ax0.ax1.fused.ax2.fused.ax3.fused, 0, 401408) {
    T_cast[ax0.ax1.fused.ax2.fused.ax3.fused] = uint8(max(min(tir.q_multiply_shift(((DepthwiseConv2d[ax0.ax1.fused.ax2.fused.ax3.fused] + placeholder[floormod(ax0.ax1.fused.ax2.fused.ax3.fused, 128)]) - placeholder[ax0.ax1.fused.ax2.fused.ax3.fused]), 2080045879, 31, -4), 255), 0))
  }
}

[Alex-grovety]

2.) We can use compute_at for DepthwiseConv2d and get:

PrimFunc([placeholder, placeholder, placeholder, placeholder, T_cast]) attrs={"from_legacy_te_schedule": (bool)1, "global_symbol": "tvmgen_default_fused_nn_conv2d_subtract_add_fixed_point_multiply_clip_cast_10", "tir.noalias": (bool)1} {
  allocate DepthwiseConv2d[int32 * 1], storage_scope = global
  for (ax0.ax1.fused.ax2.fused.ax3.fused, 0, 401408) {
    DepthwiseConv2d[0] = 0
    for (di, 0, 3) {
      for (dj, 0, 3) {
        DepthwiseConv2d[0] = (DepthwiseConv2d[0] + (int32(placeholder[((((floordiv(ax0.ax1.fused.ax2.fused.ax3.fused, 7168)*7424) + (di*7424)) + (dj*128)) + floormod(ax0.ax1.fused.ax2.fused.ax3.fused, 7168))])*int32(placeholder[(((di*384) + (dj*128)) + floormod(ax0.ax1.fused.ax2.fused.ax3.fused, 128))])))
      }
    }
    T_cast[ax0.ax1.fused.ax2.fused.ax3.fused] = uint8(max(min(tir.q_multiply_shift(((DepthwiseConv2d[0] + placeholder[floormod(ax0.ax1.fused.ax2.fused.ax3.fused, 128)]) - placeholder[ax0.ax1.fused.ax2.fused.ax3.fused]), 2080045879, 31, -4), 255), 0))
  }
}

[manupak]

2.) is impressive :) -- 2.4MB to 4 bytes.
Also the tensor going out is a 8-bit tensor.

Any catch ?

cc : @u99127 @Mousius

[Alex-grovety]

output tensor is 8-bit for case when hardware has support for fast int8 arithmetic operations, if hardware hasn't support for fast int8 arithmetic operations we get:

PrimFunc([placeholder, placeholder, placeholder, T_cast]) attrs={"from_legacy_te_schedule": (bool)1, "global_symbol": "tvmgen_default_fused_nn_conv2d_add_fixed_point_multiply_clip_cast_cast_21", "tir.noalias": (bool)1} {
  allocate DepthwiseConv2d[int32 * 1], storage_scope = global
  for (ax0.ax1.fused.ax2.fused.ax3.fused, 0, 401408) {
    DepthwiseConv2d[0] = 0
    for (di, 0, 3) {
      for (dj, 0, 3) {
        DepthwiseConv2d[0] = (DepthwiseConv2d[0] + (int32(tir.if_then_else(((((1 <= (floordiv(ax0.ax1.fused.ax2.fused.ax3.fused, 7168) + di)) && ((floordiv(ax0.ax1.fused.ax2.fused.ax3.fused, 7168) + di) < 57)) && (1 <= (floordiv(floormod(ax0.ax1.fused.ax2.fused.ax3.fused, 7168), 128) + dj))) && ((floordiv(floormod(ax0.ax1.fused.ax2.fused.ax3.fused, 7168), 128) + dj) < 57)), placeholder[((((di*7168) + (dj*128)) + ax0.ax1.fused.ax2.fused.ax3.fused) - 7296)], (int16)0))*int32(placeholder[(((di*384) + (dj*128)) + floormod(ax0.ax1.fused.ax2.fused.ax3.fused, 128))])))
      }
    }
    T_cast[ax0.ax1.fused.ax2.fused.ax3.fused] = int16(uint8(max(min(tir.q_multiply_shift((DepthwiseConv2d[0] + placeholder[floormod(ax0.ax1.fused.ax2.fused.ax3.fused, 128)]), 2080045879, 31, -4), 255), 0)))
  }
}

[Alex-grovety]

To get this result:

PrimFunc([placeholder, placeholder, placeholder, T_cast]) attrs={"from_legacy_te_schedule": (bool)1, "global_symbol": "tvmgen_default_fused_nn_conv2d_add_fixed_point_multiply_clip_cast_cast_21", "tir.noalias": (bool)1} {
  allocate DepthwiseConv2d[int32 * 1], storage_scope = global
  for (ax0.ax1.fused.ax2.fused.ax3.fused, 0, 401408) {
    DepthwiseConv2d[0] = 0
    for (di, 0, 3) {
      for (dj, 0, 3) {
        DepthwiseConv2d[0] = (DepthwiseConv2d[0] + (int32(tir.if_then_else(((((1 <= (floordiv(ax0.ax1.fused.ax2.fused.ax3.fused, 7168) + di)) && ((floordiv(ax0.ax1.fused.ax2.fused.ax3.fused, 7168) + di) < 57)) && (1 <= (floordiv(floormod(ax0.ax1.fused.ax2.fused.ax3.fused, 7168), 128) + dj))) && ((floordiv(floormod(ax0.ax1.fused.ax2.fused.ax3.fused, 7168), 128) + dj) < 57)), placeholder[((((di*7168) + (dj*128)) + ax0.ax1.fused.ax2.fused.ax3.fused) - 7296)], (int16)0))*int32(placeholder[(((di*384) + (dj*128)) + floormod(ax0.ax1.fused.ax2.fused.ax3.fused, 128))])))
      }
    }
    T_cast[ax0.ax1.fused.ax2.fused.ax3.fused] = int16(uint8(max(min(tir.q_multiply_shift((DepthwiseConv2d[0] + placeholder[floormod(ax0.ax1.fused.ax2.fused.ax3.fused, 128)]), 2080045879, 31, -4), 255), 0)))
  }
}

I added depthwise_conv2d_nhwc schedule for x86

[Mousius]

This looks great @Alex-grovety, a huge improvement!

As far as next steps, ideally from here:

tvm/python/tvm/relay/qnn/op/legalizations.py

Lines 384 to 388 in 02fbaf0

	def _qnn_dense_legalize_arm_cpu(attrs, inputs, types):
	# ARM prefers the dtypes to be same.
	if is_fast_int8_on_arm():
	return helper_change_dtypes_to_be_same(attrs, inputs, types, relay.qnn.op.dense)
	return helper_no_fast_int8_hw_legalization(attrs, inputs, types, relay.nn.dense)

We should default to this schedule if the architecture is not is_aarch64_arm() and not "+neon" in target.mattr, something like:

@qnn_dense_legalize.register("arm_cpu")
def _qnn_dense_legalize_arm_cpu(attrs, inputs, types):
    # Advanced SIMD present and no dot product extension
    if (is_aarch64_arm() or "+neon" in target.mattr) and not is_fast_int8_on_arm():
        return helper_no_fast_int8_hw_legalization(attrs, inputs, types, relay.nn.dense)
    return helper_change_dtypes_to_be_same(attrs, inputs, types, relay.qnn.op.dense)

Which would be the schedule you've suggested, as the default schedule for the Arm architecture.

[areusch]

Adding a pointer to the implementation: sergio-grovety@6d7bdaa

[Alex-grovety]

If I understand correctly it will be

elif layout == "NHWC":
    assert kernel_layout == "HWOI"
    is_aarch64 = topi.arm_cpu.arm_utils.is_aarch64_arm()
    if is_aarch64 or "+neon" in target.mattr:
        strategy.add_implementation(
            wrap_compute_conv2d(topi.arm_cpu.compute_depthwise_conv2d_nhwc),
            wrap_topi_schedule(topi.arm_cpu.schedule_depthwise_conv2d_nhwc),
            name="depthwise_conv2d_nhwc.arm_cpu",
        )
    else:
        strategy.add_implementation(
            wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc),
            wrap_topi_schedule(topi.x86.schedule_depthwise_conv2d_nhwc),
            name="depthwise_conv2d_nhwc.arm_cpu",
        )

instead of

tvm/python/tvm/relay/op/strategy/arm_cpu.py

Lines 198 to 204 in 02fbaf0

	elif layout == "NHWC":
	assert kernel_layout == "HWOI"
	strategy.add_implementation(
	wrap_compute_conv2d(topi.arm_cpu.compute_depthwise_conv2d_nhwc),
	wrap_topi_schedule(topi.arm_cpu.schedule_depthwise_conv2d_nhwc),
	name="depthwise_conv2d_nhwc.arm_cpu",
	)