fix(TensorSliceOp::fold): ignore DenseResourceElementsAttr (iree-org#…

…19182) Fixes: iree-org#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 <[email protected]>
nithinsubbiah · Nov 20, 2024 · 1aada43 · 1aada43
1 parent 901db6e
commit 1aada43
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git 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
@@ -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>
+}
 
 // -----