[StaticRuntime] Support a new pattern for ClipRangesToGatherToOffsets (pytorch#146931)

Summary:
Support the following new pattern for ClipRangesToGatherToOffsets:

Before optimization:
```
%18267 : Tensor, %18268 : Tensor = fb::clip_ranges_gather(%int_77.1, %getitem_2484.1, %493)
%getattr_368.1 : int = prim::dtype(%18267)
%to_443.1 : Tensor = aten::to(%18268, %getattr_368.1, %self._maybe_compute_kjt_to_jt_dict.is_weighted, %self._maybe_compute_kjt_to_jt_dict.is_weighted)
%lengths_to_offsets_490.1 : Tensor = fb::lengths_to_offsets(%to_443.1, %8)
```

After optimization:
```
%18297 : int = prim::dtype(%int_77.1)
%18298 : Tensor, %18299 : Tensor = fb::clip_ranges_gather_to_offsets(%int_77.1, %getitem_2484.1, %493, %8, %18297)
```

Reviewed By: garroud

Differential Revision: D69373835

Pull Request resolved: pytorch#146931
Approved by: https://github.com/hanyilou123

Author: coufon

torch/csrc/jit/runtime/static/passes.cpp

@@ -255,6 +255,27 @@ namespace {
   fuse.runOnGraph(graph);
 }
 
+// Similar to ClipRangesToGatherToOffsets, but for the case where type of aten::to is from
+// gather_ranges's data output instead of the graph input.
+[[maybe_unused]] void ClipRangesToGatherToOffsetsV2(
+    std::shared_ptr<torch::jit::Graph>& graph) {
+  std::string pattern = R"IR(
+    graph(%a, %b, %c, %d, %to0_in0):
+        %y0 : Tensor, %y1 : Tensor = fb::clip_ranges_gather(%a, %b, %c)
+        %y0_type : int = prim::dtype(%y0)
+        %y2 : Tensor = aten::to(%y1, %y0_type, %to0_in0, %to0_in0)
+        %y3 : Tensor = fb::lengths_to_offsets(%y2, %d)
+        return (%y3, %y0))IR";
+  std::string fused_pattern = R"IR(
+    graph(%a, %b, %c, %d, %to0_in0):
+        %a_type : int = prim::dtype(%a)
+        %y0 : Tensor, %y1 : Tensor = fb::clip_ranges_gather_to_offsets(%a, %b, %c, %d, %a_type)
+        return (%y1, %y0))IR";
+  SubgraphRewriter fuse;
+  fuse.RegisterRewritePattern(pattern, fused_pattern);
+  fuse.runOnGraph(graph);
+}
+
 [[maybe_unused]] void ToLengthsToOffsets(
     std::shared_ptr<torch::jit::Graph>& graph) {
   std::string pattern = R"IR(
@@ -389,7 +410,8 @@ void FuseInferenceOpsForSparseNN(std::shared_ptr<torch::jit::Graph>& graph) {
     // prioritize clip_ranges+gather_ranges+sigrid_hash fusion over
     // clip_ranges+gather_ranges
     ClipRangesGather(graph);
-
+    // Must run before ClipRangesToGatherToOffsets.
+    ClipRangesToGatherToOffsetsV2(graph);
     ClipRangesToGatherToOffsets(graph);
   }