[WIP] [stream] Add resource refcounting analysis pass

compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td

@@ -1869,6 +1869,61 @@ def Stream_AsyncAllocaOp : Stream_Op<"async.alloca", [
   let hasCanonicalizer = 1;
 }
 
+
+def Stream_ResourceAddRefOp : Stream_Op<"async.resource.add_ref"> {
+  let summary = [{inserts ref counting for a resource}];
+  let description = [{
+    During whole-progam analysis, this op reference counts a resource
+    and propogates through the program.
+  }];
+
+  let arguments = (ins
+    Stream_AnyStreamResource:$resource,
+    I64:$count
+  );
+
+  let assemblyFormat = [{
+    $resource attr-dict `:` type($resource) `,` $count
+  }];
+
+}
+
+def Stream_ResourceDecRefOp : Stream_Op<"async.resource.dec_ref"> {
+  let summary = [{decrements ref counting for a resource}];
+  let description = [{
+    During whole-progam analysis, this op reference counts a resource
+    via tracking resource liveness and propogates through the program.
+  }];
+
+  let arguments = (ins
+    Stream_AnyStreamResource:$resource,
+    I64:$count
+  );
+
+  let assemblyFormat = [{
+    $resource attr-dict `:` type($resource) `,` $count
+  }];
+
+}
+
+
+def Stream_ResourceLastUseOp : Stream_Op<"async.resource.last_use"> {
+  let summary = [{Last-use witness op}];
+  let description = [{
+    Used to mark last-use witness op that turns into dealloca when
+    refcount drops to 0
+  }];
+
+  let arguments = (ins
+    Stream_AnyStreamResource:$resource
+  );
+
+  let assemblyFormat = [{
+    $resource attr-dict `:` type($resource)
+  }];
+
+}
+
 def Stream_AsyncConstantOp : Stream_PureOp<"async.constant", [
   Stream_AffinityOp,
   Stream_AsyncPhaseOp,

compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel

@@ -36,6 +36,7 @@ iree_compiler_cc_library(
         "Passes.h.inc",
         "PropagateTimepoints.cpp",
         "RefineUsage.cpp",
+        "ResourceRefCounting.cpp",
         "ScheduleAllocation.cpp",
         "ScheduleConcurrency.cpp",
         "ScheduleExecution.cpp",

compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt

@@ -37,6 +37,7 @@ iree_cc_library(
     "Passes.h.inc"
     "PropagateTimepoints.cpp"
     "RefineUsage.cpp"
+    "ResourceRefCounting.cpp"
     "ScheduleAllocation.cpp"
     "ScheduleConcurrency.cpp"
     "ScheduleExecution.cpp"

compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp

@@ -266,6 +266,11 @@ void buildStreamCmdPassPipeline(OpPassManager &passManager,
   // After propagation many resource SSA values can be deduped or folded by the
   // cleanup patterns.
   passManager.addPass(IREE::Util::createPropagateSubrangesPass());
+
+  // Tracks cross-device resource liveness, performs whole program
+  // analysis, reference counting
+  passManager.addPass(IREE::Stream::createResourceRefCountingPass());
+
   buildStreamCleanupPassPipeline(passManager, transformOptions);
 
   // TODO(benvanik): outline streams (ala dispatch regions). Note that we may

compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td

@@ -210,6 +210,20 @@ def RefineUsagePass :
   ];
 }
 
+
+def ResourceRefCountingPass :
+    Pass<"iree-stream-resource-refcounting", "mlir::ModuleOp"> {
+  let summary = "Tracks resource liveness, and inserts deallcas where appropriate.";
+  let description = [{
+    Performs whole-program analysis to assign refcounting to
+    `!stream.resource<*>`
+     Upon completion all resources should have reference counted and deallocas identified
+  }];
+  let dependentDialects = [
+    "IREE::Stream::StreamDialect",
+  ];
+}
+
 //===----------------------------------------------------------------------===//
 // Stream formation and scheduling
 //===----------------------------------------------------------------------===//