-
Notifications
You must be signed in to change notification settings - Fork 641
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Support i1 datatype with an experimental flag. (#18713)
Enable packed i1 datatype storage This commit introduces support for packed storage of the `i1` (bit) datatype. When subbyte type packing is enabled via the `--iree-experimental-packed-i1-storage` option, vectors of `i1` elements will be stored in a compact packed representation. For example, a `vector<6xi1>` will occupy a single byte of memory with the 6 bit elements packed together and 2 padding bits. A `vector<3x3xi1>` will take up 2 bytes, with the 9 bit elements packed across the bytes and 7 padding bits. Limitations: - To ensure correct behavior, the tiling configuration aligns the innermost dimension data loads with byte boundaries. This is necessitated by the current lack of emulation for unaligned subbyte vector loading/storing. - Unaligned subbyte emulation support can be added in the future, though it may incur some performance overhead. This change requires corresponding updates in the frontend to utilize the packed `i1` storage format. Signed-off-by: Alan Li <me@alanli.org>
- Loading branch information
Showing
10 changed files
with
191 additions
and
8 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
20 changes: 20 additions & 0 deletions
20
...iler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors_packing_i1.mlir
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| // RUN: iree-opt --split-input-file --iree-stream-encode-host-tensors --iree-experimental-packed-i1-storage %s | FileCheck %s | ||
|
|
||
| func.func @unaligned_i1_size() -> index { | ||
| %0 = stream.tensor.sizeof tensor<12xi1> : index | ||
| return %0 : index | ||
| } | ||
| // CHECK: func @unaligned_i1_size() -> index { | ||
| // CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index | ||
| // CHECK: return %[[C2]] : index | ||
|
|
||
| // ----- | ||
|
|
||
| func.func @aligned_i1_size() -> index { | ||
| %0 = stream.tensor.sizeof tensor<24xi1> : index | ||
| return %0 : index | ||
| } | ||
|
|
||
| // CHECK: func @aligned_i1_size() -> index { | ||
| // CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index | ||
| // CHECK: return %[[C3]] : index |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,50 @@ | ||
| # Copyright 2024 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 | ||
|
|
||
| # Tests of end-to-end IREE support for individual ops in the TOSA dialect. | ||
| # Each test file should have a name matching the corresponding TOSA op and test only the | ||
| # functionality of that op (though may make use of other ops where necessary). Tests should be | ||
| # written using the IREE Check framework. | ||
| # See https://iree.dev/developers/general/testing-guide/#iree-core-end-to-end-e2e-tests. | ||
|
|
||
| load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob") | ||
| load("//build_tools/bazel:iree_check_test.bzl", "iree_check_single_backend_test_suite") | ||
|
|
||
| package( | ||
| features = ["layering_check"], | ||
| licenses = ["notice"], # Apache 2.0 | ||
| ) | ||
|
|
||
| LLVM_SRCS = enforce_glob( | ||
| # keep sorted | ||
| [ | ||
| "subbyte_types.mlir", | ||
| ], | ||
| include = ["*.mlir"], | ||
| exclude = [], | ||
| ) | ||
|
|
||
| iree_check_single_backend_test_suite( | ||
| name = "check_llvm-cpu_subbyte_emulation", | ||
| srcs = LLVM_SRCS, | ||
| compiler_flags = [ | ||
| "--iree-llvmcpu-target-cpu=generic", | ||
| "--iree-experimental-packed-i1-storage", | ||
| ], | ||
| driver = "local-task", | ||
| tags = [ | ||
| # subbyte support for wasm is not on priorities. | ||
| "nowasm", | ||
| ], | ||
| target_backend = "llvm-cpu", | ||
| ) | ||
|
|
||
| test_suite( | ||
| name = "check", | ||
| tests = [ | ||
| ":check_llvm-cpu_subbyte_emulation", | ||
| ], | ||
| ) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,29 @@ | ||
| ################################################################################ | ||
| # Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from # | ||
| # tests/e2e/subbyte_types/BUILD.bazel # | ||
| # # | ||
| # Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary # | ||
| # CMake-only content. # | ||
| # # | ||
| # To disable autogeneration for this file entirely, delete this header. # | ||
| ################################################################################ | ||
|
|
||
| iree_add_all_subdirs() | ||
|
|
||
| iree_check_single_backend_test_suite( | ||
| NAME | ||
| check_llvm-cpu_subbyte_emulation | ||
| SRCS | ||
| "subbyte_types.mlir" | ||
| TARGET_BACKEND | ||
| "llvm-cpu" | ||
| DRIVER | ||
| "local-task" | ||
| COMPILER_FLAGS | ||
| "--iree-llvmcpu-target-cpu=generic" | ||
| "--iree-experimental-packed-i1-storage" | ||
| LABELS | ||
| "nowasm" | ||
| ) | ||
|
|
||
| ### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ### |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,28 @@ | ||
| func.func @i1_type() { | ||
| %c0 = arith.constant 0 : index | ||
| %c255 = arith.constant 255 : i8 | ||
| %input1 = util.unfoldable_constant dense<[85]> : tensor<1xi8> // b01010101 | ||
| %input2 = util.unfoldable_constant dense<[170]> : tensor<1xi8> // b10101010 | ||
| %lhs = flow.tensor.bitcast %input1 : tensor<1xi8> -> tensor<8xi1> | ||
| %rhs = flow.tensor.bitcast %input2 : tensor<1xi8> -> tensor<8xi1> | ||
| %empty = tensor.empty() : tensor<8xi1> | ||
| %res = linalg.generic | ||
| {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} | ||
| ins(%lhs, %rhs : tensor<8xi1>, tensor<8xi1>) outs(%empty: tensor<8xi1>) { | ||
| ^bb0(%inlhs: i1, %inrhs: i1, %out: i1): | ||
| %inres = arith.xori %inlhs, %inrhs: i1 | ||
| linalg.yield %inres : i1 | ||
| } -> tensor<8xi1> | ||
| %tensor_res = flow.tensor.bitcast %res : tensor<8xi1> -> tensor<1xi8> | ||
| check.expect_eq_const(%tensor_res, dense<[255]> : tensor<1xi8>) : tensor<1xi8> | ||
| return | ||
| } | ||
|
|
||
| func.func @i1_type_slice() { | ||
| %input = util.unfoldable_constant dense<[0, 255, 0]> : tensor<3xi8> | ||
| %flat_input_all = flow.tensor.bitcast %input : tensor<3xi8> -> tensor<24xi1> | ||
| %slice = tensor.extract_slice %flat_input_all[8][8][1] : tensor<24xi1> to tensor<8xi1> | ||
| %tensor_res = flow.tensor.bitcast %slice : tensor<8xi1> -> tensor<1xi8> | ||
| check.expect_eq_const(%tensor_res, dense<[255]> : tensor<1xi8>) : tensor<1xi8> | ||
| return | ||
| } |