[StaticRuntime] Fusion pass for ClipRanges/GatherRanges/LengthsToOffsets

Hao Lu · facebook-github-bot · commit 78c4e98191c7 · 2020-12-09T13:35:25.000-08:00
Differential Revision: D25388512

fbshipit-source-id: 078d34cf0fb79d82050cfd00b2dba0bbb8bb3f2c
diff --git a/torch/csrc/jit/runtime/static/impl.cpp b/torch/csrc/jit/runtime/static/impl.cpp
@@ -13,8 +13,8 @@
 
 namespace torch {
 namespace jit {
-
 namespace {
+
 void OptimizeGraph(std::shared_ptr<torch::jit::Graph>& graph) {
   Inline(*graph);
   ConstantPropagation(graph);
diff --git a/torch/csrc/jit/runtime/static/passes.cpp b/torch/csrc/jit/runtime/static/passes.cpp
@@ -71,11 +71,30 @@ void ConcatBatchMatMulBatchGather(std::shared_ptr<torch::jit::Graph>& graph) {
   fuse.runOnGraph(graph);
 }
 
+void ClipRangesGatherRangesLengthsToOffsets(
+    std::shared_ptr<torch::jit::Graph>& graph) {
+  // TODO:: check restrictions for inputs; outputs not used elsewhere
+  std::string pattern = R"IR(
+    graph(%a, %b, %c, %d):
+        %y0 : Tensor = fb::clip_ranges(%b, %c)
+        %y1 : Tensor, %y2 : Tensor = fb::gather_ranges(%a, %y0)
+        %y3 : Tensor = fb::lengths_to_offsets(%y2, %d)
+        return (%y3, %y1))IR";
+  std::string fused_pattern = R"IR(
+    graph(%a, %b, %c, %d):
+        %y0 : Tensor, %y1 : Tensor = fb::clip_ranges_gather_lengths_to_offsets(%a, %b, %c, %d)
+        return (%y1, %y0))IR";
+  SubgraphRewriter fuse;
+  fuse.RegisterRewritePattern(pattern, fused_pattern);
+  fuse.runOnGraph(graph);
+}
+
 void FuseInferenceOpsForSparseNN(std::shared_ptr<torch::jit::Graph>& graph) {
 #ifdef FBCODE_CAFFE2
   ConcatAddMulReplaceNaNClip(graph);
   CastedBatchOneHotLengths(graph);
   ConcatBatchMatMulBatchGather(graph);
+  ClipRangesGatherRangesLengthsToOffsets(graph);
 #endif
 }
 
