estimate_bitwidth.m

function [bits, info] = estimate_bitwidth(A, B, opts)
  %ESTIMATE_BITWIDTH Estimate fixed-point bit width for emulation.
  %   [BITS,INFO]=ESTIMATE_BITWIDTH(A,B) estimates the bit width required
  %   to emulate precision-p floating-point computation of the product A*B.
  %
  %   [BITS,INFO]=ESTIMATE_BITWIDTH(A,B,OPTS) accepts the following
  %   name-value options via the struct-like argument OPTS.
  %       OPTS.variant     - 'unbalanced' or 'balanced' (default).
  %       OPTS.sharpness   - 'conservative' or 'sharp' (default).
  %       OPTS.nBlocks     - number of blocks, 1 <= k <= size(A,2).
  %       OPTS.blockSize   - size of blocks, 1 <= k <= size(A,2).
  %       OPTS.precision   - precision of target format (default 53).
  %       OPTS.formulation - 'exponent' (default) or 'log2'.
  %       OPTS.extraBits   - additional bits. If omitted, defaults to 2
  %                          for 'log2' and 4 for 'exponent'.
  %       OPTS.estimator   - 'sum' or 'max' (default)
  %
  %   The 'nBlocks' field has precedence over the 'blockSize' field. The
  %   latter is only used if the former is not specified. If neither is
  %   specified, the default behaviour is to use a single block.
  %
  %   Outputs:
  %       E    - Total bit width required for the product A*B. For balanced
  %              variants, this is a scalar. For unbalanced variants, this
  %              is a struct with fields 'bitsA' and 'bitsB'.
  %       INFO - diagnostic information (currently unused).

  arguments
    A {mustBeFloat, mustBeReal, mustBeFinite}
    B {mustBeFloat, mustBeReal, mustBeFinite}
    opts.variant (1,1) string = "balanced"
    opts.sharpness (1,1) string = "sharp"
    opts.nBlocks (1,1) double {mustBeInteger, mustBeNonnegative} = 0
    opts.blockSize (1,1) double {mustBeInteger, mustBeNonnegative} = 0
    opts.precision (1,1) double {mustBeInteger, mustBeNonnegative} = 53
    opts.formulation (1,1) string = "exponent"
    opts.extraBits (1,1) double {mustBeInteger} = -1
    opts.estimator (1,1) string = "max"
  end

  %% Parse and validate input.
  [m, n] = size(A);
  [n2, p] = size(B);
  if n ~= n2
    error('estimate_bitwidth:DimensionMismatch', ...
          'Matrices must be conformable');
  end

  sharpness = lower(opts.sharpness);
  mustBeMember(sharpness, ["conservative", "sharp"]);
  sharp = sharpness == "sharp";

  variant = lower(opts.variant);
  mustBeMember(variant, ["balanced", "unbalanced"]);
  unbalanced = variant == "unbalanced";

  nBlocks = opts.nBlocks; % This is k in the note.
  if nBlocks > n
    error('estimate_bitwidth:InvalidnBlocks',...
          'nBlocks must satisfy 1 <= nBlocks <= size(A,2).');
  end
  blockSize = opts.blockSize;
  if blockSize > n
    error('estimate_bitwidth:InvalidBlockSize',...
          'blockSize must satisfy 1 <= blockSize <= size(A,2).');
  end

  % If nBlocks is specified, use it and ignore blockSize.
  % If blockSize only is specified, use it.
  % If neither is specified, default to a single block.
  if nBlocks ~= 0
    blockSize = ceil (n / nBlocks);
  elseif blockSize ~= 0
    nBlocks = ceil(n / blockSize);
  else
    nBlocks = 1;
    blockSize = ceil(n / nBlocks);
  end

  precision = opts.precision;

  formulation = lower(opts.formulation);
  mustBeMember(formulation, ["exponent", "log2"]);
  useExponent = formulation == "exponent";

  if opts.extraBits == -1
    if useExponent
      extraBits = 3;
    else
      extraBits = 1;
    end
  elseif opts.extraBits < 0
    error('estimate_bitwidth:InvalidExtraBits',...
          'extraBits must be nonnegative when specified.');
  else
    extraBits = opts.extraBits;
  end

  estimatorString = lower(opts.estimator);
  mustBeMember(estimatorString, ["max", "sum"])
  if estimatorString == "max"
    estFun = @(x)max(x);
    estCols = @(M) max(M, [], 1);
    estRows = @(M) max(M, [], 2);
    estElementwise = @(A, B)(max(A, B));
  else
    estFun = @(x)sum(x);
    estCols = @(M) sum(M, 1);
    estRows = @(M) sum(M, 2);
    estElementwise = @(A, B)(A + B);
  end

  if useExponent
    candidateFromRatio = @(numLeft, numRight, den) ...
      fpExponent(numLeft) + fpExponent(numRight) - fpExponent(den);
  else
    candidateFromRatio = @(numLeft, numRight, den) ...
      ceil(log2((numLeft .* numRight) ./ den));
  end

  %% Initialisation.
  if unbalanced
    bitsA = 0;
    bitsB = 0;
  else
    esc = 0;
  end

  threshold = 2^-precision;
  absA = abs(A);
  absB = abs(B);
  Amax = max(absA, [], 2);
  Bmax = max(absB, [], 1);

  if blockSize == 1

    %% Fine case (blocks have size 1).
    if unbalanced % Fine, unbalanced case.
      % Consider all inner products, one of row of C at a time.
      for i = 1:m
        Z    = absA(i,:).' .* absB;
        zest = estCols(Z);
        den  = threshold * zest;
        mask = Z > den;

        % In this case, we can skip zeros.
        valid = mask & den ~= 0;
        if any(valid, "all")
          candA = candidateFromRatio(Amax(i), absB, den);
          candB = candidateFromRatio(absA(i,:).', Bmax, den);
          bitsA = max(bitsA, max(candA(valid)));
          bitsB = max(bitsB, max(candB(valid)));
        end
      end
    else % Fine, balanced case.
      % Consider all inner products, one of row of C at a time.
      for i = 1:m
        z = estCols(absA(i,:).' .* absB);
        % In this case, we can skip zeros.
        mask = z ~= 0;
        candidates = candidateFromRatio(Amax(i), Bmax, z);
        esc = max(esc, max(candidates(mask)));
      end
    end

  else

    %% Coarse case.
    if sharp
      if unbalanced % Coarse, sharp, and unbalanced.
        %  xA(:,i) contains the nBlocks largest elements in row i of A.
        %  pA(:,i) are their indices.
        [topA, idxA] = maxk(absA, nBlocks + 1, 2);
        xA   = topA(:,1:nBlocks).';
        pA   = idxA(:,1:nBlocks).';
        Akp1 = topA(:,nBlocks + 1);

        %  yB(:,j) contains the nBlocks largest elements in column j of B.
        %  pB(:,j) are their indices.
        [topB, idxB] = maxk(absB, nBlocks + 1, 1);
        yB   = topB(1:nBlocks,:);
        pB   = idxB(1:nBlocks,:);
        Bkp1 = topB(nBlocks + 1,:);

        zest = zeros(m, p, 'like', absA);
        for i = 1:m
          zest(i,:) = estCols(xA(:,i) .* absB(pA(:,i),:));
        end
        for j = 1:p
          cand      = absA(:,pB(:,j)) .* yB(:,j).';
          zest(:,j) = max(zest(:,j), estRows(cand));
        end

        den = threshold * zest;
        for i = 1:m
          BA   = absB(pA(:,i),:);
          ZA   = xA(:,i) .* BA;
          mask = ZA > den(i,:);

          if any(mask, "all")
            candA = candidateFromRatio(Amax(i), BA, den(i,:));
            candB = candidateFromRatio(absA(i,pA(:,i)).', Bmax, den(i,:));

            bitsA = max(bitsA, max(candA(mask)));
            bitsB = max(bitsB, max(candB(mask)));
          end
        end
        for j = 1:p
          XA   = absA(:,pB(:,j));
          ZB   = XA .* yB(:,j).';
          mask = ZB > den(:,j);

          if any(mask, "all")
            candA = candidateFromRatio(Amax, yB(:,j).', den(:,j));
            candB = candidateFromRatio(Bmax(j), XA, den(:,j));

            bitsA = max(bitsA, max(candA(mask)));
            bitsB = max(bitsB, max(candB(mask)));
          end
        end

        candA = candidateFromRatio(Amax, Bkp1, den);
        candB = candidateFromRatio(Akp1, Bmax, den);
        bitsA = max(bitsA, max(candA, [], "all"));
        bitsB = max(bitsB, max(candB, [], "all"));
      else % Coarse, sharp, and balanced.
        [~, idxA] = maxk(absA, nBlocks, 2);
        [~, idxB] = maxk(absB, nBlocks, 1);

        esc = 0;
        for i = 1 : m
          pi = idxA(i,:);
          for j = 1 : p
            sigma = idxB(:,j);
            candA = absA(i,pi) .* absB(pi,j).';
            candB = absA(i,sigma) .* absB(sigma,j).';
            gammaTilde = max([estFun(candA), estFun(candB)]);
            esc  = max([esc, candidateFromRatio(Amax(i), Bmax(j), gammaTilde)]);
          end
        end
      end
    else
      if (unbalanced) % Coarse, conservative, and unbalanced.
        error('estimate_bitwidth:NotImplemented',...
              'Coarse, conservative, unbalanced algorithm not implemented.');
      else % Coarse, conservative, and balanced.
           % Approach of Schwarz et al.
        pad = blockSize * nBlocks - n;

        % Find min and max of each block.
        Ablkmax = reshape(max(...
            reshape([absA, zeros(m, pad, 'like', absA)], m, blockSize, nBlocks),...
            [], 2), m, nBlocks);
        Ablkmin = reshape(min(...
            reshape([absA,  inf(m, pad, 'like', absA)], m, blockSize, nBlocks),...
            [], 2), m, nBlocks);
        Bblkmax = reshape(max(...
            reshape([absB; zeros(pad, p, 'like', absB)], blockSize, nBlocks, p),...
            [], 1), nBlocks, p);
        Bblkmin = reshape(min(...
            reshape([absB;  inf(pad, p, 'like', absB)], blockSize, nBlocks, p),...
            [], 1), nBlocks, p);

        zest = zeros(m, p, 'like', absA);
        for ell = 1:nBlocks
          C = max(Ablkmax(:,ell) * Bblkmin(ell,:), ...
                  Ablkmin(:,ell) * Bblkmax(ell,:));
          zest = estElementwise(zest, C);
        end
        esc = max(candidateFromRatio(Amax, Bmax, zest), [], "all");
      end
    end
  end

  if ~unbalanced
    bits = esc + precision + extraBits;
    if isinf(bits)
      warning('estimate_bitwidth:infiniteSlices',...
        ['Emulation is not indicated.'])
    end
  end
  if unbalanced
    bits = struct('bitsA', bitsA + extraBits, 'bitsB', bitsB + extraBits);
    if isinf(bitsA) || isinf(bitsB)
      warning('estimate_bitwidth:infiniteSlices',...
              'Emulation is not indicated.')
    end
  end

  info = struct('precision', precision);

end

function e = fpExponent(x)
  e = -inf(size(x));
  positiveMask = x > 0;
  if any(positiveMask, "all")
    [~, ep] = log2(x(positiveMask));
    e(positiveMask) = ep;
  end
end