fix(TensorSliceOp::fold): ignore DenseResourceElementsAttr (#19182)
Fixes: #17566
Certain cases involve constants who have been converted to resources.
The following change simply ignores these in the folding operation.
One possible future improvement is to perform the slicing at compile
time as we do for the DenseElements, but for now to unblock failing
models this is good enough
Signed-off-by: Christopher McGirr <mcgirr@roofline.ai>
diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp
index 3219c6e..88ac020 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp
@@ -1226,6 +1226,10 @@
OpFoldResult TensorSliceOp::fold(FoldAdaptor operands) {
if (llvm::count(operands.getOperands(), nullptr) == 0) {
+ // Ignore DenseResources for now and do not perfom folding on them.
+ if (isa<DenseResourceElementsAttr>(operands.getSource())) {
+ return {};
+ }
// Fully constant arguments so we can perform the slice here.
auto tensor = llvm::cast<ElementsAttr>(operands.getSource());
int64_t rank = llvm::cast<ShapedType>(getSource().getType()).getRank();
diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir b/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir
index 6e92e0b..d77efbf 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir
@@ -539,6 +539,22 @@
// CHECK-NEXT: util.return %[[C]]
util.return %1 : tensor<1x2x3xi32>
}
+// -----
+
+// CHECK-LABEL: @sliceDenseResourceAttr
+util.func public @sliceDenseResourceAttr() -> tensor<1x1x3xf32> {
+ %0 = arith.constant dense_resource<__elided__> : tensor<1x2x3xf32>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c3 = arith.constant 3 : index
+ %1 = flow.tensor.slice %0[%c0, %c0, %c0 for %c1, %c1, %c3] : tensor<1x2x3xf32> -> tensor<1x1x3xf32>
+ // CHECK: %[[C:.+]] = arith.constant dense_resource
+ // CHECK-SAME: : tensor<1x2x3xf32>
+ // CHECK: %[[S:.+]] = flow.tensor.slice %[[C]]
+ // CHECK-SAME: -> tensor<1x1x3xf32>
+ // CHECK-NEXT: util.return %[[S]]
+ util.return %1 : tensor<1x1x3xf32>
+}
// -----
