Fix optional bias and batch handling in cadence::fully_connected

Differential Revision: D102821213

Pull Request resolved: pytorch#19194

Author: hsharma35

backends/cadence/aot/ops_registrations.py

@@ -2527,7 +2527,7 @@ def quantized_max_pool2d_nhwc_meta(
 def fully_connected_meta(
     src: torch.Tensor,
     weight: torch.Tensor,
-    bias: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     # src comes in shape [leading_dims, in_dim]
     # weight comes in shape [out_dim, in_dim]

backends/cadence/aot/ref_implementations.py

@@ -633,10 +633,8 @@ def quantized_fully_connected_asym8uxasym8u_asym8u_per_tensor() -> torch.Tensor:
 def fully_connected(
     input_tensor: torch.Tensor,
     weight: torch.Tensor,
-    bias: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
-    if input_tensor.shape[0] != 1:
-        raise ValueError("Fully connected linear only supports batch size of 1")
     return F.linear(input_tensor, weight, bias)
 
 

backends/cadence/generic/operators/op_fully_connected.cpp

@@ -27,7 +27,8 @@ void linear(
     Tensor& output) {
   const float* __restrict__ input_data = input.const_data_ptr<float>();
   const float* __restrict__ weight_data = weight.const_data_ptr<float>();
-  const float* __restrict__ bias_data = bias.value().const_data_ptr<float>();
+  const float* __restrict__ bias_data =
+      bias.has_value() ? bias.value().const_data_ptr<float>() : nullptr;
   float* __restrict__ output_data = output.mutable_data_ptr<float>();
 
   // input comes in shape [batch_size, in_dim]
@@ -43,7 +44,7 @@ void linear(
 
   for (int i = 0; i < leading_dims; ++i) {
     for (int j = 0; j < M; ++j) {
-      float sum = bias_data[j];
+      float sum = bias_data != nullptr ? bias_data[j] : 0.0f;
       for (int k = 0; k < N; ++k) {
         sum += input_data[i * N + k] * weight_data[j * N + k];
       }