wehub-resource-sync
182 lines
6.1 KiB
Python
182 lines
6.1 KiB
Python
# 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"),
|
|
)
|