Data transformation¶
General¶
The transform ukernel allows users to convert data from one format to the other, similar to what reorder primitive provides functionally.
The only output data format supported by this routine is packed format, which is required by B matrices in BRGeMM ukernel. This is an out-of-place operation.
Data Types¶
The transform ukernel does not allow data type conversion.
Data Representation¶
src |
dst |
|---|---|
f32 |
f32 |
f16 |
f16 |
bf16 |
bf16 |
s8 |
s8 |
u8 |
u8 |
Attributes¶
No attribute is supported for transform ukernel.
Implementation limitations¶
Destination leading dimension only supported values are: 16, 32, 48, or 64.
Examples¶
This C++ API example demonstrates how to create and execute a BRGeMM ukernel.
/******************************************************************************* * Copyright 2024 Intel Corporation * * Licensed 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. *******************************************************************************/ #include <algorithm> #include <cmath> #include <iostream> #include <string> #include <utility> #include <vector> #include "example_utils.hpp" #include "oneapi/dnnl/dnnl_ukernel.hpp" using namespace dnnl; using namespace dnnl::ukernel; using tag = memory::format_tag; using dt = memory::data_type; void brgemm_example() { // Create execution dnnl::engine. Needed for reorders to operate over input // data. dnnl::engine engine(engine::kind::cpu, 0); // Create dnnl::stream. Needed for reorders for the same reason. dnnl::stream engine_stream(engine); // ukernel dimensions. // K is for a whole tensor, K_k is for a single ukernel. const memory::dim M = 8, K = 128, K_k = 64, N = 48; if (K % K_k != 0) { printf("K_k must divide K.\n"); return; } const memory::dim n_calls = K / K_k; const memory::dim lda = K; const memory::dim ldb = N; const memory::dim ldc = N; // Leading dimension for accumulator. const memory::dim ldd = N; // Leading dimension for an actual output. const memory::dim batch_size = n_calls - 1; memory::data_type a_dt = dt::u8; memory::data_type b_dt = dt::s8; memory::data_type c_dt = dt::s32; // Accumulator data type. memory::data_type d_dt = dt::f32; // Output data type. // A, B, and C tensors dimensions. memory::dims A_dims = {M, K}; memory::dims B_dims = {K, N}; memory::dims C_dims = {M, N}; memory::dims D_dims = {M, N}; memory::dims binary_add_dims = {1, 1}; memory::dims B_scales_dims = {1, N}; // Allocate buffers with user data. std::vector<float> A_user_data(product(A_dims)); std::vector<float> B_user_data(product(B_dims)); std::vector<float> binary_add_user_data(product(binary_add_dims)); std::vector<float> B_scales_user_data(product(B_scales_dims)); std::vector<float> D_data(product(D_dims)); // For reference comparison std::vector<float> D_user_data(product(D_dims)); // For reference comparison // Initialize A. std::generate(A_user_data.begin(), A_user_data.end(), []() { static int i = 0; return i++ % 4; }); // Initialize B. std::generate(B_user_data.begin(), B_user_data.end(), []() { static int i = 6; static int sign_gen = 0; int sign = (sign_gen++ % 2) ? -1 : 1; float val = sign * (i++ % 5); return val; }); // Initialize binary_add. std::generate( binary_add_user_data.begin(), binary_add_user_data.end(), []() { static int i = 3; return i++ % 6; }); // Initialize B scales. std::generate(B_scales_user_data.begin(), B_scales_user_data.end(), []() { static int i = 4; return (float)(i++ % 16) / 8.f; }); // Create f32 memories. They are used as data holders and reorder into // memories passed to the ukernel. auto A_f32_md = memory::desc(A_dims, dt::f32, tag::ab); auto B_f32_md = memory::desc(B_dims, dt::f32, tag::ab); auto binary_add_f32_md = memory::desc(binary_add_dims, dt::f32, tag::ab); auto B_scales_f32_md = memory::desc(B_scales_dims, dt::f32, tag::ab); auto D_f32_md = memory::desc(D_dims, dt::f32, tag::ab); auto A_f32_mem = memory(A_f32_md, engine, A_user_data.data()); auto B_f32_mem = memory(B_f32_md, engine, B_user_data.data()); auto binary_add_f32_mem = memory(binary_add_f32_md, engine, binary_add_user_data.data()); auto B_scales_f32_mem = memory(B_scales_f32_md, engine, B_scales_user_data.data()); auto D_f32_mem = memory(D_f32_md, engine, D_user_data.data()); // Create ukernel memories in requested data types. // Note that all formats are `ab`. auto A_md = memory::desc(A_dims, a_dt, tag::ab); auto B_md = memory::desc(B_dims, b_dt, tag::ab); auto binary_add_md = memory::desc(binary_add_dims, dt::f32, tag::ab); auto B_scales_md = memory::desc(B_scales_dims, dt::f32, tag::ab); auto C_md = memory::desc(C_dims, c_dt, tag::ab); auto D_md = memory::desc(D_dims, d_dt, tag::ab); auto A_mem = memory(A_md, engine); auto B_mem = memory(B_md, engine); auto binary_add_mem = memory(binary_add_md, engine); auto B_scales_mem = memory(B_scales_md, engine); auto C_mem = memory(C_md, engine); auto D_mem = memory(D_md, engine); const auto *A_ptr = reinterpret_cast<uint8_t *>(A_mem.get_data_handle()); auto *B_ptr = reinterpret_cast<uint8_t *>(B_mem.get_data_handle()); const size_t a_dt_size = memory::data_type_size(A_mem.get_desc().get_data_type()); const size_t b_dt_size = memory::data_type_size(B_mem.get_desc().get_data_type()); // Reorder user data into buffers passed to ukernels in target data types. reorder(A_f32_mem, A_mem).execute(engine_stream, A_f32_mem, A_mem); reorder(B_f32_mem, B_mem).execute(engine_stream, B_f32_mem, B_mem); reorder(binary_add_f32_mem, binary_add_mem) .execute(engine_stream, binary_add_f32_mem, binary_add_mem); reorder(B_scales_f32_mem, B_scales_mem) .execute(engine_stream, B_scales_f32_mem, B_scales_mem); reorder(D_f32_mem, D_mem).execute(engine_stream, D_f32_mem, D_mem); // Prepare C buffer. Needed to use a single ukernel in the example with // `beta = 1.f`. // Note: to avoid this step, the first ukernel should run `beta = 0`, and it // will initialize C buffer with intermediate values. float *C_ptr = reinterpret_cast<float *>(C_mem.get_data_handle()); for (memory::dim i = 0; i < M * N; i++) { C_ptr[i] = 0; } // Create ukernel post-ops (ReLU + Add). // It reuses `primitive_attr` abstraction. post_ops brgemm_ops; brgemm_ops.append_eltwise( algorithm::eltwise_relu, /* alpha = */ 0.f, /* beta = */ 0.f); brgemm_ops.append_binary(algorithm::binary_add, binary_add_md); // Create BRGeMM ukernel objects. // There are two objects: // * `brg` is the main one which operates over partitioned K dimension. It // utilizes `beta = 1.f` to accumulate into the same buffer. It also uses // `batch_size` to process as much as `n_calls - 1` iterations. // * `brg_po` is the ukernel that would be called the last in the chain // since it has attributes attached to the object and those will execute // after all accumulation over K dimension is done. // Note: `beta = 1.f` makes a ukernel reusable over K but will require // zeroing the correspondent piece of accumulation buffer. brgemm brg, brg_po; if (batch_size > 0) { try { // Construct a basic brgemm object. brg = brgemm( M, N, K_k, batch_size, lda, ldb, ldc, a_dt, b_dt, c_dt); // Instruct the kernel to append the result to C tensor. brg.set_add_C(true); // Finalize the initialization. brg.finalize(); // Generate the executable JIT code for the objects. brg.generate(); } catch (error &e) { if (e.status == dnnl_unimplemented) throw example_allows_unimplemented { "Kernel is not supported on this platform.\n"}; // on any other error just re-throw throw; } } try { // Construct a basic brgemm object. brg_po = brgemm(M, N, K_k, 1, lda, ldb, ldc, a_dt, b_dt, c_dt); // Instruct the kernel to append the result to C tensor. brg_po.set_add_C(true); // Specify post-ops for the brgemm object. brg_po.set_post_ops(ldd, d_dt, brgemm_ops); // Specify quantization scales for B. if (b_dt == dt::s8 || b_dt == dt::u8) { brg_po.set_B_scales(/* mask = */ 2); } // Finalize the initialization. brg_po.finalize(); // Generate the executable JIT code for the objects. brg_po.generate(); } catch (error &e) { if (e.status == dnnl_unimplemented) throw example_allows_unimplemented { "Kernel is not supported on this platform.\n"}; // on any other error just re-throw throw; } // Query a scratchpad size and initialize a scratchpad buffer if the ukernel // is expecting it. This is a service space needed, has nothing in common // with accumulation buffer. size_t scratchpad_size = brg_po.get_scratchpad_size(); std::vector<uint8_t> scratchpad(scratchpad_size); uint8_t *B_blocked = nullptr; void *B_base_ptr = B_ptr; size_t blocked_B_size = 0; // Query the packing requirement from the kernel. It's enough to query // packing requirements from a single object as long as only dimension // settings change between objects. // Note: example uses the one that always present regardless of dimensions. const bool need_pack = brg_po.get_B_pack_type() == pack_type::pack32; // If packing is needed, create a dedicated object for data transformation. if (need_pack) { // Packing B tensor routine. The BRGeMM ukernel expects B passed in a // special VNNI format for low precision data types, e.g., bfloat16_t. // Note: the routine doesn't provide a `batch_size` argument in the // constructor as it can be either incorporated into `K` dimension, or // manually iterated over in a for-loop on the user side. transform pack_B(/* K = */ K_k * n_calls, /* N = */ N, /* in_pack_type = */ pack_type::no_trans, /* in_ld = */ N, /* out_ld = */ ldb, /* in_dt = */ b_dt, /* out_dt = */ b_dt); // Size of the packed tensor. blocked_B_size = ldb * K_k * memory::data_type_size(b_dt); B_blocked = new uint8_t[blocked_B_size * n_calls]; B_base_ptr = B_blocked; // Pack B routine execution. // Note: usually should be split to process only that part of B that the // ukernel will execute. pack_B.generate(); pack_B.execute(B_ptr, B_blocked); } // BRGeMM ukernel execute section. // Prepare buffers for execution. std::vector<std::pair<memory::dim, memory::dim>> A_B_offsets(batch_size); for (memory::dim i = 0; i < batch_size; i++) { const memory::dim A_offset_i = i * K_k * a_dt_size; const memory::dim B_offset_i = need_pack ? i * blocked_B_size : i * N * K_k * b_dt_size; A_B_offsets[i] = std::make_pair(A_offset_i, B_offset_i); } if (brg) { // Make an object to call HW specialized routines. For example, prepare // AMX unit. brg.set_hw_context(); // An execute call. `A_B` is a vector of pointers to A and packed B // tensors. `acc_ptr` is a pointer to an accumulator buffer. brg.execute(A_ptr, B_base_ptr, A_B_offsets, C_ptr, scratchpad.data()); } // Same set of operations for a ukernel with post-ops. std::vector<std::pair<memory::dim, memory::dim>> A_B_po_offsets; const memory::dim A_offset_po = batch_size * K_k * a_dt_size; const memory::dim B_offset_po = need_pack ? batch_size * blocked_B_size : batch_size * N * K_k * b_dt_size; A_B_po_offsets.emplace_back(A_offset_po, B_offset_po); // This object also requires this call. brg_po.set_hw_context(); // Prepare post-ops arguments and put them in a vector to make sure pointers // are sitting side by side. std::vector<const void *> bin_po_ptrs; bin_po_ptrs.push_back(binary_add_mem.get_data_handle()); // Setting post-ops arguments into an attributes arguments storage. attr_params params; params.set_post_ops_args(bin_po_ptrs.data()); params.set_B_scales(B_scales_mem.get_data_handle()); // An execute call. The difference here is an additional D tensor pointer // to store final output result after finishing accumulation and post-ops // application. brg_po.execute(A_ptr, B_base_ptr, A_B_po_offsets, C_ptr, D_mem.get_data_handle(), scratchpad.data(), params); // Once all computations are done, need to release HW context. brgemm::release_hw_context(); // Clean up an extra buffer. delete B_blocked; // Used for verification results, need unconditional reorder. auto user_D_mem = memory(D_f32_md, engine, D_data.data()); reorder(D_mem, user_D_mem).execute(engine_stream, D_mem, user_D_mem); // Skip the check by default as data filling doesn't help with proper // verification of the result. Negative result doesn't necessarily mean // the functionality is broken. This is just a general sanity check. if (true) return; // A simplified fast verification that ukernel returned expected results. // Note: potential off-by-1 or 2 errors may pop up. This could be solved // with more sparse filling. bool to_throw = false; for (int m = 0; m < M; m++) { for (int n = 0; n < N; n++) { D_user_data[m * N + n] = 0; for (int k = 0; k < K; k++) { D_user_data[m * N + n] += A_user_data[m * K + k] * B_user_data[k * N + n]; } // B scales ref D_user_data[m * N + n] *= B_scales_user_data[n]; // Relu post-op ref D_user_data[m * N + n] = std::max(D_user_data[m * N + n], 0.f); // Binary post-op ref D_user_data[m * N + n] += binary_add_user_data[0]; const float diff = fabsf(D_user_data[m * N + n] - D_data[m * N + n]); if (diff > 1.19e-7) { to_throw = true; if (true) { printf("Error: [%3d:%3d] Ref:%12g Got:%12g Diff:%12g\n", m, n, D_user_data[m * N + n], D_data[m * N + n], diff); } } } } if (to_throw) { throw status::runtime_error; } } int main(int argc, char **argv) { return handle_example_errors({dnnl::engine::kind::cpu}, brgemm_example); }