[MLIR] Fix shape inference in toy tutorial

[frgossen]

Summary

The implementation of shape inference in the toy tutorial did not conform to the correct algorithmic description.
The result was only correct because all operations appear to be processed in sequence.

[herhut]

Thanks!

mlir/examples/toy/Ch4/mlir/ShapeInferencePass.cpp

92

bit: operand -> operands. Here and below.

[frgossen]

Fix typo

mlir/examples/toy/Ch4/mlir/ShapeInferencePass.cpp

f.walk([&](mlir::Operation *op) {
         opWorklist.insert(op);
     });
 
     // Iterate on the operations in the worklist until all operations have been
     // inferred or no change happened (fix point).
     while (!opWorklist.empty()) {
       // Find the next operation ready for inference, that is an operation
       // with all operands already resolved (non-generic).
-      auto nextop = llvm::find_if(opWorklist, returnsDynamicShape);
+      auto nextop = llvm::find_if(opWorklist, allOperandsInferred);
       if (nextop == opWorklist.end())
         break;
 
       Operation *op = *nextop;
       opWorklist.erase(op);
 
       // Ask the operation to infer its output shapes.
       LLVM_DEBUG(llvm::dbgs() << "Inferring shape for: " << *op << "\n");
void runOnFunction() override {
     // If the operation worklist isn't empty, this indicates a failure.
     if (!opWorklist.empty()) {
       f.emitError("Shape inference failed, ")
           << opWorklist.size() << " operations couldn't be inferred\n";
       signalPassFailure();
     }
   }
 
+  /// A utility method that returns if the given operation has all of its
+  /// operands inferred.
+  static bool allOperandsInferred(Operation *op) {
+    return llvm::all_of(op->getOperandTypes(), [](Type operandType) {
+      return operandType.isa<RankedTensorType>();
+    });
+  }
+
   /// A utility method that returns if the given operation has a dynamically
   /// shaped result.
   static bool returnsDynamicShape(Operation *op) {
     return llvm::any_of(op->getResultTypes(), [](Type resultType) {
       return !resultType.isa<RankedTensorType>();
     });
   }
 };
 } // end anonymous namespace
 
 /// Create a Shape Inference pass.
 std::unique_ptr<mlir::Pass> mlir::toy::createShapeInferencePass() {
   return std::make_unique<ShapeInferencePass>();
 }

mlir/examples/toy/Ch5/mlir/ShapeInferencePass.cpp

f.walk([&](mlir::Operation *op) {
         opWorklist.insert(op);
     });
 
     // Iterate on the operations in the worklist until all operations have been
     // inferred or no change happened (fix point).
     while (!opWorklist.empty()) {
       // Find the next operation ready for inference, that is an operation
       // with all operands already resolved (non-generic).
-      auto nextop = llvm::find_if(opWorklist, returnsDynamicShape);
+      auto nextop = llvm::find_if(opWorklist, allOperandsInferred);
       if (nextop == opWorklist.end())
         break;
 
       Operation *op = *nextop;
       opWorklist.erase(op);
 
       // Ask the operation to infer its output shapes.
       LLVM_DEBUG(llvm::dbgs() << "Inferring shape for: " << *op << "\n");
void runOnFunction() override {
     // If the operation worklist isn't empty, this indicates a failure.
     if (!opWorklist.empty()) {
       f.emitError("Shape inference failed, ")
           << opWorklist.size() << " operations couldn't be inferred\n";
       signalPassFailure();
     }
   }
 
+  /// A utility method that returns if the given operation has all of its
+  /// operands inferred.
+  static bool allOperandsInferred(Operation *op) {
+    return llvm::all_of(op->getOperandTypes(), [](Type operandType) {
+      return operandType.isa<RankedTensorType>();
+    });
+  }
+
   /// A utility method that returns if the given operation has a dynamically
   /// shaped result.
   static bool returnsDynamicShape(Operation *op) {
     return llvm::any_of(op->getResultTypes(), [](Type resultType) {
       return !resultType.isa<RankedTensorType>();
     });
   }
 };
 } // end anonymous namespace
 
 /// Create a Shape Inference pass.
 std::unique_ptr<mlir::Pass> mlir::toy::createShapeInferencePass() {
   return std::make_unique<ShapeInferencePass>();
 }

mlir/examples/toy/Ch6/mlir/ShapeInferencePass.cpp

f.walk([&](mlir::Operation *op) {
         opWorklist.insert(op);
     });
 
     // Iterate on the operations in the worklist until all operations have been
     // inferred or no change happened (fix point).
     while (!opWorklist.empty()) {
       // Find the next operation ready for inference, that is an operation
       // with all operands already resolved (non-generic).
-      auto nextop = llvm::find_if(opWorklist, returnsDynamicShape);
+      auto nextop = llvm::find_if(opWorklist, allOperandsInferred);
       if (nextop == opWorklist.end())
         break;
 
       Operation *op = *nextop;
       opWorklist.erase(op);
 
       // Ask the operation to infer its output shapes.
       LLVM_DEBUG(llvm::dbgs() << "Inferring shape for: " << *op << "\n");
void runOnFunction() override {
     // If the operation worklist isn't empty, this indicates a failure.
     if (!opWorklist.empty()) {
       f.emitError("Shape inference failed, ")
           << opWorklist.size() << " operations couldn't be inferred\n";
       signalPassFailure();
     }
   }
 
+  /// A utility method that returns if the given operation has all of its
+  /// operands inferred.
+  static bool allOperandsInferred(Operation *op) {
+    return llvm::all_of(op->getOperandTypes(), [](Type operandType) {
+      return operandType.isa<RankedTensorType>();
+    });
+  }
+
   /// A utility method that returns if the given operation has a dynamically
   /// shaped result.
   static bool returnsDynamicShape(Operation *op) {
     return llvm::any_of(op->getResultTypes(), [](Type resultType) {
       return !resultType.isa<RankedTensorType>();
     });
   }
 };
 } // end anonymous namespace
 
 /// Create a Shape Inference pass.
 std::unique_ptr<mlir::Pass> mlir::toy::createShapeInferencePass() {
   return std::make_unique<ShapeInferencePass>();
 }

mlir/examples/toy/Ch7/mlir/ShapeInferencePass.cpp

f.walk([&](mlir::Operation *op) {
         opWorklist.insert(op);
     });
 
     // Iterate on the operations in the worklist until all operations have been
     // inferred or no change happened (fix point).
     while (!opWorklist.empty()) {
       // Find the next operation ready for inference, that is an operation
       // with all operands already resolved (non-generic).
-      auto nextop = llvm::find_if(opWorklist, returnsDynamicShape);
+      auto nextop = llvm::find_if(opWorklist, allOperandsInferred);
       if (nextop == opWorklist.end())
         break;
 
       Operation *op = *nextop;
       opWorklist.erase(op);
 
       // Ask the operation to infer its output shapes.
       LLVM_DEBUG(llvm::dbgs() << "Inferring shape for: " << *op << "\n");
void runOnFunction() override {
     // If the operation worklist isn't empty, this indicates a failure.
     if (!opWorklist.empty()) {
       f.emitError("Shape inference failed, ")
           << opWorklist.size() << " operations couldn't be inferred\n";
       signalPassFailure();
     }
   }
 
+  /// A utility method that returns if the given operation has all of its
+  /// operands inferred.
+  static bool allOperandsInferred(Operation *op) {
+    return llvm::all_of(op->getOperandTypes(), [](Type operandType) {
+      return operandType.isa<RankedTensorType>();
+    });
+  }
+
   /// A utility method that returns if the given operation has a dynamically
   /// shaped result.
   static bool returnsDynamicShape(Operation *op) {
     return llvm::any_of(op->getResultTypes(), [](Type resultType) {
       return !resultType.isa<RankedTensorType>();
     });
   }
 };
 } // end anonymous namespace
 
 /// Create a Shape Inference pass.
 std::unique_ptr<mlir::Pass> mlir::toy::createShapeInferencePass() {
   return std::make_unique<ShapeInferencePass>();
 }