[mlir] support !shape.value_shape when replace WithOp in OutlineShapeComputationPass.

[python3kgae]

Summary

For https://github.com/llvm/llvm-project/issues/60069

In case like:

%1 = shape.with_shape %arg1, %0 : !shape.value_shape, !shape.shape
%2 = shape.value_of %1 : tensor<?xf32>

cannot replace %2 with %arg1.
Transform it into

%2 = shape.value_of %arg1 : tensor<?xf32>

[jpienaar]

mlir/lib/Dialect/Shape/Transforms/OutlineShapeComputation.cpp

237

Could you add a comment here to explain?

[python3kgae]

add comments about the transform.

[python3kgae]

mlir/lib/Dialect/Shape/Transforms/OutlineShapeComputation.cpp

237

Done.

[jpienaar]

Looks good, as Mehdi mentioned using Fixes #60069 will result in closing the issue upon commit.

mlir/test/Dialect/Shape/arg_with_shape.mlir

10

Nit: lets use something like @repro_issue_60069 or some such to give additional context for future folks.

[python3kgae]

Fix clang-format missed by arc lint

mlir/lib/Dialect/Shape/Transforms/OutlineShapeComputation.cpp

moduleOp.walk([&](func::FuncOp funcOp) {
 
     DenseMap<Value, SmallVector<Operation *, 8>> clusters =
         constructClustersForEachShape(allWithOps, funcOp);
     constructShapeFunc(allWithOps, context, clusters, symbolTable,
                        dynShape2ShapeFunc, funcOp, shapeMappingAnalysis);
 
     for (shape::WithOp withOp : allWithOps) {
       Value value = withOp.getOperand();
-      for (Operation *user : withOp.getResult().getUsers()) {
-        if (Value valueOf = llvm::dyn_cast<shape::ValueOfOp>(user))
-          valueOf.replaceAllUsesExcept(value, withOp);
+      for (Operation *user : llvm::make_early_inc_range(withOp.getResult().getUsers())) {
+        if (auto valueOf = llvm::dyn_cast<shape::ValueOfOp>(user)) {
+          // For pattern like
+          //   %1 = shape.with_shape %arg1, %0
+          //   %2 = shape.value_of %1
+          // because shape.value doesn't care the shape, the shape.with_shape is
+          // redundant.
+          // If type of %arg1 and %2 has same type, just
+          //   replaced %2 with %arg1.
+          // If type of %arg1 has different type like !shape.value_shape,
+          // transform into
+          //   %2 = shape.value_of %arg1
+          if (valueOf.getType() == value.getType())
+            valueOf.replaceAllUsesWith(value);
+          else
+            valueOf.setOperand(value);
+        }
       }
     }
 
     // Apply patterns, note this also performs DCE.
     if (failed(applyPatternsAndFoldGreedily(funcOp, {})))
       return signalPassFailure();
   });
 }

mlir/test/Dialect/Shape/arg_with_shape.mlir

+// RUN: mlir-opt -outline-shape-computation -split-input-file %s 2>%t | FileCheck %s
+
+func.func @func1(%arg0: !shape.value_shape, %arg1: !shape.value_shape) -> !shape.shape {
+  %0 = shape.shape_of %arg0 : !shape.value_shape -> !shape.shape
+  %1 = shape.shape_of %arg1 : !shape.value_shape -> !shape.shape
+  %2 = shape.meet %0, %1 : !shape.shape, !shape.shape -> !shape.shape
+  return %2 : !shape.shape
+}
+// Make sure with_shape used by call not crash.
+// CHECK-LABEL:func.func @func
+func.func @func(%arg0: !shape.value_shape, %arg1: !shape.value_shape) -> !shape.shape {
+  %0 = shape.shape_of %arg0 : !shape.value_shape -> !shape.shape
+  %1 = shape.with_shape %arg1, %0 : !shape.value_shape, !shape.shape
+  %2 = call @func1(%arg0, %1) : (!shape.value_shape, !shape.value_shape) -> !shape.shape
+  return %2 : !shape.shape
+}

mlir/test/Dialect/Shape/outline-shape-computation.mlir

 // CHECK-DAG:   %[[V0:.*]] = shape_of %arg0 : tensor<?x4xf32> -> tensor<2xindex>
 // CHECK-DAG:   %[[V1:.*]] = shape_of %arg1 : tensor<?x4xf32> -> tensor<2xindex>
 // CHECK-DAG:   %[[V2:.*]] = get_extent %[[V0]], %c0 : tensor<2xindex>, index -> index
 // CHECK-DAG:   %[[V3:.*]] = get_extent %[[V1]], %c0 : tensor<2xindex>, index -> index
 // CHECK-DAG:   %[[V4:.*]] = arith.addi %[[V2]], %[[V3]] : index
 // CHECK-DAG:   %[[V5:.*]] = from_extents %[[V4]], %c4 : index, index
 // CHECK-DAG:   return %[[V5]] : !shape.shape
 
+// Make sure redundant with_shape is removed when with_shape input is !shape.value_shape.
+func.func @value_shape_with_shape(%arg0: !shape.value_shape, %arg1: !shape.value_shape) -> tensor<?xf32> {
+  %1 = shape.shape_of %arg0 : !shape.value_shape -> !shape.shape
+  %2 = shape.with_shape %arg1, %1 : !shape.value_shape, !shape.shape
+  %3 = shape.value_of %2 : tensor<?xf32>
+  return %3 : tensor<?xf32>
+}
+// CHECK-LABEL:func.func @value_shape_with_shape
+// CHECK-NEXT:%0 = shape.value_of %arg1 : tensor<?xf32>
+// CHECK-NEXT:return %0 : tensor<?xf32>