Rewrite fully connected and batch matrix multiply tests

- Independent of operator API
- Uncovered bugs in bf16 packing
- Add error check for XNN_FLAG_TRANSPOSE_WEIGHTS in qc4w, which currently just crashes via calling a NULL function pointer (not obvious even under sanitizers).
- This reveals an issue, F32-QC4W and QD8-F32-QC4W actually have different weights formats (int4 vs. uint4).

PiperOrigin-RevId: 743573333

Author: dsharletg

src/datatype.c

@@ -88,6 +88,60 @@ bool xnn_datatype_is_quantized(enum xnn_datatype t) {
   return false;
 }
 
+bool xnn_datatype_is_channelwise_quantized(enum xnn_datatype t) {
+  switch (t) {
+    case xnn_datatype_qcint8:
+    case xnn_datatype_qcint32:
+    case xnn_datatype_qcint4:
+      return true;
+    case xnn_datatype_invalid:
+    case xnn_datatype_qint8:
+    case xnn_datatype_pqint8:
+    case xnn_datatype_quint8:
+    case xnn_datatype_qint32:
+    case xnn_datatype_qdint8:
+    case xnn_datatype_qduint8:
+    case xnn_datatype_qpint8:
+    case xnn_datatype_qbint4:
+    case xnn_datatype_fp32:
+    case xnn_datatype_fp16:
+    case xnn_datatype_bf16:
+    case xnn_datatype_int32:
+    case xnn_datatype_pfp16:
+    case xnn_datatype_pfp32:
+      return false;
+  }
+  XNN_UNREACHABLE;
+  return false;
+}
+
+bool xnn_datatype_is_blockwise_quantized(enum xnn_datatype t) {
+  switch (t) {
+    case xnn_datatype_qbint4:
+      return true;
+    case xnn_datatype_invalid:
+    case xnn_datatype_qint8:
+    case xnn_datatype_pqint8:
+    case xnn_datatype_quint8:
+    case xnn_datatype_qint32:
+    case xnn_datatype_qcint8:
+    case xnn_datatype_qcint32:
+    case xnn_datatype_qcint4:
+    case xnn_datatype_qdint8:
+    case xnn_datatype_qduint8:
+    case xnn_datatype_qpint8:
+    case xnn_datatype_fp32:
+    case xnn_datatype_fp16:
+    case xnn_datatype_bf16:
+    case xnn_datatype_int32:
+    case xnn_datatype_pfp16:
+    case xnn_datatype_pfp32:
+      return false;
+  }
+  XNN_UNREACHABLE;
+  return false;
+}
+
 
 size_t xnn_datatype_log2_size_bits(enum xnn_datatype t) {
   switch (t) {

src/operators/fully-connected-nc.c

@@ -1826,6 +1826,13 @@ enum xnn_status xnn_create_fully_connected_nc_f32_qc4w(
     return xnn_status_invalid_parameter;
   }
 
+  if (flags & XNN_FLAG_TRANSPOSE_WEIGHTS) {
+    xnn_log_error(
+      "failed to create %s operator with XNN_FLAG_TRANSPOSE_WEIGHTS: not supported",
+      xnn_operator_type_to_string(xnn_operator_type_fully_connected_nc_f32_qc4w));
+    return xnn_status_unsupported_parameter;
+  }
+
   for (size_t output_channel = 0; output_channel < output_channels; output_channel++) {
     if (kernel_scale[output_channel] <= 0.0f || !isnormal(kernel_scale[output_channel])) {
       xnn_log_error(

src/xnnpack/buffer.h

@@ -77,8 +77,8 @@ class NumericLimits<xnn_bfloat16> {
   static xnn_bfloat16 max_identity() { return -infinity(); }
 };
 
-template <typename T>
-class NumericLimits<quantized<T>> {
+template <typename T, typename Kind>
+class NumericLimits<quantized<T, Kind>> {
  public:
   static quantized<T> min() { return {std::numeric_limits<T>::lowest()}; }
   static quantized<T> max() { return {std::numeric_limits<T>::max()}; }
@@ -359,7 +359,7 @@ class Tensor {
   }
   size_t size() const {
     assert(is_contiguous());
-    return data_->size();
+    return end_ - begin_;
   }
   T* begin() { return data(); }
   T* end() { return end_; }
@@ -374,12 +374,18 @@ class Tensor {
   // Tensor, they do not affect the memory addressed by the Tensor. To realize
   // the effect of these operations, make a copy with `deep_copy`.
 
-  // Reorder the dimensions to extents = {extent(i) for i in perm}, and similar
-  // for strides.
-  Tensor<T, Alignment> transpose(const std::vector<size_t>& perm) const {
+  // Reorder the dimensions in `dims`. Dimensions not in dims maintain their
+  // relative ordering.
+  Tensor<T, Alignment> transpose(std::vector<size_t> perm) const {
+    // Sort idx to get the new locations
+    std::vector<size_t> sorted = perm;
+    std::sort(sorted.begin(), sorted.end());
+
     Tensor<T, Alignment> result(*this);
-    result.extents_ = permute(perm, extents_);
-    result.strides_ = permute(perm, strides_);
+    for (size_t i = 0; i < sorted.size(); i++) {
+      result.extents_[sorted[i]] = extent(perm[i]);
+      result.strides_[sorted[i]] = stride(perm[i]);
+    }
     return result;
   }
 
@@ -429,14 +435,16 @@ class Tensor {
 
     Tensor<T, Alignment> result(*this);
     std::vector<size_t> offsets(rank());
+    std::vector<size_t> maxs(rank());
     for (size_t i = 0; i < rank(); ++i) {
       offsets[i] = begins[i] < 0 ? extents_[i] + begins[i] : begins[i];
       result.extents_[i] =
           (ends[i] <= 0 ? extents_[i] + ends[i] : ends[i]) - offsets[i];
+      maxs[i] = result.extents_[i] - 1;
     }
 
     result.begin_ = begin_ + flat_offset(offsets);
-    result.end_ = result.begin_ + result.flat_offset(result.extents_);
+    result.end_ = result.begin_ + result.flat_offset(maxs) + 1;
 
     return result;
   }
@@ -460,6 +468,18 @@ class Tensor {
     return slice(dim, at, at + 1);
   }
 
+  // Slice the leading dimensions at the indices of `at`.
+  Tensor<T, Alignment> slice_leading(std::vector<size_t> at) const {
+    std::vector<int64_t> begins(rank());
+    std::vector<int64_t> ends(rank());
+    std::copy(at.begin(), at.end(), begins.begin());
+    std::copy(at.begin(), at.end(), ends.begin());
+    for (size_t i = 0; i < at.size(); ++i) {
+      ends[i] += 1;
+    }
+    return slice(begins, ends);
+  }
+
   // Remove `pre` elements from the beginning of each dimension, and `post`
   // elements from the end of each dimension.
   Tensor<T, Alignment> crop_padding(const index_type& pre,
@@ -730,7 +750,9 @@ xnn_quantization_params random_quantization(xnn_datatype datatype, Rng& rng,
   std::uniform_real_distribution<float> scale_dist{min_scale, max_scale};
   switch (datatype) {
     case xnn_datatype_qint8:
-      // int8 quantization assumes zero point is 0.
+    case xnn_datatype_qcint8:
+    case xnn_datatype_qcint4:
+      // signed integer quantization assumes zero point is 0.
       return {0, scale_dist(rng)};
     case xnn_datatype_quint8:
       return {u8_dist(rng), scale_dist(rng)};

src/xnnpack/datatype.h

@@ -28,6 +28,8 @@ bool xnn_datatype_is_integral(enum xnn_datatype t);
 
 // Returns true if the datatype is a quantized real datatype.
 bool xnn_datatype_is_quantized(enum xnn_datatype t);
+bool xnn_datatype_is_channelwise_quantized(enum xnn_datatype t);
+bool xnn_datatype_is_blockwise_quantized(enum xnn_datatype t);
 
 // Returns the size of an element of the datatype.
 size_t xnn_datatype_log2_size_bits(enum xnn_datatype t);
@@ -46,18 +48,22 @@ const char* xnn_datatype_to_string(enum xnn_datatype type);
 
 namespace xnnpack {
 
+struct channelwise {};
+
 // We need a type that distinguishes an intX_t from a quantized intX_t. We can't
 // do arithmetic on these, because we don't know the quantization parameters.
-template <typename T>
+template <typename T, typename Kind = void>
 struct quantized {
   T value;
   using type = T;
 
   operator T() const { return value; }
+  // Forward operator[] in case T is a sub-byte packed value.
+  auto operator[](size_t i) const { return value[i]; }
 
   quantized() = default;
   quantized(T t) : value(t) {}
-  quantized<T>& operator=(T t) {
+  quantized<T, Kind>& operator=(T t) {
     value = t;
     return *this;
   }
@@ -66,26 +72,26 @@ struct quantized {
 template <typename T>
 struct is_quantized : std::false_type {};
 
-template <typename T>
-struct is_quantized<quantized<T>> : std::true_type {};
+template <typename T, typename Kind>
+struct is_quantized<quantized<T, Kind>> : std::true_type {};
 
 template <typename T>
 struct unwrap_quantized {
   using type = T;
 };
 
-template <>
-struct unwrap_quantized<quantized<int8_t>> {
+template <typename Kind>
+struct unwrap_quantized<quantized<int8_t, Kind>> {
   using type = int8_t;
 };
 
-template <>
-struct unwrap_quantized<quantized<uint8_t>> {
+template <typename Kind>
+struct unwrap_quantized<quantized<uint8_t, Kind>> {
   using type = uint8_t;
 };
 
-template <>
-struct unwrap_quantized<quantized<int32_t>> {
+template <typename Kind>
+struct unwrap_quantized<quantized<int32_t, Kind>> {
   using type = int32_t;
 };
 
@@ -97,6 +103,9 @@ xnn_datatype xnn_datatype_of() {
     return xnn_datatype_quint8;
   } else if (std::is_same<T, xnnpack::quantized<int8_t>>::value) {
     return xnn_datatype_qint8;
+  } else if (std::is_same<
+                 T, xnnpack::quantized<int8_t, xnnpack::channelwise>>::value) {
+    return xnn_datatype_qcint8;
   } else if (std::is_same<T, xnnpack::quantized<int32_t>>::value) {
     return xnn_datatype_qint32;
   } else if (std::is_same<T, xnn_float16>::value) {