Update pass documentation

Author: nicolasvasilache

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td

@@ -183,19 +183,38 @@ def ResolveShapedTypeResultDimsPass : Pass<"resolve-shaped-type-result-dims"> {
 }
 
 def ReifyResultShapesPass : Pass<"reify-result-shapes"> {
-  let summary = "Reifies the results of all `ReifyRankedShapedTypeOpInterface` operations";
+  let summary ="Reifies the results of `tensor::PadOp` and `tensor::ConcatOp`.";
   let description = [{
-    This pass reifies the shapes of every `ReifyRankedShapedTypeOpInterface`
-    operation with ranked `memref` and `tensor` results. Replacing the
-    operations with their reified versions, and inserting casts when results
-    shapes are updated.
+    This pass reifies the shapes of a subset of `ReifyRankedShapedTypeOpInterface`
+    ops with `tensor` results.
+    
+    The pass currently only supports result shape type reification for:
+      - tensor::PadOp
+      - tensor::ConcatOp
+    It addresses a representation gap where implicit op semantics are needed to
+    infer static result types from dynamic operands.
+    But it does so by using `ReifyRankedShapedTypeOpInterface` as the source of
+    truth rather than the op itself. As a consequence, this cannot generalize
+    today.
+
+    TODO: in the future, we should consider coupling this information with op
+    "transfer functions" (e.g. `IndexingMapOpInterface`) to provide a source of
+    truth that can work across result shape inference, canonicalization and op
+    verifiers.
+
+    The pass replaces the operations with their reified versions, when more
+    static information can be derived, and inserts casts when results shapes
+    are updated.
 
     Example:
     ```mlir
     #map = affine_map<(d0) -> (-d0 + 256)>
-    func.func @func(%arg0: f32, %arg1: index, %arg2: tensor<64x?x64xf32>) -> tensor<1x?x64xf32> {
+    func.func @func(%arg0: f32, %arg1: index, %arg2: tensor<64x?x64xf32>)
+        -> tensor<1x?x64xf32>
+    {
       %0 = affine.apply #map(%arg1)
-      %extracted_slice = tensor.extract_slice %arg2[0, 0, 0] [1, %arg1, 64] [1, 1, 1] : tensor<64x?x64xf32> to tensor<1x?x64xf32>
+      %extracted_slice = tensor.extract_slice %arg2[0, 0, 0] [1, %arg1, 64] [1, 1, 1]
+        : tensor<64x?x64xf32> to tensor<1x?x64xf32>
       %padded = tensor.pad %extracted_slice low[0, 0, 0] high[0, %0, 0] {
       ^bb0(%arg3: index, %arg4: index, %arg5: index):
         tensor.yield %arg0 : f32
@@ -205,9 +224,12 @@ def ReifyResultShapesPass : Pass<"reify-result-shapes"> {
 
     // mlir-opt --reify-result-shapes
     #map = affine_map<()[s0] -> (-s0 + 256)>
-    func.func @func(%arg0: f32, %arg1: index, %arg2: tensor<64x?x64xf32>) -> tensor<1x?x64xf32> {
+    func.func @func(%arg0: f32, %arg1: index, %arg2: tensor<64x?x64xf32>)
+        -> tensor<1x?x64xf32>
+    {
       %0 = affine.apply #map()[%arg1]
-      %extracted_slice = tensor.extract_slice %arg2[0, 0, 0] [1, %arg1, 64] [1, 1, 1] : tensor<64x?x64xf32> to tensor<1x?x64xf32>
+      %extracted_slice = tensor.extract_slice %arg2[0, 0, 0] [1, %arg1, 64] [1, 1, 1]
+        : tensor<64x?x64xf32> to tensor<1x?x64xf32>
       %padded = tensor.pad %extracted_slice low[0, 0, 0] high[0, %0, 0] {
       ^bb0(%arg3: index, %arg4: index, %arg5: index):
         tensor.yield %arg0 : f32

mlir/lib/Dialect/MemRef/Transforms/ReifyResultShapes.cpp

@@ -17,6 +17,7 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "llvm/Support/InterleavedRange.h"
 
@@ -134,8 +135,10 @@ struct ReifyResultShapesPass final
 void ReifyResultShapesPass::runOnOperation() {
   SmallVector<ReifyRankedShapedTypeOpInterface> ops;
   getOperation()->walk([&](ReifyRankedShapedTypeOpInterface op) {
-    // Some ops have rigid type checkers and need to update their operands.
-    // Only admit the ones that are explicitly supported for now.
+    // Handle ops that are not DPS and that do not carry an tied operand shapes.
+    // For now, limit to tensor::PadOp and tensor::ConcatOp.
+    if (isa<DestinationStyleOpInterface>(op.getOperation()))
+      return;
     if (!isa<tensor::PadOp, tensor::ConcatOp>(op.getOperation()))
       return;
     ops.push_back(op);