Add bias support to torchao kernels (#1879)

Add bias support to torchao kernels (#1879)

Summary:
Pull Request resolved: #1879

This diff adds bias support for torchao kernels / quantizer.

Differential Revision: D71093679

Author: metascroy

torchao/experimental/kernels/cpu/aarch64/benchmarks/benchmark_linear.cpp

@@ -34,30 +34,38 @@ channelwise_8bit_activation_groupwise_lowbit_weight_1x1x32_f32_neondot(
           has_clamp);
 
   std::vector<char> activation_data(
-      activation_data_size<has_weight_zeros>(m, k, group_size));
-  prepare_activation_data<has_weight_zeros>(
+      activation_data_size(m, k, group_size, has_weight_zeros));
+  prepare_activation_data(
       (void*)activation_data.data(),
       m,
       k,
       group_size,
-      test_case.activations.data());
+      test_case.activations.data(),
+      has_weight_zeros);
 
-  std::vector<char> weight_data(
-      weight_data_size<weight_nbit, has_weight_zeros, has_bias>(
-          n, k, group_size));
-  prepare_weight_data<weight_nbit, has_weight_zeros, has_bias>(
+  std::vector<char> weight_data(weight_data_size<weight_nbit>(
+      n, k, group_size, has_weight_zeros, has_bias));
+  int8_t* weight_zeros_ptr = nullptr;
+  if (has_weight_zeros) {
+    weight_zeros_ptr = test_case.weight_zeros.data();
+  }
+  float* bias_ptr = nullptr;
+  if (has_bias) {
+    bias_ptr = test_case.bias.data();
+  }
+  prepare_weight_data<weight_nbit>(
       (void*)weight_data.data(),
       n,
       k,
       group_size,
       test_case.weight_qvals.data(),
       test_case.weight_scales.data(),
-      test_case.weight_zeros.data(),
-      test_case.bias.data());
+      weight_zeros_ptr,
+      bias_ptr);
 
   std::vector<float> output(m * k);
   for (auto _ : state) {
-    kernel<weight_nbit, has_weight_zeros, has_bias, has_clamp>(
+    kernel<weight_nbit>(
         output.data(),
         /*output_m_stride=*/n,
         m,
@@ -67,7 +75,10 @@ channelwise_8bit_activation_groupwise_lowbit_weight_1x1x32_f32_neondot(
         weight_data.data(),
         activation_data.data(),
         test_case.clamp_min,
-        test_case.clamp_max);
+        test_case.clamp_max,
+        has_weight_zeros,
+        has_bias,
+        has_clamp);
   }
 }
 
@@ -95,30 +106,38 @@ channelwise_8bit_activation_groupwise_lowbit_weight_1x4x16_f32_neondot(
           has_clamp);
 
   std::vector<char> activation_data(
-      activation_data_size<has_weight_zeros>(m, k, group_size));
-  prepare_activation_data<has_weight_zeros>(
+      activation_data_size(m, k, group_size, has_weight_zeros));
+  prepare_activation_data(
       (void*)activation_data.data(),
       m,
       k,
       group_size,
-      test_case.activations.data());
+      test_case.activations.data(),
+      has_weight_zeros);
 
-  std::vector<char> weight_data(
-      weight_data_size<weight_nbit, has_weight_zeros, has_bias>(
-          n, k, group_size));
-  prepare_weight_data<weight_nbit, has_weight_zeros, has_bias>(
+  std::vector<char> weight_data(weight_data_size<weight_nbit>(
+      n, k, group_size, has_weight_zeros, has_bias));
+  int8_t* weight_zeros_ptr = nullptr;
+  if (has_weight_zeros) {
+    weight_zeros_ptr = test_case.weight_zeros.data();
+  }
+  float* bias_ptr = nullptr;
+  if (has_bias) {
+    bias_ptr = test_case.bias.data();
+  }
+  prepare_weight_data<weight_nbit>(
       (void*)weight_data.data(),
       n,
       k,
       group_size,
       test_case.weight_qvals.data(),
       test_case.weight_scales.data(),
-      test_case.weight_zeros.data(),
-      test_case.bias.data());
+      weight_zeros_ptr,
+      bias_ptr);
 
   std::vector<float> output(m * k);
   for (auto _ : state) {
-    kernel<weight_nbit, has_weight_zeros, has_bias, has_clamp>(
+    kernel<weight_nbit>(
         output.data(),
         /*output_m_stride=*/n,
         m,
@@ -128,7 +147,10 @@ channelwise_8bit_activation_groupwise_lowbit_weight_1x4x16_f32_neondot(
         weight_data.data(),
         activation_data.data(),
         test_case.clamp_min,
-        test_case.clamp_max);
+        test_case.clamp_max,
+        has_weight_zeros,
+        has_bias,
+        has_clamp);
   }
 }
 
@@ -156,30 +178,38 @@ channelwise_8bit_activation_groupwise_lowbit_weight_1x8x16_f32_neondot(
           has_clamp);
 
   std::vector<char> activation_data(
-      activation_data_size<has_weight_zeros>(m, k, group_size));
-  prepare_activation_data<has_weight_zeros>(
+      activation_data_size(m, k, group_size, has_weight_zeros));
+  prepare_activation_data(
       (void*)activation_data.data(),
       m,
       k,
       group_size,
-      test_case.activations.data());
+      test_case.activations.data(),
+      has_weight_zeros);
 
-  std::vector<char> weight_data(
-      weight_data_size<weight_nbit, has_weight_zeros, has_bias>(
-          n, k, group_size));
-  prepare_weight_data<weight_nbit, has_weight_zeros, has_bias>(
+  std::vector<char> weight_data(weight_data_size<weight_nbit>(
+      n, k, group_size, has_weight_zeros, has_bias));
+  int8_t* weight_zeros_ptr = nullptr;
+  if (has_weight_zeros) {
+    weight_zeros_ptr = test_case.weight_zeros.data();
+  }
+  float* bias_ptr = nullptr;
+  if (has_bias) {
+    bias_ptr = test_case.bias.data();
+  }
+  prepare_weight_data<weight_nbit>(
       (void*)weight_data.data(),
       n,
       k,
       group_size,
       test_case.weight_qvals.data(),
       test_case.weight_scales.data(),
-      test_case.weight_zeros.data(),
-      test_case.bias.data());
+      weight_zeros_ptr,
+      bias_ptr);
 
   std::vector<float> output(m * k);
   for (auto _ : state) {
-    kernel<weight_nbit, has_weight_zeros, has_bias, has_clamp>(
+    kernel<weight_nbit>(
         output.data(),
         /*output_m_stride=*/n,
         m,
@@ -189,7 +219,10 @@ channelwise_8bit_activation_groupwise_lowbit_weight_1x8x16_f32_neondot(
         weight_data.data(),
         activation_data.data(),
         test_case.clamp_min,
-        test_case.clamp_max);
+        test_case.clamp_max,
+        has_weight_zeros,
+        has_bias,
+        has_clamp);
   }
 }