Support block-wise quantization

[huningxin]

Block-wise quantization divides input tensors into smaller blocks that are independently quantized, resulting in faster optimization and high precision quantization. It is used for popular language models, such as phi-3 mini int4 quantized model.

Native ML API's support

DML DML_OPERATOR_QUANTIZE and DML_OPERATOR_DEQUANTIZE introduced in Feature Level 6.3
CoreML constexpr_blockwise_shift_scale
TFLite: ?

Proposal

No API signature changes regarding to @fdwr 's proposal of dequantizeLinear and quantizeLinear ops.

MLOperand dequantizeLinear(MLOperand input, MLOperand scale, MLOperand zeroPoint, optional MLOperatorOptions options = {});
MLOperand quantizeLinear(MLOperand input, MLOperand scale, MLOperand zeroPoint, optional MLOperatorOptions options = {});

The block_size is an integer and implied by block_size = input_size / scale_size (where input_size % scale_size == 0) along a dimension. zeroPoint and scale should have the same shape.

[fdwr]

Thanks for the paper link. I'd be surprised if TFLite didn't have some blockwise support somewhere, but if not, it might need decomposition (e.g. scale and zeroPoint blockwise expanded up to the input shape via tf.tile or tf.repeats or tf.imaging.resize or some other similar function, then dq = (input - zeroPoint) * scale).

[fdwr]

One realization while decomposing the emulation for dequantizeLinear in #805 is that it's quite verbose (see below). The entire thing could be reduced to just the dequantizeLinear function by using expand if expand was augmented to accept any from-shape that was an integer multiple of the to-shape (or alternately via nearest-neighbor upsampling if resample was augmented from 2D to ND, like ONNX Resize). Blockwise expansion seems useful as a standalone concept (independent of quantization) at least for clean decomposability. Will create separate issue...

Current emulation code 🥲:

    function dequantizeLinear(builder, input, scale, zeroPoint, options) {
      // output = (input - zeroPoint) * scale
      const floatInput = builder.cast(input, scale.dataType);
      const floatZeroPoint = builder.cast(zeroPoint, scale.dataType);
      const upsampledScale = blockwiseExpand(builder, scale, input.shape);
      const upsampledZeroPoint = blockwiseExpand(builder, floatZeroPoint, input.shape);
      return builder.mul(builder.sub(floatInput, upsampledZeroPoint), upsampledScale);
    }

    function blockwiseExpand(builder, input, targetShape) {
      // This expands each axis by repeating the block the number of times per that axis, given the
      // original input shape and target shape. However, backend implementations might have much more
      // efficient upsampling operators that can accept multiple dimensions to upsample all
      // dimensions at once by integer multiples (like tile) using nearest neighbor resampling:
      // output = resample(scale, {sizes: input.shape})

      let expandedInput = input;

      for (let axis = 0; axis < input.shape.length; ++axis) {
        const inputShape = expandedInput.shape;
        const oldDimensionLength = inputShape[axis];
        const newDimensionLength = targetShape[axis];

        if (newDimensionLength != oldDimensionLength) {
          // Since tile/expand can only accept repetitions of entire dimension slices (not repeating
          // individual elements along an axis), temporarily reshape the tensor to enable them to broadcast
          // the elements up to the full block size, utilizing an inserted dimension of size 1.
          const elementRepeatCount = newDimensionLength / oldDimensionLength;
          const flattenedShape = getFlattenedShapeAroundAxis(inputShape, axis);
          const unexpandedShape = [flattenedShape[0], flattenedShape[1], 1, flattenedShape[2]];
          const expandedShape = [flattenedShape[0], flattenedShape[1], elementRepeatCount, flattenedShape[2]];
          const reshapedInput = builder.reshape(expandedInput, unexpandedShape);
          expandedInput = builder.expand(reshapedInput, expandedShape);
        }

        let newInputShape = [...inputShape];
        newInputShape[axis] = newDimensionLength;
        expandedInput = builder.reshape(expandedInput, newInputShape);
      }

      return expandedInput;
    }

    // Compute the flattened shape before and after the given axis, yielding a 3-element list.
    // e.g. inputShape = [2,3,4,5,6] with axis = 2 yields shape [6,4,30].
    // e.g. inputShape = [4] with axis = 0 yields shape [1,4,1].
    function getFlattenedShapeAroundAxis(inputShape, axis) {
      axis = Math.max(Math.min(axis, input.shape.length - 1), 0);
      const countBefore = axis.slice(0, axis).reduce((a, b) => a * b, 1);
      const countAfter = axis.slice(axis + 1, input.shape.length).reduce((a, b) => a * b, 1);
      return [countBefore, inputShape[axis], countAfter];
    }