Changes to refine types

- Add `!torch.optional` knowledge tracking
- Changes to improve type propagation for branches and terminators. See
examples in `refine-types-branch.mlir`
- Refator to separate handling of different ops from `visitOperation`
- Add refine types for a few new ops

include/npcomp/Dialect/Torch/IR/TorchOps.td

@@ -461,6 +461,7 @@ def Torch_PrimLoopOp : Torch_Op<"prim.Loop", [
 }
 
 def Torch_PrimLoopConditionOp : Torch_Op<"prim.Loop.condition", [
+    DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface>,
     Terminator,
     HasParent<"::mlir::NPCOMP::Torch::PrimLoopOp">]> {
   let summary = "yield-like terminator for torch.prim.Loop";

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -252,6 +252,17 @@ void PrimLoopOp::getSuccessorRegions(
   regions.emplace_back(getResults());
 }
 
+//===----------------------------------------------------------------------===//
+// PrimLoopConditionOp
+//===----------------------------------------------------------------------===//
+
+MutableOperandRange
+PrimLoopConditionOp::getMutableSuccessorOperands(Optional<unsigned> index) {
+  // Pass all operands except the condition to the successor which is the
+  // parent loop op.
+  return iterArgsMutable();
+}
+
 //===----------------------------------------------------------------------===//
 // PrimIfOp
 //===----------------------------------------------------------------------===//
@@ -357,10 +368,11 @@ bool DerefineOp::areCastCompatible(mlir::TypeRange inputs,
 void DerefineOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                              MLIRContext *context) {
   patterns.add(+[](DerefineOp op, PatternRewriter &rewriter) {
-    // TODO: Extend RefineTypes for this case and delete this canonicalization,
-    // since we don't want control flow or calls to randomly block this fold
-    // (this canonicalization pattern makes the compiler brittle to control flow
-    // and calls).
+    // TODO: This pattern should be removed because type refine does a better
+    // job dealing with control flow. However, removing this would expose an
+    // issue with ReduceOpVariants. DerefineOp doesn't have value semantics and
+    // if not removed eagerly by canonicalizer would prevent ReduceOpVariants
+    // from converting certain tensors value semantics.
     bool allAllowRefinement =
         llvm::all_of(op.getResult().getUsers(), allowsTypeRefinement);
     if (!allAllowRefinement)