| # 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. |
| # pylint: disable=invalid-name |
| |
| """Hexagon testing infrastructure""" |
| |
| import numpy |
| import tvm |
| from tvm import te |
| |
| |
| def ceildiv(o, d): |
| assert o >= 0 |
| assert d >= 0 |
| return tvm.tir.floordiv(o + d - 1, d) |
| |
| |
| # defines inner block shape: 8h8w32c |
| def get_block_shape(): |
| return 8, 8, 32 |
| |
| |
| # defines inner filter block shape: 8i32o41 |
| def get_filter_block_shape(): |
| return 8, 32, 4 |
| |
| |
| # input: locgical shape in nhwc layout |
| # output: physical packed shape in nhw8h8w32c layout |
| def get_packed_shape(logical_shape_nhwc): |
| assert len(logical_shape_nhwc) == 4 |
| physical_shape_nhwc8h8w32c = [logical_shape_nhwc[0]] |
| block_shape = get_block_shape() |
| off_h, off_w, off_c = block_shape |
| physical_shape_nhwc8h8w32c.append(ceildiv(logical_shape_nhwc[1], off_h)) |
| physical_shape_nhwc8h8w32c.append(ceildiv(logical_shape_nhwc[2], off_w)) |
| physical_shape_nhwc8h8w32c.append(ceildiv(logical_shape_nhwc[3], off_c)) |
| physical_shape_nhwc8h8w32c.extend(block_shape) |
| return physical_shape_nhwc8h8w32c |
| |
| |
| # input: physical packed shape in nhw8h8w32c layout |
| # output: logical shape in nhwc layout |
| def get_logical_shape(physical_shape_nhwc8h8w32c): |
| assert len(physical_shape_nhwc8h8w32c) == 7 |
| logical_shape_nhwc = [physical_shape_nhwc8h8w32c[0]] |
| logical_shape_nhwc.append(physical_shape_nhwc8h8w32c[1] * physical_shape_nhwc8h8w32c[4]) |
| logical_shape_nhwc.append(physical_shape_nhwc8h8w32c[2] * physical_shape_nhwc8h8w32c[5]) |
| logical_shape_nhwc.append(physical_shape_nhwc8h8w32c[3] * physical_shape_nhwc8h8w32c[6]) |
| return logical_shape_nhwc |
| |
| |
| def get_packed_filter_shape(logical_shape_oihw): |
| """return packed filter shape |
| |
| Parameters |
| ---------- |
| logical_shape_oihw : |
| logical shape in oihw layout |
| |
| Returns |
| ------- |
| physical_shape_oihw8i32o4i : |
| physical packed shape in oihw8i3204i layout |
| """ |
| assert len(logical_shape_oihw) == 4 |
| filter_block_shape = get_filter_block_shape() |
| filter_Cio, filter_Ki, filter_Cii = filter_block_shape |
| filter_Ci = filter_Cio * filter_Cii |
| physical_shape_oihw8i32o4i = [] |
| physical_shape_oihw8i32o4i.append(int(ceildiv(logical_shape_oihw[0], filter_Ki))) |
| physical_shape_oihw8i32o4i.append(int(ceildiv(logical_shape_oihw[1], filter_Ci))) |
| physical_shape_oihw8i32o4i.append(logical_shape_oihw[2]) |
| physical_shape_oihw8i32o4i.append(logical_shape_oihw[3]) |
| physical_shape_oihw8i32o4i.extend(filter_block_shape) |
| return physical_shape_oihw8i32o4i |
| |
| |
| def build_and_run(inputs, func, target: str, target_host: str, *args, **kwargs): |
| """build and run the function func""" |
| schedule, placeholders, binds = func(*args, **kwargs) |
| |
| func = tvm.compile( |
| schedule, placeholders, target=tvm.target.Target(target, host=target_host), binds=binds |
| ) |
| dev = tvm.device(target) |
| tensors = [] |
| for tensor in inputs: |
| tensors.append(tvm.runtime.tensor(tensor, dev)) |
| tensors.append( |
| tvm.runtime.tensor( |
| numpy.zeros([i.value for i in placeholders[-1].shape], dtype=placeholders[-1].dtype), |
| dev, |
| ) |
| ) |
| func(*tensors) |
| |
| return tensors[-1].numpy() |
| |
| |
| def run_module(mod, inputs): |
| """invokes run function of specified module with inputs provided""" |
| mod.set_input(**inputs) |
| mod.run() |
| output = mod.get_output(0).numpy() |
| return output |
| |
| |
| def get_conv2d_nhwc_shape(shape_nhwc, kernel_size, strides, padding, dilation, out_channels): |
| assert len(shape_nhwc) == 4 |
| kernel = [] |
| kernel.append((kernel_size[0] - 1) * dilation[0] + 1) |
| kernel.append((kernel_size[1] - 1) * dilation[1] + 1) |
| return ( |
| shape_nhwc[0], |
| (shape_nhwc[1] - kernel[0] + padding[0] + padding[1]) // strides[0] + 1, |
| (shape_nhwc[2] - kernel[1] + padding[2] + padding[3]) // strides[1] + 1, |
| out_channels, |
| ) |
| |
| |
| def conv2d_verify(output, ref_output, dtype): |
| """transpose and reshape output and compare with ref_output""" |
| # nhwc8h8w32c -> nhwc |
| logical_output_shape = get_logical_shape(output.shape) |
| output = output.transpose(0, 1, 4, 2, 5, 3, 6).reshape(logical_output_shape) |
| |
| # slice output to match ref_output shape |
| # e.g. 8x8 spatial 3x3 filter = 6x6 ref output |
| # but still 8x8 output given the blocked layout |
| output = output[ |
| 0 : ref_output.shape[0] : 1, |
| 0 : ref_output.shape[1] : 1, |
| 0 : ref_output.shape[2] : 1, |
| 0 : ref_output.shape[3] : 1, |
| ] |
| |
| if "int" in dtype: |
| tol = {"atol": 0, "rtol": 0} |
| elif dtype == "float32": |
| tol = {"rtol": 1e-4, "atol": 2e-4} |
| tvm.testing.assert_allclose(output, ref_output, **tol) |
| |
| |
| def conv2d_compute(X, filt, pad, stride, dilation): |
| """Define conv2d compute""" |
| block_shape = get_block_shape() |
| block_H, block_W, block_C = block_shape |
| filter_c_io, _, filter_c_ii = get_filter_block_shape() |
| filter_c_i = filter_c_io * filter_c_ii |
| |
| shape_filter = filt.shape |
| kernel_size = tuple(shape_filter[2:4]) |
| out_channels = shape_filter[0] * shape_filter[5] |
| |
| logical_input_shape = get_logical_shape(X.shape) |
| logical_output_shape = get_conv2d_nhwc_shape( |
| logical_input_shape, |
| kernel_size, |
| stride, |
| pad, |
| dilation, |
| out_channels, |
| ) |
| |
| output_shape = get_packed_shape(logical_output_shape) |
| rh = te.reduce_axis((0, kernel_size[0]), name="rh") |
| rw = te.reduce_axis((0, kernel_size[1]), name="rw") |
| rc = te.reduce_axis((0, logical_input_shape[3]), name="rc") |
| |
| def compute(n, ho, wo, ko, hi, wi, ki): |
| h = ho * block_H + hi |
| h_contig = h * stride[0] + rh |
| h_block_id = h_contig // block_H |
| h_block_offset = h_contig % block_H |
| |
| w = wo * block_W + wi |
| w_contig = w * stride[1] + rw |
| w_block_id = w_contig // block_W |
| w_block_offset = w_contig % block_W |
| |
| c_block_id = rc // block_C |
| c_block_offset = rc % block_C |
| |
| rco = rc // filter_c_i |
| rcio = (rc % filter_c_i) // filter_c_ii |
| rcii = rc % filter_c_ii |
| |
| return te.sum( |
| X[ |
| n, |
| h_block_id, |
| w_block_id, |
| c_block_id, |
| h_block_offset, |
| w_block_offset, |
| c_block_offset, |
| ] |
| * filt[ko, rco, rh, rw, rcio, ki, rcii], |
| axis=[rh, rw, rc], |
| ) |
| |
| return output_shape, compute |
| |
| |
| def transform_numpy(arr_np, current_layout: str, new_layout: str): |
| """Reshape and transpose numpy array according to the specified layout""" |
| if current_layout == "nhwc": |
| if new_layout == "nhwc": |
| return arr_np |
| if new_layout in ["nhwc-8h2w32c2w-2d", "nhwc-8h2w32c2w-1d"]: |
| n, h, w, c = arr_np.shape |
| return arr_np.reshape([n, h // 8, 8, w // 4, 2, 2, c // 32, 32]).transpose( |
| 0, 1, 3, 6, 2, 4, 7, 5 |
| ) |
| if new_layout in ["nhwc-4h2w32c2w-2d"]: |
| n, h, w, c = arr_np.shape |
| return arr_np.reshape([n, h // 4, 4, w // 4, 2, 2, c // 32, 32]).transpose( |
| 0, 1, 3, 6, 2, 4, 7, 5 |
| ) |
| if new_layout in ["n11c-1024c-2d", "n11c-1024c-1d"]: |
| n, h, w, c = arr_np.shape |
| assert h == 1 and w == 1, "The size of h and w must be 1" |
| return arr_np.reshape([n, 1, 1, c // 1024, 1024]) |
| if new_layout == "nc-1024-2d": |
| n, c = arr_np.shape |
| return arr_np.reshape([n, c // 1024, 1024]) |
| if new_layout == "nhwc-1024c-2d": |
| N, H, W, C = arr_np.shape |
| return arr_np.reshape([N, H, W, C // 1024, 1024]) |
| if new_layout == "nc-2048-2d": |
| N, C = arr_np.shape |
| return arr_np.reshape([N, C // 2048, 2048]) |
| if new_layout == "nhwc-2048c-2d": |
| N, H, W, C = arr_np.shape |
| return arr_np.reshape([N, H, W, C // 2048, 2048]) |
| if new_layout == "nhwc-8h8w32c-2d": |
| n, h, w, c = arr_np.shape |
| return arr_np.reshape([n, h // 8, 8, w // 8, 8, c // 32, 32]).transpose( |
| 0, 1, 3, 5, 2, 4, 6 |
| ) |
| if new_layout == "n11c-2048c-2d": |
| n, h, w, c = arr_np.shape |
| assert h == 1 and w == 1, "The size of h and w must be 1" |
| return arr_np.reshape([n, h, w, c // 2048, 2048]) |
| raise RuntimeError(f"Unexpected new_layout '{new_layout}'") |
| |
| if current_layout == "nc": |
| n, c = arr_np.shape |
| if new_layout in ["nc-2048c-1d"]: |
| return arr_np.reshape([n, c // 2048, 2048]) |
| if new_layout in ["nc-2048c-2d"]: |
| return arr_np.reshape([n, c // 2048, 2048]) |
| if new_layout in ["nc-1024c-2d"]: |
| return arr_np.reshape([n, c // 1024, 1024]) |
| if new_layout in ["nc-1024c-1d"]: |
| return arr_np.reshape([n, c // 1024, 1024]) |
| if new_layout in ["nc-512c-2d"]: |
| return arr_np.reshape([n, c // 512, 512]) |
| if new_layout in ["nc-2048c-2d"]: |
| return arr_np.reshape([n, c // 2048, 2048]) |
| raise RuntimeError(f"Unexpected new_layout '{new_layout}'") |
| |
| if current_layout == "nhw": |
| if new_layout in ["nhw-32h16w-2d"]: |
| n, h, w = arr_np.shape |
| return arr_np.reshape([n, h // 32, 32, w // 16, 16]).transpose(0, 1, 3, 2, 4) |
| |
| raise RuntimeError(f"Unexpected new_layout '{new_layout}'") |
| |
| if current_layout == "ncw": |
| if new_layout == "ncw": |
| return arr_np |
| if new_layout in ["ncw-32c64w-2d"]: |
| n, c, w = arr_np.shape |
| return arr_np.reshape([n, c // 32, 32, w // 64, 64]).transpose(0, 1, 3, 2, 4) |
| |
| raise RuntimeError(f"Unexpected new_layout '{new_layout}'") |
| |
| if current_layout == "nchw": |
| if new_layout in ["nchw-32c8h8w-2d", "nchw-32c8h8w-1d"]: |
| n, c, h, w = arr_np.shape |
| return arr_np.reshape([n, c // 32, 32, h // 8, 8, w // 8, 8]).transpose( |
| 0, 1, 3, 5, 2, 4, 6 |
| ) |
| if new_layout in ["nchw-32c8h4w-2d", "nchw-32c8h4w-1d"]: |
| n, c, h, w = arr_np.shape |
| return arr_np.reshape([n, c // 32, 32, h // 8, 8, w // 4, 4]).transpose( |
| 0, 1, 3, 5, 2, 4, 6 |
| ) |
| raise RuntimeError(f"Unexpected new_layout '{new_layout}'") |
| |
| raise RuntimeError(f"Unexpected current_layout '{current_layout}'") |
| |
| |
| def quantize_np(arr_np: numpy.ndarray, dtype: str): |
| """ |
| Returns quantized array along with scale and zero-point |
| |
| Parameters |
| ---------- |
| arr_np: numpy.ndarray |
| Input numpy array to be quantized |
| dtype: str |
| dtype of the quantized array: "uint8", "int8", etc |
| |
| Returns |
| ------- |
| quant_np: numpy.ndarray |
| Quantized numpy array |
| scale: float |
| Scale |
| zero_point: int |
| Value corresponding to float 0 |
| |
| """ |
| if dtype == "uint8": |
| qmax = 255 |
| qmin = 0 |
| elif dtype == "int8": |
| qmax = 127 |
| qmin = -128 |
| else: |
| raise RuntimeError(f"Unsupported quantized data type '{dtype}'") |
| fmin = numpy.amin(arr_np) |
| fmax = numpy.amax(arr_np) |
| |
| # Include floating-point zero in the range |
| if fmax < 0: |
| fmax = 0.0 |
| elif fmin > 0: |
| fmin = 0.0 |
| |
| scale = (fmax - fmin) / (qmax - qmin) |
| zero_point = numpy.rint((fmax * qmin - fmin * qmax) / (fmax - fmin)).astype("int32") |
| quant_np = numpy.clip(((arr_np / scale).round() + zero_point), qmin, qmax).astype(dtype) |
| return quant_np, scale, zero_point |
| |
| |
| def get_hexagon_target(cpu_ver: str, **kwargs) -> tvm.target.Target: |
| """Creates a Hexagon target""" |
| target = tvm.target.hexagon(cpu_ver, **kwargs) |
| return tvm.target.Target(target, host=target) |
