From 3099d0b114121952f10c049f1c2b1f9a77575a1d Mon Sep 17 00:00:00 2001 From: Alan Li Date: Thu, 19 Dec 2024 13:18:52 +0000 Subject: [PATCH] update according to comments --- .../Conversion/HALToStream/Patterns.cpp | 4 +- .../compiler/Utils/ElementPackingUtils.cpp | 45 ++++++++++--------- .../iree/compiler/Utils/ElementPackingUtils.h | 10 ----- 3 files changed, 27 insertions(+), 32 deletions(-) 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 9b8ca8e35570..b27c8a516e43 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.cpp +++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.cpp @@ -91,7 +91,9 @@ struct ConvertTensorImportOp RankedTensorType tensorType, ValueRange dynamicDims, OpBuilder &builder) { - // If the encoding attr is about packed storage then we don't need all this + // 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(); } diff --git a/compiler/src/iree/compiler/Utils/ElementPackingUtils.cpp b/compiler/src/iree/compiler/Utils/ElementPackingUtils.cpp index 2d53b8c47a38..f4a40cce1187 100644 --- a/compiler/src/iree/compiler/Utils/ElementPackingUtils.cpp +++ b/compiler/src/iree/compiler/Utils/ElementPackingUtils.cpp @@ -29,12 +29,8 @@ llvm::cl::opt clEnableI1Support( namespace mlir::iree_compiler { -bool needToPackSubByteElementBitWidth(unsigned bitWidth) { - return needToPackSubByteElementBitWidth( - bitWidth, /*isPackedStorage=*/clEnableI1Support); -} - -bool needToPackSubByteElementBitWidth(unsigned bitWidth, bool isPackedStorage) { +static bool needToPackSubByteElementBitWidthImpl(unsigned bitWidth, + bool isPackedStorage) { // Enable i1 support if requested. if (isPackedStorage && bitWidth == 1) { return true; @@ -46,6 +42,11 @@ bool needToPackSubByteElementBitWidth(unsigned bitWidth, bool isPackedStorage) { 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()); // Two paths to enable packed storage for i1 tensors: the attribute or cl @@ -53,15 +54,11 @@ bool needToPackSubByteElements(RankedTensorType shapedType) { // tensors with attributes. bool isPackedStorage = IREE::Encoding::hasPackedStorageAttr(shapedType) || clEnableI1Support; - return needToPackSubByteElementBitWidth(bitWidth, isPackedStorage); + return needToPackSubByteElementBitWidthImpl(bitWidth, isPackedStorage); } -Type legalizeStorageElementType(Type elementType) { - return legalizeStorageElementType(elementType, - /*isPackedStorage=*/clEnableI1Support); -} - -Type legalizeStorageElementType(Type elementType, bool isPackedStorage) { +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) @@ -69,7 +66,7 @@ Type legalizeStorageElementType(Type elementType, bool isPackedStorage) { // For sub-byte elements, default to pack them into bytes. unsigned bitWidth = intType.getWidth(); - if (needToPackSubByteElementBitWidth(bitWidth, isPackedStorage)) + if (needToPackSubByteElementBitWidthImpl(bitWidth, isPackedStorage)) return elementType; // Otherwise, extend them to the next power-of-two bit width. @@ -81,6 +78,12 @@ Type legalizeStorageElementType(Type elementType, bool isPackedStorage) { 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, @@ -93,16 +96,16 @@ Value calculateStorageElementCountInBytes(Location loc, loc, builder, shapedType, dynamicDims); } - // TODO: remove cl options once frontend can emit packed i1 tensors. + // TODO(lialan): remove cl options once frontend can emit packed i1 tensors. bool isPackedStorage = IREE::Encoding::hasPackedStorageAttr(shapedType) || clEnableI1Support; - Type alignedElementType = - legalizeStorageElementType(shapedType.getElementType(), isPackedStorage); + 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, isPackedStorage)) { + if (!needToPackSubByteElementBitWidthImpl(elementBits, isPackedStorage)) { staticCount *= IREE::Util::getRoundedElementByteWidth(alignedElementType); } @@ -117,7 +120,7 @@ Value calculateStorageElementCountInBytes(Location loc, value = builder.createOrFold(loc, value, dim); } // Sub-byte packing requires putting multiple elements in the same byte. - if (needToPackSubByteElementBitWidth(elementBits, isPackedStorage)) { + 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 @@ -140,12 +143,12 @@ Value calculateStorageElementOffsetInBytes(Location loc, // TODO: remove cl options once frontend can emit packed i1 tensors. bool isPackedStorage = IREE::Encoding::hasPackedStorageAttr(originalType) || clEnableI1Support; - Type alignedElementType = legalizeStorageElementType( + 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, isPackedStorage)) { + 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 5dfa37381d3d..9de6ea70c26a 100644 --- a/compiler/src/iree/compiler/Utils/ElementPackingUtils.h +++ b/compiler/src/iree/compiler/Utils/ElementPackingUtils.h @@ -15,11 +15,6 @@ 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, bool isPackedStorage); - -/// Temporary wrapper for the above function. `isPackedStorage` will be -/// determined by the cl option. This allows enabling packed storage for i1 -/// in both attribute and cl option ways. bool needToPackSubByteElementBitWidth(unsigned bitWidth); /// Returns true if the given |shapedType|, if appearing at runtime-kernel @@ -33,11 +28,6 @@ bool needToPackSubByteElements(RankedTensorType shapedType); /// runtime and kernel. For such cases, we perform tight packing for supported /// sub-byte elements, and expand to the next power-of-two bit width for other /// cases. -Type legalizeStorageElementType(Type elementType, bool isPackedStorage); - -/// Temporary wrapper for the above function. `isPackedStorage` will be -/// determined by the cl option. This allows enabling packed storage for i1 -/// in both attribute and cl option ways. Type legalizeStorageElementType(Type elementType); /// Emits IR with the given |builder| to calculate the total number of bytes