From 419ee6b7c08a15e982cf48570c9ec83676b99eac Mon Sep 17 00:00:00 2001 From: Alan Li Date: Tue, 3 Dec 2024 09:49:29 +0000 Subject: [PATCH] [i1] Implement `packed_storage` layout encoding attribute * make `packed_storage` as a type of `iree_encoding` attribute, and make type converters accept it. * `i1` tensors with `#iree_encoding.packed_storage` will be interpreted as packed i1 type, same as specifying `--iree-experimental-packed-i1-storage`. * `--iree-experimental-packed-i1-storage` are kept for testing purposes. We can drop this option after frontend enables emitting `i1` tensors with attributes. Signed-off-by: Alan Li --- .../compiler/Codegen/Common/EncodingUtils.cpp | 11 ++-- .../Dialect/Encoding/IR/EncodingAttrs.cpp | 4 ++ .../Dialect/Encoding/IR/EncodingAttrs.td | 11 ++++ .../Dialect/Encoding/IR/EncodingTypes.h | 3 + .../Conversion/HALToStream/Patterns.cpp | 8 +++ .../Stream/Transforms/EncodeTensors.cpp | 3 +- .../Stream/Transforms/test/BUILD.bazel | 3 +- .../Stream/Transforms/test/CMakeLists.txt | 3 +- .../encode_host_tensors_packing_i1_attr.mlir | 22 +++++++ ...ensors_packing_i1_experimental_clopt.mlir} | 0 .../compiler/Utils/ElementPackingUtils.cpp | 62 ++++++++++++++----- .../iree/compiler/Utils/ElementPackingUtils.h | 1 + 12 files changed, 108 insertions(+), 23 deletions(-) create mode 100644 compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors_packing_i1_attr.mlir rename compiler/src/iree/compiler/Dialect/Stream/Transforms/test/{encode_host_tensors_packing_i1.mlir => encode_host_tensors_packing_i1_experimental_clopt.mlir} (100%) diff --git a/compiler/src/iree/compiler/Codegen/Common/EncodingUtils.cpp b/compiler/src/iree/compiler/Codegen/Common/EncodingUtils.cpp index e3ca734a964b0..bab399484e986 100644 --- a/compiler/src/iree/compiler/Codegen/Common/EncodingUtils.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/EncodingUtils.cpp @@ -6,6 +6,7 @@ #include "iree/compiler/Codegen/Common/EncodingUtils.h" #include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h" +#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h" #include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Utils/IndexingUtils.h" @@ -62,16 +63,18 @@ MaterializeEncodingConversionTarget::MaterializeEncodingConversionTarget( // Mark any operation that has operands/results with encoding as // illegal. markUnknownOpDynamicallyLegal([](Operation *op) { - auto typeHasEncoding = [](Type t) -> bool { + auto typeHasDataTilingEncoding = [](Type t) -> bool { auto tensorType = dyn_cast(t); - return tensorType && tensorType.getEncoding(); + if (!tensorType) + return false; + return getEncodingAttr(tensorType) != nullptr; }; auto valueHasEncoding = [=](Value v) -> bool { - return typeHasEncoding(v.getType()); + return typeHasDataTilingEncoding(v.getType()); }; bool hasOperandOrResultsWithEncoding = llvm::any_of(op->getOperands(), valueHasEncoding) || - llvm::any_of(op->getResultTypes(), typeHasEncoding); + llvm::any_of(op->getResultTypes(), typeHasDataTilingEncoding); return !hasOperandOrResultsWithEncoding; }); } diff --git a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp index 593d9b8fc5c62..26d7d23d6d92a 100644 --- a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp +++ b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp @@ -243,6 +243,10 @@ EncodingAttr getEncodingAttr(RankedTensorType type) { return dyn_cast_or_null(type.getEncoding()); } +bool hasPackedStorageAttr(RankedTensorType type) { + return dyn_cast_or_null(type.getEncoding()) != nullptr; +} + FailureOr getEncodingContractionDims(EncodingAttr encoding) { auto indexingMapsAttr = encoding.getUserIndexingMaps(); diff --git a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.td b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.td index 54829b68e2cf6..9c2b1f2eca01d 100644 --- a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.td +++ b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.td @@ -41,6 +41,17 @@ def EncodingOpType : IREEEncoding_I32EnumAttr<"EncodingOpType", def EncodingOpTypeAttr: IREEEncoding_EnumAttr; + +def PackedStorageAttr : IREEEncoding_Attr<"PackedStorage"> { + let mnemonic = "packed_storage"; + let summary = [{Indicates packed storage data type.}]; + let description = [{ + This attribute indicates this is a back-to-back packed storage in memory. + This attribute takes no arguments. + }]; + let genVerifyDecl = 0; +} + def EncodingAttr : IREEEncoding_Attr<"Encoding", [ DeclareAttrInterfaceMethods getEncodingContractionDims(EncodingAttr encoding); diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.cpp b/compiler/src/iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.cpp index e597aaffba8fe..433f559bfe74f 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.cpp +++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.cpp @@ -6,6 +6,7 @@ #include "iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.h" +#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h" #include "iree/compiler/Dialect/HAL/IR/HALOps.h" #include "iree/compiler/Dialect/Stream/Conversion/PatternUtils.h" #include "iree/compiler/Dialect/Stream/IR/StreamDialect.h" @@ -100,6 +101,13 @@ struct ConvertTensorImportOp RankedTensorType tensorType, ValueRange dynamicDims, OpBuilder &builder) { + // If the encoding attr is about packed storage then we don't need + // assertion, because packed storage attribute is about memory layout and it + // doesn't affect the tensor shape. + if (IREE::Encoding::hasPackedStorageAttr(tensorType)) { + return success(); + } + auto expectedElementType = builder.create( loc, tensorType.getElementType()); auto expectedEncodingType = builder.create( diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/EncodeTensors.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/EncodeTensors.cpp index 5cb95099bbbc3..42b359784d0cd 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/EncodeTensors.cpp +++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/EncodeTensors.cpp @@ -46,7 +46,8 @@ static LogicalResult checkEncoding(Operation *op, RankedTensorType encodingType, ValueRange encodingDims, PatternRewriter &rewriter) { auto encoding = encodingType.getEncoding(); - if (encoding && !llvm::isa(encoding)) { + if (encoding && !llvm::isa(encoding)) { return rewriter.notifyMatchFailure(op, [=](Diagnostic &d) { d << "unsupported tensor encoding: " << encodingType; }); diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel index 87d6bea609773..a244adbe6e9ae 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel +++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel @@ -29,7 +29,8 @@ iree_lit_test_suite( "encode_device_tensors_packing.mlir", "encode_host_tensors.mlir", "encode_host_tensors_packing.mlir", - "encode_host_tensors_packing_i1.mlir", + "encode_host_tensors_packing_i1_attr.mlir", + "encode_host_tensors_packing_i1_experimental_clopt.mlir", "fold_globals.mlir", "fold_uniform_operands.mlir", "fuse_dispatch_bindings.mlir", diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt index 8c4ca85927c53..12cf7e0d8f352 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt +++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt @@ -27,7 +27,8 @@ iree_lit_test_suite( "encode_device_tensors_packing.mlir" "encode_host_tensors.mlir" "encode_host_tensors_packing.mlir" - "encode_host_tensors_packing_i1.mlir" + "encode_host_tensors_packing_i1_attr.mlir" + "encode_host_tensors_packing_i1_experimental_clopt.mlir" "fold_globals.mlir" "fold_uniform_operands.mlir" "fuse_dispatch_bindings.mlir" diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors_packing_i1_attr.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors_packing_i1_attr.mlir new file mode 100644 index 0000000000000..eefc9810aed58 --- /dev/null +++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors_packing_i1_attr.mlir @@ -0,0 +1,22 @@ +// RUN: iree-opt --split-input-file --iree-stream-encode-host-tensors %s | FileCheck %s + +#packed = #iree_encoding.packed_storage +func.func @unaligned_i1_size() -> index { + %0 = stream.tensor.sizeof tensor<12xi1, #packed> : index + return %0 : index +} +// CHECK: func @unaligned_i1_size() -> index { +// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index +// CHECK: return %[[C2]] : index + +// ----- + +#packed = #iree_encoding.packed_storage +func.func @aligned_i1_size() -> index { + %0 = stream.tensor.sizeof tensor<24xi1, #packed> : index + return %0 : index +} + +// CHECK: func @aligned_i1_size() -> index { +// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index +// CHECK: return %[[C3]] : index diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors_packing_i1.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors_packing_i1_experimental_clopt.mlir similarity index 100% rename from compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors_packing_i1.mlir rename to compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors_packing_i1_experimental_clopt.mlir diff --git a/compiler/src/iree/compiler/Utils/ElementPackingUtils.cpp b/compiler/src/iree/compiler/Utils/ElementPackingUtils.cpp index 2f92ffa629135..f4a40cce11875 100644 --- a/compiler/src/iree/compiler/Utils/ElementPackingUtils.cpp +++ b/compiler/src/iree/compiler/Utils/ElementPackingUtils.cpp @@ -15,17 +15,24 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/IR/BuiltinTypes.h" -namespace mlir::iree_compiler { - llvm::cl::opt clEnableI1Support( "iree-experimental-packed-i1-storage", llvm::cl::desc( - "Experimental feature: enable i1 data type support in codegen"), + "Experimental feature: force to use packed storage for i1 tensors." + "Turning on this option will see i1 tensors as if it has " + "#iree_encoding.packed_storage attribute." + "This is to allow an alternative way to test the packed storage " + "feature before frontend can emit packed i1 tensors." + "This option can be dropped once the frontend can emit packed i1 " + "tensors."), llvm::cl::init(false)); -bool needToPackSubByteElementBitWidth(unsigned bitWidth) { +namespace mlir::iree_compiler { + +static bool needToPackSubByteElementBitWidthImpl(unsigned bitWidth, + bool isPackedStorage) { // Enable i1 support if requested. - if (clEnableI1Support && bitWidth == 1) { + if (isPackedStorage && bitWidth == 1) { return true; } // Require the original bit width to be some power of two for now to avoid @@ -35,12 +42,23 @@ bool needToPackSubByteElementBitWidth(unsigned bitWidth) { return bitWidth < 8 && llvm::isPowerOf2_32(bitWidth) && bitWidth != 1; } +bool needToPackSubByteElementBitWidth(unsigned bitWidth) { + return needToPackSubByteElementBitWidthImpl( + bitWidth, /*isPackedStorage=*/clEnableI1Support); +} + bool needToPackSubByteElements(RankedTensorType shapedType) { unsigned bitWidth = IREE::Util::getTypeBitWidth(shapedType.getElementType()); - return needToPackSubByteElementBitWidth(bitWidth); + // Two paths to enable packed storage for i1 tensors: the attribute or cl + // option. The cl option will be dropped once frontend supports emitting + // tensors with attributes. + bool isPackedStorage = + IREE::Encoding::hasPackedStorageAttr(shapedType) || clEnableI1Support; + return needToPackSubByteElementBitWidthImpl(bitWidth, isPackedStorage); } -Type legalizeStorageElementType(Type elementType) { +static Type legalizeStorageElementTypeImpl(Type elementType, + bool isPackedStorage) { // Only handle integers; floats in MLIR all have aligned widths (today). auto intType = dyn_cast(elementType); if (!intType) @@ -48,7 +66,7 @@ Type legalizeStorageElementType(Type elementType) { // For sub-byte elements, default to pack them into bytes. unsigned bitWidth = intType.getWidth(); - if (needToPackSubByteElementBitWidth(bitWidth)) + if (needToPackSubByteElementBitWidthImpl(bitWidth, isPackedStorage)) return elementType; // Otherwise, extend them to the next power-of-two bit width. @@ -60,6 +78,12 @@ Type legalizeStorageElementType(Type elementType) { intType.getSignedness()); } +Type legalizeStorageElementType(Type elementType) { + // Consider packed storage for i1 tensors if cl opt is set. + return legalizeStorageElementTypeImpl(elementType, + /*isPackedStorage=*/clEnableI1Support); +} + Value calculateStorageElementCountInBytes(Location loc, RankedTensorType shapedType, ValueRange dynamicDims, @@ -72,13 +96,16 @@ Value calculateStorageElementCountInBytes(Location loc, loc, builder, shapedType, dynamicDims); } - Type alignedElementType = - legalizeStorageElementType(shapedType.getElementType()); + // TODO(lialan): remove cl options once frontend can emit packed i1 tensors. + bool isPackedStorage = + IREE::Encoding::hasPackedStorageAttr(shapedType) || clEnableI1Support; + Type alignedElementType = legalizeStorageElementTypeImpl( + shapedType.getElementType(), isPackedStorage); unsigned elementBits = IREE::Util::getTypeBitWidth(alignedElementType); // Calculate all static dims first, if any. int64_t staticCount = 1; - if (!needToPackSubByteElementBitWidth(elementBits)) { + if (!needToPackSubByteElementBitWidthImpl(elementBits, isPackedStorage)) { staticCount *= IREE::Util::getRoundedElementByteWidth(alignedElementType); } @@ -93,13 +120,13 @@ Value calculateStorageElementCountInBytes(Location loc, value = builder.createOrFold(loc, value, dim); } // Sub-byte packing requires putting multiple elements in the same byte. - if (needToPackSubByteElementBitWidth(elementBits)) { + if (needToPackSubByteElementBitWidthImpl(elementBits, isPackedStorage)) { assert(8 % elementBits == 0); unsigned byteElements = 8 / elementBits; // TODO(antiagainst): We may want to emit runtime check to make sure this is // divisible. auto divisor = builder.create(loc, byteElements); - if (!clEnableI1Support && dynamicDims.empty() && + if (!isPackedStorage && dynamicDims.empty() && (staticCount * elementBits) % 8 != 0) { return nullptr; } @@ -113,12 +140,15 @@ Value calculateStorageElementOffsetInBytes(Location loc, RankedTensorType originalType, Value linearizedIndex, OpBuilder &builder) { - Type alignedElementType = - legalizeStorageElementType(originalType.getElementType()); + // TODO: remove cl options once frontend can emit packed i1 tensors. + bool isPackedStorage = + IREE::Encoding::hasPackedStorageAttr(originalType) || clEnableI1Support; + Type alignedElementType = legalizeStorageElementTypeImpl( + originalType.getElementType(), isPackedStorage); unsigned elementBits = IREE::Util::getTypeBitWidth(alignedElementType); // Sub-byte packing requires putting multiple elements in the same byte. - if (needToPackSubByteElementBitWidth(elementBits)) { + if (needToPackSubByteElementBitWidthImpl(elementBits, isPackedStorage)) { Value byteElements = builder.create(loc, 8 / elementBits); // TODO(antiagainst): We may want to emit runtime check to make sure this is diff --git a/compiler/src/iree/compiler/Utils/ElementPackingUtils.h b/compiler/src/iree/compiler/Utils/ElementPackingUtils.h index 6d5d18dcc7180..9de6ea70c26a7 100644 --- a/compiler/src/iree/compiler/Utils/ElementPackingUtils.h +++ b/compiler/src/iree/compiler/Utils/ElementPackingUtils.h @@ -16,6 +16,7 @@ namespace mlir::iree_compiler { /// Returns true if the given |bitWidth|, if appearing at runtime-kernel /// interface, is less than a byte that should be tightly packed together. bool needToPackSubByteElementBitWidth(unsigned bitWidth); + /// Returns true if the given |shapedType|, if appearing at runtime-kernel /// interface, has sub-byte element types that should be tightly packed /// together.