[i1] Remove command line option to enable packed storage #19528

lialan · 2024-12-19T15:46:49Z

only use #iree_encoding.packed_storage to designate if an i1 tensor is of packed memory layout.
remove iree-experimental-packed-i1-storage command line option.
teach type converters to allow casting into packed tensor types

lialan · 2024-12-31T03:11:53Z

This is a successor, improving PR to #19354

hanhanW · 2025-01-07T07:51:30Z

The diff is off. There are some changes from #19354. E.g., the below is not part of this PR, right?

Can you rebase and fix it?

lialan · 2025-01-10T05:11:09Z

Note that changes that are not in this PR: #19618

hanhanW

First round of the review comments, thanks for pushing on this!

compiler/src/iree/compiler/Codegen/Common/EncodingUtils.cpp

compiler/src/iree/compiler/Utils/ElementPackingUtils.h

compiler/src/iree/compiler/Utils/ElementPackingUtils.cpp

hanhanW · 2025-01-13T05:59:02Z

tests/e2e/subbyte_types/subbyte_types_attr.mlir

+#packed = #iree_encoding.packed_storage
+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, #packed>


I'm not familiar with how flow.tensor.bitcast works. Is the encoding dropped in the TensorExportBufferViewOpPattern pattern? How does it become a tensor_export_buffer_view op? I think we need some input from Ben about the change and the test.

The encoding dropped while it is calculating streamOp's allocated sizes.

flow.tensor.bitcast is just an omni-potent caster here to cast a tensor to another tensor with #packed_storage attribute.

hanhanW

Thanks for the update! Please also add a test to https://github.com/iree-org/iree/blob/main/compiler/src/iree/compiler/Codegen/Common/test/type_propagation_packing.mlir

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.td

compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp

compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp

hanhanW · 2025-01-15T05:19:33Z

compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp

@@ -1512,7 +1517,7 @@ LogicalResult TensorCloneOp::verify() {
  // information.
  auto sourceEncoding = llvm::cast<RankedTensorType>(op.getSourceEncoding());
  auto resultEncoding = llvm::cast<RankedTensorType>(op.getResultEncoding());
-  if (sourceEncoding.getEncoding() != resultEncoding.getEncoding()) {
+  if (getEncodingAttr(sourceEncoding) != getEncodingAttr(resultEncoding)) {


Is this required by flow.tensor.bitcast lowering? I.e., the op is lowered to flow.tensor.clone and you need to bypass the check?

I'm not convinced that the change is correct. Because the getEncodingAttr checks if the tensor type has IREE::Encoding::EncodingAttr attribute. There could be other encodings and it is a bug if we introduce new encodings. E.g., tensor<3x4xi32, #whatever_other_encoding_with_padding_semantic> can not be cloned to tensor<3x4xi32>. I think we need to have a stronger restriction. Perhaps just relax the check for packed_storage encoding?

iree/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp

Lines 278 to 280 in 27e7a90

EncodingAttr getEncodingAttr(RankedTensorType type) {

return dyn_cast_or_null<EncodingAttr>(type.getEncoding());

}

This is a little tricky. And yes this is required by flow.tensor.bitcast, and we would sometimes cast from a tensor without attribute to another tensor with packed attribute. In such case, we shouldn't check if both the source and result matches packed attribute.

I have slightly updated it to exclude packed attribute comparison. Suggestions are welcome.

compiler/src/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp

compiler/src/iree/compiler/Utils/ElementPackingUtils.cpp

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp

hanhanW

Reminder: the test is not added yet. #19528 (review)

compiler/src/iree/compiler/Codegen/Common/EncodingUtils.cpp

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingTypes.h

compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp

hanhanW · 2025-01-15T14:57:18Z

compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp

+  if (sourceEncoding.getEncoding() != resultEncoding.getEncoding() &&
+      !IREE::Encoding::hasPackedStorageAttr(sourceEncoding) &&
+      !IREE::Encoding::hasPackedStorageAttr(resultEncoding)) {


It looks better to me. I think we need some input from @benvanik

The previous comment is not tracked in the review mode. Here was the previous comment: #19528 (comment)

TLDR is that we use flow.tensor.bitcast to prepare packed data for e2e testing. The op becomes TensorClone op during the lowering. So we want to relax the verification here.

%flat_input_all = flow.tensor.bitcast %input : tensor<3xi8> -> tensor<24xi1, #packed>

@benvanik Do you have opinions on this?

lialan · 2025-01-16T07:59:16Z

@hanhanW I've added a test to test type propagation.

hanhanW

Just a few final nits. LGTM except the TensorCloneOp::verify part. Please coordinate with @benvanik about it.

compiler/src/iree/compiler/Codegen/Common/test/type_propagation_packing.mlir

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingTypes.h

* only use `#iree_encoding.packed_storage` to designate if an `i1` tensor is of packed memory layout. * remove `iree-experimental-packed-i1-storage` command line option. * teach type converters to allow casting into packed tensor types Signed-off-by: Alan Li <[email protected]>

lialan force-pushed the lialan/i1_attr branch from 3099d0b to 419ee6b Compare December 29, 2024 02:50

lialan force-pushed the lialan/bitcast branch from 9159d90 to e27a766 Compare December 29, 2024 02:57

lialan force-pushed the lialan/i1_attr branch from d26a392 to 7ac00fe Compare December 29, 2024 09:10

lialan force-pushed the lialan/bitcast branch from e27a766 to 7f752a4 Compare December 29, 2024 10:48

lialan force-pushed the lialan/bitcast branch from bd53c95 to b17ea2f Compare December 31, 2024 04:50

lialan marked this pull request as ready for review December 31, 2024 09:30

lialan requested review from bjacob, hanhanW and benvanik as code owners December 31, 2024 09:30

lialan mentioned this pull request Jan 6, 2025

General support of i1 mask in attention #19380

Open

lialan force-pushed the lialan/i1_attr branch from 7ac00fe to 7fc957d Compare January 6, 2025 08:39

Base automatically changed from lialan/i1_attr to main January 10, 2025 02:02

lialan force-pushed the lialan/bitcast branch 4 times, most recently from a6f4a4c to 6468ea1 Compare January 10, 2025 04:56

hanhanW requested changes Jan 13, 2025

View reviewed changes

lialan force-pushed the lialan/bitcast branch 3 times, most recently from c1711b1 to fa8e6d6 Compare January 14, 2025 05:49

lialan requested a review from hanhanW January 14, 2025 09:38

hanhanW requested changes Jan 15, 2025

View reviewed changes

lialan force-pushed the lialan/bitcast branch 3 times, most recently from 09316fa to 6f232c9 Compare January 15, 2025 13:06

hanhanW requested changes Jan 15, 2025

View reviewed changes

lialan force-pushed the lialan/bitcast branch from aad9619 to 2e8d28b Compare January 16, 2025 02:48

lialan force-pushed the lialan/bitcast branch 2 times, most recently from f4b5379 to 4a626d6 Compare January 16, 2025 05:26

hanhanW reviewed Jan 16, 2025

View reviewed changes

lialan force-pushed the lialan/bitcast branch from 51d3f0f to 1a6f454 Compare January 17, 2025 04:19

lialan force-pushed the lialan/bitcast branch 2 times, most recently from dbcaa0a to 0e60383 Compare February 5, 2025 17:41

remove deps

6e9fa41

lialan force-pushed the lialan/bitcast branch from df53753 to 6e9fa41 Compare February 5, 2025 20:11

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[i1] Remove command line option to enable packed storage #19528

[i1] Remove command line option to enable packed storage #19528

lialan commented Dec 19, 2024

lialan commented Dec 31, 2024

hanhanW commented Jan 7, 2025

lialan commented Jan 10, 2025 •

edited

Loading

hanhanW left a comment

hanhanW Jan 13, 2025

lialan Jan 14, 2025 •

edited

Loading

hanhanW left a comment

hanhanW Jan 15, 2025

lialan Jan 15, 2025

hanhanW left a comment

hanhanW Jan 15, 2025

lialan Feb 5, 2025

lialan commented Jan 16, 2025

hanhanW left a comment

	EncodingAttr getEncodingAttr(RankedTensorType type) {
	return dyn_cast_or_null<EncodingAttr>(type.getEncoding());
	}

[i1] Remove command line option to enable packed storage #19528

Are you sure you want to change the base?

[i1] Remove command line option to enable packed storage #19528

Conversation

lialan commented Dec 19, 2024

lialan commented Dec 31, 2024

hanhanW commented Jan 7, 2025

lialan commented Jan 10, 2025 • edited Loading

hanhanW left a comment

Choose a reason for hiding this comment

hanhanW Jan 13, 2025

Choose a reason for hiding this comment

lialan Jan 14, 2025 • edited Loading

Choose a reason for hiding this comment

hanhanW left a comment

Choose a reason for hiding this comment

hanhanW Jan 15, 2025

Choose a reason for hiding this comment

lialan Jan 15, 2025

Choose a reason for hiding this comment

hanhanW left a comment

Choose a reason for hiding this comment

hanhanW Jan 15, 2025

Choose a reason for hiding this comment

lialan Feb 5, 2025

Choose a reason for hiding this comment

lialan commented Jan 16, 2025

hanhanW left a comment

Choose a reason for hiding this comment

lialan commented Jan 10, 2025 •

edited

Loading

lialan Jan 14, 2025 •

edited

Loading