[Flow] Fix TensorSliceOp::fold for parameter attributes (#24762)

[Flow] Fix tensor folding for parameter attributes
    
TensorSliceOp::fold and TensorUpdateOp::fold assume that fully folded
tensor operands are ElementsAttr values. However, #flow.parameter.named
    folds to a parameter attribute instead, resulting in invalid casts.
    
Skip folding when the required operands are not ElementsAttr values and
add regression coverage for slicing and updating parameter-backed tensors.
    
Fixes #24751

Signed-off-by: zhangpan <98025960+zhangpan2001@users.noreply.github.com>

compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp

diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp
index 311120f..bea682a 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp
@@ -1023,7 +1023,10 @@
       return {};
     }
     // Fully constant arguments so we can perform the slice here.
-    auto tensor = cast<ElementsAttr>(operands.getSource());
+    auto tensor = dyn_cast<ElementsAttr>(operands.getSource());
+    if (!tensor) {
+      return {};
+    }
     int64_t rank = cast<ShapedType>(getSource().getType()).getRank();
     auto start =
         llvm::map_to_vector(operands.getStartIndices(), [](Attribute value) {
@@ -1153,9 +1156,12 @@
       llvm::count(operands.getStartIndices(), nullptr) == 0;
   if (operands.getUpdate() && operands.getTarget() && allIndicesConstant) {
     // Fully constant arguments so we can perform the update here.
-    return tensorUpdate(cast<ElementsAttr>(operands.getUpdate()),
-                        cast<ElementsAttr>(operands.getTarget()),
-                        operands.getStartIndices());
+    auto update = dyn_cast<ElementsAttr>(operands.getUpdate());
+    auto target = dyn_cast<ElementsAttr>(operands.getTarget());
+    if (!update || !target) {
+      return {};
+    }
+    return tensorUpdate(update, target, operands.getStartIndices());
   } else {
     // Replace the entire tensor when the sizes match.
     auto updateType = cast<ShapedType>(getUpdate().getType());

compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir

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 af7731f..aa69163 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
@@ -690,6 +690,21 @@
 
 // -----
 
+// CHECK-LABEL: @noFoldSliceParameterAttr
+util.func public @noFoldSliceParameterAttr() -> tensor<16xf32> {
+  // CHECK: %[[C:.+]] = flow.tensor.constant #flow.parameter.named<"m"::"w"> : tensor<20xf32>
+  %0 = flow.tensor.constant #flow.parameter.named<"m"::"w"> : tensor<20xf32>
+  %c0 = arith.constant 0 : index
+  %c16 = arith.constant 16 : index
+  // CHECK: %[[S:.+]] = flow.tensor.slice %[[C]]
+  // CHECK-SAME:             : tensor<20xf32> -> tensor<16xf32>
+  %1 = flow.tensor.slice %0[%c0 for %c16] : tensor<20xf32> -> tensor<16xf32>
+  // CHECK-NEXT: util.return %[[S]]
+  util.return %1 : tensor<16xf32>
+}
+
+// -----
+
 // CHECK-LABEL: @sliceOfSliceStatic2D
 util.func public @sliceOfSliceStatic2D(%src: tensor<3x3xf32>) -> tensor<1x1xf32> {
   %c0 = arith.constant 0 : index
@@ -857,6 +872,21 @@
 
 // -----
 
+// CHECK-LABEL: @noFoldUpdateParameterAttr
+util.func public @noFoldUpdateParameterAttr() -> tensor<4xf32> {
+  // CHECK: %[[UPDATE:.+]] = flow.tensor.constant dense<1.000000e+00> : tensor<1xf32>
+  %0 = flow.tensor.constant dense<1.0> : tensor<1xf32>
+  // CHECK: %[[TARGET:.+]] = flow.tensor.constant #flow.parameter.named<"m"::"w"> : tensor<4xf32>
+  %1 = flow.tensor.constant #flow.parameter.named<"m"::"w"> : tensor<4xf32>
+  %c0 = arith.constant 0 : index
+  // CHECK: %[[RESULT:.+]] = flow.tensor.update %[[UPDATE]], %[[TARGET]]
+  %2 = flow.tensor.update %0, %1[%c0] : tensor<1xf32> -> tensor<4xf32>
+  // CHECK-NEXT: util.return %[[RESULT]]
+  util.return %2 : tensor<4xf32>
+}
+
+// -----
+
 // CHECK-LABEL: @updateReplace
 util.func public @updateReplace(%arg0 : tensor<4xi32>, %arg1 : tensor<4xi32>) -> tensor<4xi32> {
   %c0 = arith.constant 0 : index