From 1aada43632919ee71ad1c0b26d179b2f29be0ab9 Mon Sep 17 00:00:00 2001
From: Christopher McGirr <7071833+chrsmcgrr@users.noreply.github.com>
Date: Wed, 20 Nov 2024 17:24:37 +0100
Subject: [PATCH] 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>
---
 .../compiler/Dialect/Flow/IR/FlowOpFolders.cpp   |  4 ++++
 .../Dialect/Flow/IR/test/tensor_folding.mlir     | 16 ++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp
index 3219c6e86186..88ac0205214f 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 @@ static ElementsAttr tensorSlice(ElementsAttr tensor, uint64_t dim,
 
 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 6e92e0b7f85e..d77efbf597ed 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 @@ util.func public @sliceConst3D() -> tensor<1x2x3xi32> {
   // 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>
+}
 
 // -----