# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # # ruff: noqa: E731 """Utility functions for implementing Winograd convolutions [*] Fast Algorithms for Convolutional Neural Networks Andrew Lavin, Scott Gray https://arxiv.org/abs/1509.09308 https://github.com/andravin/wincnn """ from functools import reduce from operator import mul import numpy as np from tvm.contrib.pickle_memoize import memoize from ..utils import const_matrix # pylint: disable=invalid-name def _cook_toom_convolution(a, n, r): """Compute Cook-Toom convolution A,B,G matrices""" def _F_m(a, n): f = lambda j, i: reduce(mul, ((a[i] - a[k] if k != i else 1) for k in range(0, n - 1)), 1) F = np.fromfunction(np.vectorize(f), (1, n - 1), dtype=int) F = np.diagflat(F) F = np.append(F, np.zeros((n - 1, 1), dtype=int), axis=1) f = lambda i, j: 1 if j == (n - 1) else 0 z = np.fromfunction(np.vectorize(f), (1, n), dtype=int) return np.append(F, z, axis=0) def _A_m(a, m, n): f = lambda i, j: a[i] ** j A = np.fromfunction(np.vectorize(f), (m - 1, n), dtype=int) f = lambda i, j: 1 if j == (n - 1) else 0 z = np.fromfunction(np.vectorize(f), (1, n), dtype=int) return np.append(A, z, axis=0) def _B_m(a, n): f = lambda j, i: reduce(mul, ((a[i] - a[k] if k != i else 1) for k in range(0, n - 1)), 1) Ff = np.fromfunction(np.vectorize(f), (1, n - 1), dtype=int) f = lambda i, nth: ( ( reduce(mul, [(np.poly1d([1, -a[k]]) if k != i else 1) for k in range(0, n - 1)], 1) ).coef[n - 1 - nth - 1] / Ff[0, i] ) F = np.fromfunction(np.vectorize(f), (n - 1, n - 1), dtype=int) f = lambda i, j: -(a[i] ** (n - 1)) t = np.fromfunction(np.vectorize(f), (n - 1, 1), dtype=int) T = np.append(np.eye(n - 1), t, axis=1) return np.append(F.T.dot(T), np.array([np.eye(n)[n - 1]]), axis=0) alpha = n + r - 1 f = _F_m(a, alpha) if f[0, 0] < 0: f[0, :] *= -1 A = _A_m(a, alpha, n) G = _A_m(a, alpha, r).T G = G.dot(np.linalg.inv(f)).T B = _B_m(a, alpha) B = B.dot(f.T) return (A, B, G) def _interpolation_points(degree): """Propose filter points""" assert 2 < degree < 18 # Default interpolation lookup table # # [1] Error Analysis and Improving the Accuracy of Winograd Convolution for Deep Neural Networks # Barbara Barabasz, Andrew Anderson, Kirk M. Soodhalter, David Gregg # https://arxiv.org/abs/1803.10986 # # pylint: disable=bad-whitespace,line-too-long in_pts = [ # {invalid} [], # 01 {E=4.63E-08 on conv2d [1]} [], # 02 {E=7.65E-08 on F( 2,3) [1]} [0, -1, 1], # 03 {E=2.35E-07 on F( 3,3) [1]} [0, -1, 1, 1 / 2], # 04 {E=3.29E-07 on F( 4,3) [1]} [0, -1, 1, 1 / 2, -2], # 05 {E=6.81E-07 on F( 5,3) [1]} [0, -1, 1, 1 / 2, -2, -1 / 2], # 06 {E=8.79E-07 on F( 6,3) [1]} [0, -1, 1, 1 / 2, -1 / 2, 2, -2], # 07 {E=3.71E-06 on F( 7,3) [1]} [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4], # 08 {E=7.35E-06 on F( 8,3) [1]} [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4], # 09 {E=2.20E-05 on F( 9,3) [1]} [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 3 / 4, -4 / 3], # 10 {E=3.22E-05 on F(10,3) [1]} [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 3 / 4, -4 / 3], # 11 {E=1.09E-04 on F(11,3) [1]} [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 3 / 4, -4 / 3, 1 / 4], # 12 {E=1.99E-04 on F(12,3) [1]} [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 1 / 4, -3 / 4, 4 / 3, -4], # 13 {E=5.54E-04 on F(13,3) [1]} [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 1 / 4, -3 / 4, 4 / 3, 3 / 4, -4 / 3], # 14 {E=8.80E-04 on F(14,3) [1]} [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 1 / 4, -3 / 4, 4 / 3, -4, 3 / 4, -4 / 3], # 15 {E=1.07E-02 on F(15,3) [1]} [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 1 / 4, -3 / 4, 4 / 3, -4, 2 / 3, -3 / 2, 3 / 2], # 16 {E=1.93E-02 on F(16,3) [1]} [ 0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 1 / 4, -3 / 4, 4 / 3, -4, 2 / 3, -3 / 2, -2 / 3, 3 / 2, ], ] # pylint: enable=bad-whitespace,line-too-long return np.array(in_pts[degree - 1], dtype=np.float64) @memoize("topi.nn.winograd_matrices", save_at_exit=False) def winograd_transform_matrices(tile_size, kernel_size, out_dtype): """Compute the A, B, and G transform matrices for `tile_size` as a `tvm.Expr`.""" if not 1 < tile_size < 9: raise ValueError(f"Unsupported tile size for Winograd: {tile_size}") if not 2 < kernel_size < 8: raise ValueError(f"Unsupported kernel size for Winograd: {kernel_size}") degree = tile_size + kernel_size - 2 intp_pts = _interpolation_points(degree) A_data, B_data, G_data = _cook_toom_convolution(intp_pts, tile_size, kernel_size) out_dtype = "uint16" if out_dtype == "bfloat16" else out_dtype return ( const_matrix(A_data.astype(out_dtype), "A"), const_matrix(B_data.astype(out_dtype), "B"), const_matrix(G_data.astype(out_dtype), "G"), )