runtime/src/iree/builtins/ukernel/pack.c - 3p/openxla/iree - Git at Google

 // Copyright 2022 The IREE Authors
 //
 // Licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

 #include "iree/builtins/ukernel/pack_internal.h"

 enum { iree_uk_pack_tmp_buf_size = 4096 };

 // Holds some information and a temporary buffer for performing padding.
 typedef struct iree_uk_pack_tmpbuf_helper_t {
   // Temporary buffer to pad the source data into, to pass to the tile_func.
   // Cache line alignment helps in pack_benchmark on ARM Cortex-A510/A710.
   IREE_UK_ATTRIBUTE_ALIGNED(64) char tmp_buf[iree_uk_pack_tmp_buf_size];
   // How many tiles fit in `tmp_buf`.
   int max_tiles_in_tmp_buf;
   // Is the padding value as single-byte pattern (so we can use memset).
   bool is_padding_single_byte;
 } iree_uk_pack_tmpbuf_helper_t;

 // Return x/y for x>=0 and y>0, with a fast path for when y is a power of two.
 static iree_uk_ssize_t iree_uk_div_nonneg_by_pos_and_likely_po2_i32(
     iree_uk_ssize_t x, iree_uk_int32_t y) {
   IREE_UK_ASSERT(x >= 0);
   IREE_UK_ASSERT(y > 0);
   return IREE_UK_LIKELY(iree_uk_is_po2_u32(y)) ? (x >> iree_uk_po2_log2_u32(y))
                                                : (x / y);
 }

 // Returns true if the bytes in `bytes01234567` are all equal.
 static bool iree_uk_is_single_byte_pattern(iree_uk_uint64_t bytes01234567) {
   // Most common case of zero pattern.
   if (!bytes01234567) {
     return true;
   }
   iree_uk_uint32_t bytes0123 = bytes01234567;
   iree_uk_uint32_t bytes4567 = bytes01234567 >> 32;
   iree_uk_uint16_t bytes01 = bytes0123;
   iree_uk_uint16_t bytes23 = bytes0123 >> 16;
   iree_uk_uint8_t byte0 = bytes01;
   iree_uk_uint8_t byte1 = bytes01 >> 8;
   return (bytes0123 == bytes4567) && (bytes01 == bytes23) && (byte0 == byte1);
 }

 // Initializes a `iree_uk_pack_tmpbuf_helper_t`. Asserts if the temporary buffer
 // is smaller than one tile.
 static void iree_uk_pack_tmpbuf_helper_t_init(
     iree_uk_ssize_t tile_size0, iree_uk_ssize_t tile_size1,
     iree_uk_ssize_t elem_size, iree_uk_uint64_t padding_value,
     iree_uk_pack_tmpbuf_helper_t* helper) {
   helper->max_tiles_in_tmp_buf = iree_uk_div_nonneg_by_pos_and_likely_po2_i32(
       iree_uk_pack_tmp_buf_size, tile_size0 * tile_size1 * elem_size);
   IREE_UK_ASSERT(helper->max_tiles_in_tmp_buf > 0);
   helper->is_padding_single_byte =
       (elem_size == 1) || iree_uk_is_single_byte_pattern(padding_value);
 }

 static void iree_uk_pack_validate(const iree_uk_pack_params_t* params) {
 #ifdef IREE_UK_ENABLE_ASSERTS
   const iree_uk_uint32_t allflags = IREE_UK_FLAG_PACK_TRANSPOSE_INNER |
                                     IREE_UK_FLAG_PACK_TRANSPOSE_OUTER |
                                     IREE_UK_FLAG_PACK_TYPE_MASK;
   IREE_UK_ASSERT(!(params->flags & ~allflags));
   iree_uk_uint32_t flags_type = params->flags & IREE_UK_FLAG_PACK_TYPE_MASK;
   IREE_UK_ASSERT(flags_type == IREE_UK_FLAG_PACK_TYPE_F32F32 ||
                  flags_type == IREE_UK_FLAG_PACK_TYPE_I8I8 ||
                  flags_type == IREE_UK_FLAG_PACK_TYPE_I32I32);
   IREE_UK_ASSERT(params->in_stride0 >= 0);
   IREE_UK_ASSERT(params->out_stride0 >= 0);
   IREE_UK_ASSERT(params->in_size0 >= 0);
   IREE_UK_ASSERT(params->in_size1 >= 0);
   IREE_UK_ASSERT(params->out_size0 >= 0);
   IREE_UK_ASSERT(params->out_size1 >= 0);
   IREE_UK_ASSERT(params->out_size2 >= 0);
   IREE_UK_ASSERT(params->out_size3 >= 0);
   // Check that the input and output shapes match, give or take padding that
   // must not exceed the inner tile size.s
   iree_uk_ssize_t outer_size0 = params->out_size0;
   iree_uk_ssize_t outer_size1 = params->out_size1;
   iree_uk_ssize_t tile_size0 = params->out_size2;
   iree_uk_ssize_t tile_size1 = params->out_size3;
   if (params->flags & IREE_UK_FLAG_PACK_TRANSPOSE_OUTER) {
     iree_uk_ssize_swap(&outer_size0, &outer_size1);
   }
   if (params->flags & IREE_UK_FLAG_PACK_TRANSPOSE_INNER) {
     iree_uk_ssize_swap(&tile_size0, &tile_size1);
   }
   IREE_UK_ASSERT(outer_size0 * tile_size0 >= params->in_size0);
   IREE_UK_ASSERT(outer_size1 * tile_size1 >= params->in_size1);
   // TODO(#11632): reenable these conditions.
   // IREE_UK_ASSERT((outer_size0 - 1) * tile_size0 < params->in_size0);
   // IREE_UK_ASSERT((outer_size1 - 1) * tile_size1 < params->in_size1);

   // Initialize a padding helper, just to get the assertion that the tile size
   // does not exceed the internal temporary buffer size, without having to
   // duplicate this arithmetic. Generally, we want to hit all failure modes
   // in the validation function so that the subsequent ukernel code can be
   // treated as infallible.
   iree_uk_pack_tmpbuf_helper_t padding_helper;
   iree_uk_pack_type_t pack_type = iree_uk_pack_type(params->flags);
   iree_uk_type_t elem_type = iree_uk_pack_in_type(pack_type);
   iree_uk_ssize_t elem_size = iree_uk_type_size(elem_type);
   iree_uk_pack_tmpbuf_helper_t_init(tile_size0, tile_size1, elem_size,
                                     params->padding_value, &padding_helper);
 #endif  // IREE_UK_ENABLE_ASSERTS
 }

 // Early-return implementation for this ukernel. Returns true if already done.
 static bool iree_uk_pack_early(const iree_uk_pack_params_t* params) {
   return (params->out_size0 == 0 || params->out_size1 == 0 ||
           params->out_size2 == 0 || params->out_size3 == 0);
 }

 // Fills `buf` with `num_elems` times the `pattern` of size `elem_size`.
 // If this pattern's `elem_size` bytes are all equal, then it is legal to pass
 // `is_single_byte_pattern=true`, which allows the impl to use memset.
 static void iree_uk_fill(char* IREE_UK_RESTRICT buf, iree_uk_ssize_t num_elems,
                          iree_uk_ssize_t elem_size,
                          iree_uk_uint64_t padding_value,
                          bool is_padding_single_byte) {
   if (is_padding_single_byte) {
     iree_uk_memset(buf, padding_value & 0xFF, num_elems * elem_size);
   } else if (elem_size == 2) {
     iree_uk_uint16_t padding_value_uint16 = padding_value;
     iree_uk_uint16_t* IREE_UK_RESTRICT buf_uint16 = (iree_uk_uint16_t*)buf;
     for (iree_uk_ssize_t i = 0; i < num_elems; ++i) {
       buf_uint16[i] = padding_value_uint16;
     }
   } else if (elem_size == 4) {
     iree_uk_uint32_t padding_value_uint32 = padding_value;
     iree_uk_uint32_t* IREE_UK_RESTRICT buf_uint32 = (iree_uk_uint32_t*)buf;
     for (iree_uk_ssize_t i = 0; i < num_elems; ++i) {
       buf_uint32[i] = padding_value_uint32;
     }
   } else {  // elem_size >= 8
     // While arbitrary large elem_size is allowed, padding_value remains a
     // uint64, so elem_size >= 16 only support a repeating 8-byte pattern.
     iree_uk_uint64_t* IREE_UK_RESTRICT buf_uint64 = (iree_uk_uint64_t*)buf;
     for (iree_uk_ssize_t i = 0; i < num_elems * elem_size / 8; ++i) {
       buf_uint64[i] = padding_value;
     }
   }
 }

 // Copy from a source 2D buffer to a destination 2D buffer, padding to the
 // destination size.
 static void iree_uk_copy_and_pad(
     iree_uk_ssize_t src_size0, iree_uk_ssize_t src_size1,
     iree_uk_ssize_t src_stride0, const char* src_buf, iree_uk_ssize_t dst_size0,
     iree_uk_ssize_t dst_size1, iree_uk_ssize_t dst_stride0, char* dst_buf,
     iree_uk_ssize_t elem_size, iree_uk_uint64_t padding_value,
     bool is_padding_single_byte) {
   iree_uk_fill(dst_buf, dst_size1 + (dst_size0 - 1) * dst_stride0, elem_size,
                padding_value, is_padding_single_byte);
   for (iree_uk_ssize_t in_i0 = 0; in_i0 < src_size0; in_i0++) {
     iree_uk_memcpy(dst_buf, src_buf, src_size1 * elem_size);
     dst_buf += dst_stride0 * elem_size;
     src_buf += src_stride0 * elem_size;
   }
 }

 // Pads and packs an entire row. In cases that are known not to require padding,
 // it is more efficient to call tile_func directly.
 static void iree_uk_pad_and_pack_row_using_tile_func(
     iree_uk_pack_tile_func_t tile_func, iree_uk_ssize_t dim1_tile_start,
     iree_uk_ssize_t dim1_tile_end, iree_uk_ssize_t dim0_src_read_size,
     iree_uk_ssize_t tile_size0, iree_uk_ssize_t tile_size1,
     iree_uk_ssize_t elem_size, iree_uk_ssize_t in_size1,
     iree_uk_ssize_t in_stride0, iree_uk_ssize_t out_stride1,
     iree_uk_uint64_t padding_value, iree_uk_pack_tmpbuf_helper_t* helper,
     const char* in_buf, char* out_buf) {
   iree_uk_ssize_t dim1_tile = dim1_tile_start;
   while (dim1_tile < dim1_tile_end) {
     iree_uk_ssize_t dim1_chunk_tiles = iree_uk_ssize_clamp(
         dim1_tile_end - dim1_tile, 0, helper->max_tiles_in_tmp_buf);
     iree_uk_ssize_t dim1_chunk_src_width = dim1_chunk_tiles * tile_size1;
     iree_uk_ssize_t dim1_chunk_src_pos = dim1_tile * tile_size1;
     iree_uk_ssize_t i1_read_size = iree_uk_ssize_clamp(
         in_size1 - dim1_chunk_src_pos, 0, dim1_chunk_src_width);
     iree_uk_copy_and_pad(dim0_src_read_size, i1_read_size, in_stride0,
                          in_buf + dim1_chunk_src_pos * elem_size, tile_size0,
                          dim1_chunk_src_width, dim1_chunk_src_width,
                          helper->tmp_buf, elem_size, padding_value,
                          helper->is_padding_single_byte);
     tile_func(out_buf + (dim1_tile * out_stride1 * elem_size), helper->tmp_buf,
               dim1_chunk_tiles, out_stride1, dim1_chunk_src_width, elem_size,
               tile_size0, tile_size1);
     dim1_tile += dim1_chunk_tiles;
   }
 }

 static void iree_uk_pack_using_tile_func(const iree_uk_pack_params_t* params,
                                          iree_uk_pack_tile_func_t tile_func) {
   // For now, the input and output element types are always the same.
   iree_uk_pack_type_t pack_type = iree_uk_pack_type(params->flags);
   iree_uk_type_t elem_type = iree_uk_pack_in_type(pack_type);
   iree_uk_ssize_t elem_size = iree_uk_type_size(elem_type);
   iree_uk_ssize_t outer_size0 = params->out_size0;
   iree_uk_ssize_t outer_size1 = params->out_size1;
   iree_uk_ssize_t tile_size0 = params->out_size2;
   iree_uk_ssize_t tile_size1 = params->out_size3;
   iree_uk_ssize_t out_stride_l0 = params->out_stride0;
   iree_uk_ssize_t out_stride1 = params->out_size3 * params->out_size2;
   if (params->flags & IREE_UK_FLAG_PACK_TRANSPOSE_OUTER) {
     iree_uk_ssize_swap(&outer_size0, &outer_size1);
     iree_uk_ssize_swap(&out_stride_l0, &out_stride1);
   }
   if (params->flags & IREE_UK_FLAG_PACK_TRANSPOSE_INNER) {
     iree_uk_ssize_swap(&tile_size0, &tile_size1);
   }
   const char* in_buf =
       (const char*)params->in_buffer + (params->in_offset * elem_size);
   char* out_buf = (char*)params->out_buffer + (params->out_offset * elem_size);
   // Prepare for padding.
   iree_uk_pack_tmpbuf_helper_t padding_helper;
   if (params->in_size0 < outer_size0 * tile_size0 ||
       params->in_size1 < outer_size1 * tile_size1) {
     iree_uk_pack_tmpbuf_helper_t_init(tile_size0, tile_size1, elem_size,
                                       params->padding_value, &padding_helper);
   }
   // Compute number of tiles along both dimensions that fit entirely within the
   // source buffer's boundaries.
   int dim1_full_tiles = iree_uk_div_nonneg_by_pos_and_likely_po2_i32(
       params->in_size1, tile_size1);
   iree_uk_ssize_t i0 = 0;
   for (; i0 <= params->in_size0 - tile_size0; i0 += tile_size0) {
     // Pack whole tiles that do not require padding (entirely within the source
     // buffer's boundaries).
     tile_func(out_buf, in_buf, dim1_full_tiles, out_stride1, params->in_stride0,
               elem_size, tile_size0, tile_size1);
     // Right-padding.
     iree_uk_pad_and_pack_row_using_tile_func(
         tile_func, dim1_full_tiles, outer_size1, tile_size0, tile_size0,
         tile_size1, elem_size, params->in_size1, params->in_stride0,
         out_stride1, params->padding_value, &padding_helper, in_buf, out_buf);
     out_buf += out_stride_l0 * elem_size;
     in_buf += tile_size0 * params->in_stride0 * elem_size;
   }
   // Bottom-padding.
   for (; i0 < outer_size0 * tile_size0; i0 += tile_size0) {
     iree_uk_ssize_t dim0_src_read_size =
         iree_uk_ssize_clamp(params->in_size0 - i0, 0, tile_size0);
     iree_uk_pad_and_pack_row_using_tile_func(
         tile_func, 0, outer_size1, dim0_src_read_size, tile_size0, tile_size1,
         elem_size, params->in_size1, params->in_stride0, out_stride1,
         params->padding_value, &padding_helper, in_buf, out_buf);
     out_buf += out_stride_l0 * elem_size;
     in_buf += tile_size0 * params->in_stride0 * elem_size;
   }
 }

 IREE_UK_EXPORT int iree_uk_pack(const iree_uk_pack_params_t* params) {
   iree_uk_pack_validate(params);

   if (iree_uk_pack_early(params)) return 0;

   // Select a target-specific tile_func and use that with generic outer loops.
   iree_uk_pack_tile_func_t tile_func = iree_uk_pack_select_tile_func(params);
   iree_uk_pack_using_tile_func(params, tile_func);
   return 0;
 }
	// Copyright 2022 The IREE Authors
	//
	// Licensed under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

	#include "iree/builtins/ukernel/pack_internal.h"

	enum { iree_uk_pack_tmp_buf_size = 4096 };

	// Holds some information and a temporary buffer for performing padding.
	typedef struct iree_uk_pack_tmpbuf_helper_t {
	// Temporary buffer to pad the source data into, to pass to the tile_func.
	// Cache line alignment helps in pack_benchmark on ARM Cortex-A510/A710.
	IREE_UK_ATTRIBUTE_ALIGNED(64) char tmp_buf[iree_uk_pack_tmp_buf_size];
	// How many tiles fit in `tmp_buf`.
	int max_tiles_in_tmp_buf;
	// Is the padding value as single-byte pattern (so we can use memset).
	bool is_padding_single_byte;
	} iree_uk_pack_tmpbuf_helper_t;

	// Return x/y for x>=0 and y>0, with a fast path for when y is a power of two.
	static iree_uk_ssize_t iree_uk_div_nonneg_by_pos_and_likely_po2_i32(
	iree_uk_ssize_t x, iree_uk_int32_t y) {
	IREE_UK_ASSERT(x >= 0);
	IREE_UK_ASSERT(y > 0);
	return IREE_UK_LIKELY(iree_uk_is_po2_u32(y)) ? (x >> iree_uk_po2_log2_u32(y))
	: (x / y);
	}

	// Returns true if the bytes in `bytes01234567` are all equal.
	static bool iree_uk_is_single_byte_pattern(iree_uk_uint64_t bytes01234567) {
	// Most common case of zero pattern.
	if (!bytes01234567) {
	return true;
	}
	iree_uk_uint32_t bytes0123 = bytes01234567;
	iree_uk_uint32_t bytes4567 = bytes01234567 >> 32;
	iree_uk_uint16_t bytes01 = bytes0123;
	iree_uk_uint16_t bytes23 = bytes0123 >> 16;
	iree_uk_uint8_t byte0 = bytes01;
	iree_uk_uint8_t byte1 = bytes01 >> 8;
	return (bytes0123 == bytes4567) && (bytes01 == bytes23) && (byte0 == byte1);
	}

	// Initializes a `iree_uk_pack_tmpbuf_helper_t`. Asserts if the temporary buffer
	// is smaller than one tile.
	static void iree_uk_pack_tmpbuf_helper_t_init(
	iree_uk_ssize_t tile_size0, iree_uk_ssize_t tile_size1,
	iree_uk_ssize_t elem_size, iree_uk_uint64_t padding_value,
	iree_uk_pack_tmpbuf_helper_t* helper) {
	helper->max_tiles_in_tmp_buf = iree_uk_div_nonneg_by_pos_and_likely_po2_i32(
	iree_uk_pack_tmp_buf_size, tile_size0 * tile_size1 * elem_size);
	IREE_UK_ASSERT(helper->max_tiles_in_tmp_buf > 0);
	helper->is_padding_single_byte =
	(elem_size == 1) \|\| iree_uk_is_single_byte_pattern(padding_value);
	}

	static void iree_uk_pack_validate(const iree_uk_pack_params_t* params) {
	#ifdef IREE_UK_ENABLE_ASSERTS
	const iree_uk_uint32_t allflags = IREE_UK_FLAG_PACK_TRANSPOSE_INNER \|
	IREE_UK_FLAG_PACK_TRANSPOSE_OUTER \|
	IREE_UK_FLAG_PACK_TYPE_MASK;
	IREE_UK_ASSERT(!(params->flags & ~allflags));
	iree_uk_uint32_t flags_type = params->flags & IREE_UK_FLAG_PACK_TYPE_MASK;
	IREE_UK_ASSERT(flags_type == IREE_UK_FLAG_PACK_TYPE_F32F32 \|\|
	flags_type == IREE_UK_FLAG_PACK_TYPE_I8I8 \|\|
	flags_type == IREE_UK_FLAG_PACK_TYPE_I32I32);
	IREE_UK_ASSERT(params->in_stride0 >= 0);
	IREE_UK_ASSERT(params->out_stride0 >= 0);
	IREE_UK_ASSERT(params->in_size0 >= 0);
	IREE_UK_ASSERT(params->in_size1 >= 0);
	IREE_UK_ASSERT(params->out_size0 >= 0);
	IREE_UK_ASSERT(params->out_size1 >= 0);
	IREE_UK_ASSERT(params->out_size2 >= 0);
	IREE_UK_ASSERT(params->out_size3 >= 0);
	// Check that the input and output shapes match, give or take padding that
	// must not exceed the inner tile size.s
	iree_uk_ssize_t outer_size0 = params->out_size0;
	iree_uk_ssize_t outer_size1 = params->out_size1;
	iree_uk_ssize_t tile_size0 = params->out_size2;
	iree_uk_ssize_t tile_size1 = params->out_size3;
	if (params->flags & IREE_UK_FLAG_PACK_TRANSPOSE_OUTER) {
	iree_uk_ssize_swap(&outer_size0, &outer_size1);
	}
	if (params->flags & IREE_UK_FLAG_PACK_TRANSPOSE_INNER) {
	iree_uk_ssize_swap(&tile_size0, &tile_size1);
	}
	IREE_UK_ASSERT(outer_size0 * tile_size0 >= params->in_size0);
	IREE_UK_ASSERT(outer_size1 * tile_size1 >= params->in_size1);
	// TODO(#11632): reenable these conditions.
	// IREE_UK_ASSERT((outer_size0 - 1) * tile_size0 < params->in_size0);
	// IREE_UK_ASSERT((outer_size1 - 1) * tile_size1 < params->in_size1);

	// Initialize a padding helper, just to get the assertion that the tile size
	// does not exceed the internal temporary buffer size, without having to
	// duplicate this arithmetic. Generally, we want to hit all failure modes
	// in the validation function so that the subsequent ukernel code can be
	// treated as infallible.
	iree_uk_pack_tmpbuf_helper_t padding_helper;
	iree_uk_pack_type_t pack_type = iree_uk_pack_type(params->flags);
	iree_uk_type_t elem_type = iree_uk_pack_in_type(pack_type);
	iree_uk_ssize_t elem_size = iree_uk_type_size(elem_type);
	iree_uk_pack_tmpbuf_helper_t_init(tile_size0, tile_size1, elem_size,
	params->padding_value, &padding_helper);
	#endif // IREE_UK_ENABLE_ASSERTS
	}

	// Early-return implementation for this ukernel. Returns true if already done.
	static bool iree_uk_pack_early(const iree_uk_pack_params_t* params) {
	return (params->out_size0 == 0 \|\| params->out_size1 == 0 \|\|
	params->out_size2 == 0 \|\| params->out_size3 == 0);
	}

	// Fills `buf` with `num_elems` times the `pattern` of size `elem_size`.
	// If this pattern's `elem_size` bytes are all equal, then it is legal to pass
	// `is_single_byte_pattern=true`, which allows the impl to use memset.
	static void iree_uk_fill(char* IREE_UK_RESTRICT buf, iree_uk_ssize_t num_elems,
	iree_uk_ssize_t elem_size,
	iree_uk_uint64_t padding_value,
	bool is_padding_single_byte) {
	if (is_padding_single_byte) {
	iree_uk_memset(buf, padding_value & 0xFF, num_elems * elem_size);
	} else if (elem_size == 2) {
	iree_uk_uint16_t padding_value_uint16 = padding_value;
	iree_uk_uint16_t* IREE_UK_RESTRICT buf_uint16 = (iree_uk_uint16_t*)buf;
	for (iree_uk_ssize_t i = 0; i < num_elems; ++i) {
	buf_uint16[i] = padding_value_uint16;
	}
	} else if (elem_size == 4) {
	iree_uk_uint32_t padding_value_uint32 = padding_value;
	iree_uk_uint32_t* IREE_UK_RESTRICT buf_uint32 = (iree_uk_uint32_t*)buf;
	for (iree_uk_ssize_t i = 0; i < num_elems; ++i) {
	buf_uint32[i] = padding_value_uint32;
	}
	} else { // elem_size >= 8
	// While arbitrary large elem_size is allowed, padding_value remains a
	// uint64, so elem_size >= 16 only support a repeating 8-byte pattern.
	iree_uk_uint64_t* IREE_UK_RESTRICT buf_uint64 = (iree_uk_uint64_t*)buf;
	for (iree_uk_ssize_t i = 0; i < num_elems * elem_size / 8; ++i) {
	buf_uint64[i] = padding_value;
	}
	}
	}

	// Copy from a source 2D buffer to a destination 2D buffer, padding to the
	// destination size.
	static void iree_uk_copy_and_pad(
	iree_uk_ssize_t src_size0, iree_uk_ssize_t src_size1,
	iree_uk_ssize_t src_stride0, const char* src_buf, iree_uk_ssize_t dst_size0,
	iree_uk_ssize_t dst_size1, iree_uk_ssize_t dst_stride0, char* dst_buf,
	iree_uk_ssize_t elem_size, iree_uk_uint64_t padding_value,
	bool is_padding_single_byte) {
	iree_uk_fill(dst_buf, dst_size1 + (dst_size0 - 1) * dst_stride0, elem_size,
	padding_value, is_padding_single_byte);
	for (iree_uk_ssize_t in_i0 = 0; in_i0 < src_size0; in_i0++) {
	iree_uk_memcpy(dst_buf, src_buf, src_size1 * elem_size);
	dst_buf += dst_stride0 * elem_size;
	src_buf += src_stride0 * elem_size;
	}
	}

	// Pads and packs an entire row. In cases that are known not to require padding,
	// it is more efficient to call tile_func directly.
	static void iree_uk_pad_and_pack_row_using_tile_func(
	iree_uk_pack_tile_func_t tile_func, iree_uk_ssize_t dim1_tile_start,
	iree_uk_ssize_t dim1_tile_end, iree_uk_ssize_t dim0_src_read_size,
	iree_uk_ssize_t tile_size0, iree_uk_ssize_t tile_size1,
	iree_uk_ssize_t elem_size, iree_uk_ssize_t in_size1,
	iree_uk_ssize_t in_stride0, iree_uk_ssize_t out_stride1,
	iree_uk_uint64_t padding_value, iree_uk_pack_tmpbuf_helper_t* helper,
	const char* in_buf, char* out_buf) {
	iree_uk_ssize_t dim1_tile = dim1_tile_start;
	while (dim1_tile < dim1_tile_end) {
	iree_uk_ssize_t dim1_chunk_tiles = iree_uk_ssize_clamp(
	dim1_tile_end - dim1_tile, 0, helper->max_tiles_in_tmp_buf);
	iree_uk_ssize_t dim1_chunk_src_width = dim1_chunk_tiles * tile_size1;
	iree_uk_ssize_t dim1_chunk_src_pos = dim1_tile * tile_size1;
	iree_uk_ssize_t i1_read_size = iree_uk_ssize_clamp(
	in_size1 - dim1_chunk_src_pos, 0, dim1_chunk_src_width);
	iree_uk_copy_and_pad(dim0_src_read_size, i1_read_size, in_stride0,
	in_buf + dim1_chunk_src_pos * elem_size, tile_size0,
	dim1_chunk_src_width, dim1_chunk_src_width,
	helper->tmp_buf, elem_size, padding_value,
	helper->is_padding_single_byte);
	tile_func(out_buf + (dim1_tile * out_stride1 * elem_size), helper->tmp_buf,
	dim1_chunk_tiles, out_stride1, dim1_chunk_src_width, elem_size,
	tile_size0, tile_size1);
	dim1_tile += dim1_chunk_tiles;
	}
	}

	static void iree_uk_pack_using_tile_func(const iree_uk_pack_params_t* params,
	iree_uk_pack_tile_func_t tile_func) {
	// For now, the input and output element types are always the same.
	iree_uk_pack_type_t pack_type = iree_uk_pack_type(params->flags);
	iree_uk_type_t elem_type = iree_uk_pack_in_type(pack_type);
	iree_uk_ssize_t elem_size = iree_uk_type_size(elem_type);
	iree_uk_ssize_t outer_size0 = params->out_size0;
	iree_uk_ssize_t outer_size1 = params->out_size1;
	iree_uk_ssize_t tile_size0 = params->out_size2;
	iree_uk_ssize_t tile_size1 = params->out_size3;
	iree_uk_ssize_t out_stride_l0 = params->out_stride0;
	iree_uk_ssize_t out_stride1 = params->out_size3 * params->out_size2;
	if (params->flags & IREE_UK_FLAG_PACK_TRANSPOSE_OUTER) {
	iree_uk_ssize_swap(&outer_size0, &outer_size1);
	iree_uk_ssize_swap(&out_stride_l0, &out_stride1);
	}
	if (params->flags & IREE_UK_FLAG_PACK_TRANSPOSE_INNER) {
	iree_uk_ssize_swap(&tile_size0, &tile_size1);
	}
	const char* in_buf =
	(const char)params->in_buffer + (params->in_offset elem_size);
	char* out_buf = (char)params->out_buffer + (params->out_offset elem_size);
	// Prepare for padding.
	iree_uk_pack_tmpbuf_helper_t padding_helper;
	if (params->in_size0 < outer_size0 * tile_size0 \|\|
	params->in_size1 < outer_size1 * tile_size1) {
	iree_uk_pack_tmpbuf_helper_t_init(tile_size0, tile_size1, elem_size,
	params->padding_value, &padding_helper);
	}
	// Compute number of tiles along both dimensions that fit entirely within the
	// source buffer's boundaries.
	int dim1_full_tiles = iree_uk_div_nonneg_by_pos_and_likely_po2_i32(
	params->in_size1, tile_size1);
	iree_uk_ssize_t i0 = 0;
	for (; i0 <= params->in_size0 - tile_size0; i0 += tile_size0) {
	// Pack whole tiles that do not require padding (entirely within the source
	// buffer's boundaries).
	tile_func(out_buf, in_buf, dim1_full_tiles, out_stride1, params->in_stride0,
	elem_size, tile_size0, tile_size1);
	// Right-padding.
	iree_uk_pad_and_pack_row_using_tile_func(
	tile_func, dim1_full_tiles, outer_size1, tile_size0, tile_size0,
	tile_size1, elem_size, params->in_size1, params->in_stride0,
	out_stride1, params->padding_value, &padding_helper, in_buf, out_buf);
	out_buf += out_stride_l0 * elem_size;
	in_buf += tile_size0 * params->in_stride0 * elem_size;
	}
	// Bottom-padding.
	for (; i0 < outer_size0 * tile_size0; i0 += tile_size0) {
	iree_uk_ssize_t dim0_src_read_size =
	iree_uk_ssize_clamp(params->in_size0 - i0, 0, tile_size0);
	iree_uk_pad_and_pack_row_using_tile_func(
	tile_func, 0, outer_size1, dim0_src_read_size, tile_size0, tile_size1,
	elem_size, params->in_size1, params->in_stride0, out_stride1,
	params->padding_value, &padding_helper, in_buf, out_buf);
	out_buf += out_stride_l0 * elem_size;
	in_buf += tile_size0 * params->in_stride0 * elem_size;
	}
	}

	IREE_UK_EXPORT int iree_uk_pack(const iree_uk_pack_params_t* params) {
	iree_uk_pack_validate(params);

	if (iree_uk_pack_early(params)) return 0;

	// Select a target-specific tile_func and use that with generic outer loops.
	iree_uk_pack_tile_func_t tile_func = iree_uk_pack_select_tile_func(params);
	iree_uk_pack_using_tile_func(params, tile_func);
	return 0;
	}