[MLIR] Don't drop attached discardable attributes #111261

pashu123 · 2024-10-05T18:08:50Z

The creation of pack op was dropping discardable attributes.

llvmbot · 2024-10-05T18:09:21Z

@llvm/pr-subscribers-mlir-tensor

Author: Prashant Kumar (pashu123)

Changes

The creation of pack op was dropping custom attached attributes.

Full diff: https://github.com/llvm/llvm-project/pull/111261.diff

2 Files Affected:

(modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+9-4)
(modified) mlir/test/Dialect/Tensor/canonicalize.mlir (+12-1)

diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index defac8308b9092..659eabd2e93880 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -4337,11 +4337,16 @@ LogicalResult PackOp::canonicalize(PackOp packOp, PatternRewriter &rewriter) {
       dest =
           rewriter.create<tensor::CastOp>(loc, newDestType, packOp.getDest());
     }
-    Value newOp = rewriter.create<tensor::PackOp>(
-        loc, source, dest, packOp.getInnerDimsPos(), packOp.getMixedTiles(),
-        packOp.getPaddingValue(), packOp.getOuterDimsPerm());
+    auto clonedPackOp = cast<PackOp>(rewriter.clone(*packOp));
+    Value res = clonedPackOp.getResult();
+    rewriter.startOpModification(clonedPackOp);
+    clonedPackOp.getSourceMutable().assign(source);
+    clonedPackOp.getDestMutable().assign(dest);
+    res.setType(dest.getType());
+    rewriter.finalizeOpModification(clonedPackOp);
+
     rewriter.replaceOpWithNewOp<tensor::CastOp>(
-        packOp, packOp.getResult().getType(), newOp);
+        packOp, packOp.getResult().getType(), clonedPackOp);
     return success();
   }
 
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 86754c1c37536d..03ff45380dca9b 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -2357,7 +2357,7 @@ func.func @unpack_pack_with_padding_no_canonicalization(%t: tensor<256x512xbf16>
   %tensor_empty = tensor.empty() : tensor<4x16x64x32xbf16>
   %tensor_empty1 = tensor.empty() : tensor<224x512xbf16>
   %packed = tensor.pack %t outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [64, 32] into %tensor_empty : tensor<256x512xbf16> -> tensor<4x16x64x32xbf16>
-  %unpacked = tensor.unpack %packed inner_dims_pos = [0, 1] inner_tiles = [64, 32] into %tensor_empty1 : tensor<4x16x64x32xbf16> -> tensor<224x512xbf16> 
+  %unpacked = tensor.unpack %packed inner_dims_pos = [0, 1] inner_tiles = [64, 32] into %tensor_empty1 : tensor<4x16x64x32xbf16> -> tensor<224x512xbf16>
   return %unpacked : tensor<224x512xbf16>
 }
 
@@ -2707,3 +2707,14 @@ func.func @test_destination_multiple_result(%arg0: tensor<2x2xf32>, %arg1: tenso
   %0:2 = test.destination_style_op ins(%cast : tensor<?x2xf32>) outs(%cast_0 : tensor<?x2xf32>) -> tensor<?x2xf32>, index
   return %0#1 : index
 }
+
+// -----
+
+// CHECK-LABEL:   func.func @pack_dont_drop_attributes(
+// CHECK: tensor.pack {{.*}}  {test_attr = 16 : i64}
+func.func @pack_dont_drop_attributes(%arg0: tensor<?x?x?xf16>, %arg1: tensor<128x?x100x16x1xf16>) -> tensor<128x?x100x16x1xf16> {
+  %c32_i64 = arith.constant 32 : i64
+  %cst = arith.constant 0.000000e+00 : f16
+  %pack = tensor.pack %arg0 padding_value(%cst : f16) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [16, 1] into %arg1 {test_attr = 16 : i64} : tensor<?x?x?xf16> -> tensor<128x?x100x16x1xf16>
+  return %pack : tensor<128x?x100x16x1xf16>
+}

llvmbot · 2024-10-05T18:09:21Z

@llvm/pr-subscribers-mlir

Author: Prashant Kumar (pashu123)

Changes

The creation of pack op was dropping custom attached attributes.

Full diff: https://github.com/llvm/llvm-project/pull/111261.diff

2 Files Affected:

(modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+9-4)
(modified) mlir/test/Dialect/Tensor/canonicalize.mlir (+12-1)

diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index defac8308b9092..659eabd2e93880 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -4337,11 +4337,16 @@ LogicalResult PackOp::canonicalize(PackOp packOp, PatternRewriter &rewriter) {
       dest =
           rewriter.create<tensor::CastOp>(loc, newDestType, packOp.getDest());
     }
-    Value newOp = rewriter.create<tensor::PackOp>(
-        loc, source, dest, packOp.getInnerDimsPos(), packOp.getMixedTiles(),
-        packOp.getPaddingValue(), packOp.getOuterDimsPerm());
+    auto clonedPackOp = cast<PackOp>(rewriter.clone(*packOp));
+    Value res = clonedPackOp.getResult();
+    rewriter.startOpModification(clonedPackOp);
+    clonedPackOp.getSourceMutable().assign(source);
+    clonedPackOp.getDestMutable().assign(dest);
+    res.setType(dest.getType());
+    rewriter.finalizeOpModification(clonedPackOp);
+
     rewriter.replaceOpWithNewOp<tensor::CastOp>(
-        packOp, packOp.getResult().getType(), newOp);
+        packOp, packOp.getResult().getType(), clonedPackOp);
     return success();
   }
 
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 86754c1c37536d..03ff45380dca9b 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -2357,7 +2357,7 @@ func.func @unpack_pack_with_padding_no_canonicalization(%t: tensor<256x512xbf16>
   %tensor_empty = tensor.empty() : tensor<4x16x64x32xbf16>
   %tensor_empty1 = tensor.empty() : tensor<224x512xbf16>
   %packed = tensor.pack %t outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [64, 32] into %tensor_empty : tensor<256x512xbf16> -> tensor<4x16x64x32xbf16>
-  %unpacked = tensor.unpack %packed inner_dims_pos = [0, 1] inner_tiles = [64, 32] into %tensor_empty1 : tensor<4x16x64x32xbf16> -> tensor<224x512xbf16> 
+  %unpacked = tensor.unpack %packed inner_dims_pos = [0, 1] inner_tiles = [64, 32] into %tensor_empty1 : tensor<4x16x64x32xbf16> -> tensor<224x512xbf16>
   return %unpacked : tensor<224x512xbf16>
 }
 
@@ -2707,3 +2707,14 @@ func.func @test_destination_multiple_result(%arg0: tensor<2x2xf32>, %arg1: tenso
   %0:2 = test.destination_style_op ins(%cast : tensor<?x2xf32>) outs(%cast_0 : tensor<?x2xf32>) -> tensor<?x2xf32>, index
   return %0#1 : index
 }
+
+// -----
+
+// CHECK-LABEL:   func.func @pack_dont_drop_attributes(
+// CHECK: tensor.pack {{.*}}  {test_attr = 16 : i64}
+func.func @pack_dont_drop_attributes(%arg0: tensor<?x?x?xf16>, %arg1: tensor<128x?x100x16x1xf16>) -> tensor<128x?x100x16x1xf16> {
+  %c32_i64 = arith.constant 32 : i64
+  %cst = arith.constant 0.000000e+00 : f16
+  %pack = tensor.pack %arg0 padding_value(%cst : f16) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [16, 1] into %arg1 {test_attr = 16 : i64} : tensor<?x?x?xf16> -> tensor<128x?x100x16x1xf16>
+  return %pack : tensor<128x?x100x16x1xf16>
+}

mlir/test/Dialect/Tensor/canonicalize.mlir

The creation of pack op was dropping custom attached attributes.

joker-eph · 2024-10-06T22:55:02Z

Nit: these aren't "custom" attributes, but "discardable" attributes, please update title and description.

Also, it's not clear to me why this is correct in the full generality: how do you know that the transformation isn't invalidating some of these attributes potentially?

pashu123 · 2024-10-07T15:33:59Z

Nit: these aren't "custom" attributes, but "discardable" attributes, please update title and description.

Also, it's not clear to me why this is correct in the full generality: how do you know that the transformation isn't invalidating some of these attributes potentially?

I've updated the title. Thanks for the suggestion. This is not correct in full generality. For this patch, the only concern is canonicalize dropping the discardable attrs, which shouldn't be the case.

hanhanW · 2024-10-07T16:44:26Z

Also, it's not clear to me why this is correct in the full generality: how do you know that the transformation isn't invalidating some of these attributes potentially?

This is not correct in full generality. For this patch, the only concern is canonicalize dropping the discardable attrs, which shouldn't be the case.

+1, it is a canonicalization pattern which infers static shapes when possible. The transformations are not applied yet, so we don't drop the attributes.

This is also what's happening in Linalg transformations. Some old school patterns look at linalg ops with __linalg_transformation__ attribute, which uses it as a matcher. And we also have static shape inference in LinalgOp canonicalization patterns, which do not drop the discardable attributes.

MaheshRavishankar · 2024-10-07T22:07:10Z

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

-    Value newOp = rewriter.create<tensor::PackOp>(
-        loc, source, dest, packOp.getInnerDimsPos(), packOp.getMixedTiles(),
-        packOp.getPaddingValue(), packOp.getOuterDimsPerm());
+    auto clonedPackOp = cast<PackOp>(rewriter.clone(*packOp));


Sorry didnt get to review before it landed but you could use

auto discardableAttributes = getPrunedAttributeList(packOp, PackOp::getAttributesList()); clonedPackOp.setAttrs(discardableAttributes);

Thanks for the suggestion. Looks clean!

The Operation::getDiscardableAttrDictionary() method seems like a more direct implementation than getPrunedAttributeList() (why isn't this one living in mlir/IR by the way??)

That said, I didn't notice in the review that you're cloning the op here: why is that? Why aren't you just modifying it in place? Since you were doing rewriter.startOpModification I was assuming that this is what is happening.
When we can avoid recreating an operation and destroying the original, it's just more efficient.

Could you suggest how to modify it inplace without cloning in this case?

Shortest code should be:

rewriter.modifyOpInPlace(packOp, [&] { packOp.getSourceMutable().assign(source); packOp.getDestMutable().assign(dest); packOp.getResult();.setType(dest.getType()); });

The problem was adding tensor.cast after the modification.

llvm-project/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

Line 4348 in a649e8f

rewriter.replaceOpWithNewOp<tensor::CastOp>(

I can't replace it with tensor.pack.

I think it is not trivial to do in-place modification. There are two situations:

The dest tensor type is the same. In this case, we do not need a tensor.cast consumer.

The dest tensor type is changed. In this case, we need to create the tensor.cast op which makes the types consistent.

In the (1) situation, we can do in-place modification -- which is very simple.

In the (2) situation, it is not trivial because you need to replace the original op with the new tensor.cast op. If we do in-place modification, I don't see a trivial way to replace the op. Perhaps we can replace the uses of the tensor.pack ops with the new tensor.cast op, when it is the case.

IMO, cloning an op is cheap in this case. Instead of adding complex to logics, I'm +1 on cloning the op approach.

Note that this is also what we're doing for LinalgOps and it's been there for a couple years. I'm not saying that this is the correct way, but it's more like providing data points.

llvm-project/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

Lines 2561 to 2575 in 9f3c559

// Clone op.

Operation *newOp = clone(rewriter, linalgOp, resultTypes, newOperands);

SmallVector<Value> replacements;

replacements.reserve(newOp->getNumResults());

for (auto it : llvm::zip(linalgOp->getResults(), newOp->getResults())) {

Value newResult = std::get<1>(it);

Value oldResult = std::get<0>(it);

Type newType = newResult.getType();

Type oldType = oldResult.getType();

replacements.push_back(

(newType != oldType)

? rewriter.create<tensor::CastOp>(loc, oldType, newResult)

: newResult);

}

rewriter.replaceOp(linalgOp, replacements);

There are two situations:

The dest tensor type is the same. In this case, we do not need a tensor.cast consumer.

The dest tensor type is changed. In this case, we need to create the tensor.cast op which makes the types consistent.

The current code already does not try differentiate between these I believe: it always creates the cast, which is folded later if the types were matching.

I think it is not trivial to do in-place modification.

I don't understand the complexity you're foreseeing actually?
I sent a PR implementing it: #111593

#18718) The error was fixed with llvm/llvm-project#111261

llvmbot added mlir mlir:tensor labels Oct 5, 2024

pashu123 requested a review from hanhanW October 5, 2024 18:09

pashu123 mentioned this pull request Oct 5, 2024

canonicalize dropping the lowering_config attached to an operation. iree-org/iree#18697

Closed

hanhanW reviewed Oct 6, 2024

View reviewed changes

mlir/test/Dialect/Tensor/canonicalize.mlir Outdated Show resolved Hide resolved

[MLIR] Don't drop attached custom attributes

72e9153

The creation of pack op was dropping custom attached attributes.

pashu123 force-pushed the pack_dontdrop branch from 2b54cb8 to 72e9153 Compare October 6, 2024 06:19

pashu123 requested a review from hanhanW October 6, 2024 06:19

pashu123 changed the title ~~[MLIR] Don't drop attached custom attributes~~ [MLIR] Don't drop attached discardable attributes Oct 7, 2024

hanhanW approved these changes Oct 7, 2024

View reviewed changes

pashu123 merged commit 971b579 into llvm:main Oct 7, 2024
9 checks passed

MaheshRavishankar reviewed Oct 7, 2024

View reviewed changes

pashu123 mentioned this pull request Oct 8, 2024

Revert "Revert "[CPU] Enable tileDispatchUsingForall for mmt4d and co… iree-org/iree#18718

Merged

pashu123 added a commit to iree-org/iree that referenced this pull request Oct 8, 2024

Revert "Revert "[CPU] Enable tileDispatchUsingForall for mmt4d and co… (

6c095ed

#18718) The error was fixed with llvm/llvm-project#111261

Manewing mentioned this pull request Jun 17, 2025

Enable mac OS build of iree-runtime and iree-compiler RooflineAI/iree#75

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MLIR] Don't drop attached discardable attributes #111261

[MLIR] Don't drop attached discardable attributes #111261

Uh oh!

pashu123 commented Oct 5, 2024 •

edited

Loading

Uh oh!

llvmbot commented Oct 5, 2024

Uh oh!

llvmbot commented Oct 5, 2024

Uh oh!

Uh oh!

joker-eph commented Oct 6, 2024

Uh oh!

pashu123 commented Oct 7, 2024

Uh oh!

hanhanW commented Oct 7, 2024

Uh oh!

Uh oh!

MaheshRavishankar Oct 7, 2024

Uh oh!

pashu123 Oct 8, 2024

Uh oh!

joker-eph Oct 8, 2024

Uh oh!

pashu123 Oct 8, 2024

Uh oh!

joker-eph Oct 8, 2024 •

edited

Loading

Uh oh!

pashu123 Oct 8, 2024

Uh oh!

hanhanW Oct 8, 2024

Uh oh!

joker-eph Oct 8, 2024 •

edited

Loading

Uh oh!

Uh oh!

	// Clone op.
	Operation *newOp = clone(rewriter, linalgOp, resultTypes, newOperands);
	SmallVector<Value> replacements;
	replacements.reserve(newOp->getNumResults());
	for (auto it : llvm::zip(linalgOp->getResults(), newOp->getResults())) {
	Value newResult = std::get<1>(it);
	Value oldResult = std::get<0>(it);
	Type newType = newResult.getType();
	Type oldType = oldResult.getType();
	replacements.push_back(
	(newType != oldType)
	? rewriter.create<tensor::CastOp>(loc, oldType, newResult)
	: newResult);
	}
	rewriter.replaceOp(linalgOp, replacements);

[MLIR] Don't drop attached discardable attributes #111261

[MLIR] Don't drop attached discardable attributes #111261

Uh oh!

Conversation

pashu123 commented Oct 5, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Oct 5, 2024

Uh oh!

llvmbot commented Oct 5, 2024

Uh oh!

Uh oh!

joker-eph commented Oct 6, 2024

Uh oh!

pashu123 commented Oct 7, 2024

Uh oh!

hanhanW commented Oct 7, 2024

Uh oh!

Uh oh!

MaheshRavishankar Oct 7, 2024

Choose a reason for hiding this comment

Uh oh!

pashu123 Oct 8, 2024

Choose a reason for hiding this comment

Uh oh!

joker-eph Oct 8, 2024

Choose a reason for hiding this comment

Uh oh!

pashu123 Oct 8, 2024

Choose a reason for hiding this comment

Uh oh!

joker-eph Oct 8, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pashu123 Oct 8, 2024

Choose a reason for hiding this comment

Uh oh!

hanhanW Oct 8, 2024

Choose a reason for hiding this comment

Uh oh!

joker-eph Oct 8, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

pashu123 commented Oct 5, 2024 •

edited

Loading

joker-eph Oct 8, 2024 •

edited

Loading

joker-eph Oct 8, 2024 •

edited

Loading