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";
@@ -641,8 +642,6 @@ def Torch_DerefineOp : Torch_Op<"derefine", [
   let assemblyFormat = [{
     $operand attr-dict `:` type($operand) `to` type($result)
   }];
-
-  let hasCanonicalizer = 1;
 }
 
 def Torch_OperatorOp : Torch_Op<"operator", [

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -248,6 +248,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
 //===----------------------------------------------------------------------===//
@@ -350,22 +361,6 @@ bool DerefineOp::areCastCompatible(mlir::TypeRange inputs,
   return isValidSubtype(inputs[0], outputs[0]);
 }
 
-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).
-    bool allAllowRefinement =
-        llvm::all_of(op.getResult().getUsers(), allowsTypeRefinement);
-    if (!allAllowRefinement)
-      return failure();
-    rewriter.replaceOp(op, op.getOperand());
-    return success();
-  });
-}
-
 template <typename OpTy>
 static OpFoldResult atenIsOrIsNotFoldHelper(OpTy op, bool equalIsTrue) {
   Type lhsType = op.self().getType();