Repository: alibaba/BladeDISC (language: python)

Issue:
enhancement bert training op clustering
To utilize BladeDISC accelerating bert training, we need to cluster more op into disc, as the experiment, we found 244 `FusionGroups` in the first iteration with LTC, to reduce the clusters, we should:

- enhancement op coverage, #441 lists all ops that required, I just list some priority ones:
    - native_layer_norm
    - _softmax
    - _log_softmax
    - _log_softmax_backward_data
    - _nll_loss_backward
    - _nll_loss_forward
    - tanh_backward
    - native_dropout_backward
    - gelu_backward
- enhancement shape propagation
    - native_layer_norm
    - native_dropout
    - native_dropout_backward
    

[here](https://gist.github.com/Yancey1989/2ffbdea3c48baa0243e56d436eecfc98) list the 224 FusionGroups


Discussion / hints:

Fix patch:
diff --git a/pytorch_blade/src/compiler/jit/torch/op_registry.cpp b/pytorch_blade/src/compiler/jit/torch/op_registry.cpp
index 56cf308da26..cfab9f89296 100644
--- a/pytorch_blade/src/compiler/jit/torch/op_registry.cpp
+++ b/pytorch_blade/src/compiler/jit/torch/op_registry.cpp
@@ -30,9 +30,13 @@ using namespace torch::jit;
 //     Knowing the type and device of weights or biases is usually enough to
 //     infer the output type.
 std::shared_ptr<OperatorSet> nn_ops_first_input_preserving() {
-  std::shared_ptr<OperatorSet> ops = std::make_shared<OperatorSet>(OperatorSet{
+  // clang-format off
+  std::shared_ptr<OperatorSet> ops = std::make_shared<OperatorSet>(
+    OperatorSet {
+#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION >= 12
       "aten::_log_softmax_backward_data(Tensor grad_output, Tensor output, int dim, ScalarType input_dtype) -> Tensor",
-      "aten::_softmax_backward_data(Tensor grad_output, Tensor output, int dim, ScalarType input_dtype) -> Tensor",
+      "aten::native_dropout_backward(Tensor grad_output, Tensor mask, float scale) -> Tensor",
+#endif
       "aten::batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps, bool cudnn_enabled) -> Tensor",
       "aten::layer_norm(Tensor input, int[] normalized_shape, Tensor? weight=None, Tensor? bias=None, float eps=1e-05, bool cudnn_enable=True) -> Tensor",
       "aten::conv1d(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int groups) -> Tensor",
@@ -43,8 +47,10 @@ std::shared_ptr<OperatorSet> nn_ops_first_input_preserving() {
       "aten::conv_transpose2d(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] output_padding, int groups, int[] dilation) -> Tensor",
       "aten::conv_transpose3d(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] output_padding, int groups, int[] dilation) -> Tensor",
       "aten::convolution(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, bool transposed, int[] output_padding, int groups) -> Tensor",
-      "aten::_convolution(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, bool transposed, int[] output_padding, int groups, bool benchmark, bool deterministic, bool cudnn_enabled) -> Tensor", // deprecated _convolution
+#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSOIN > 6
       "aten::_convolution(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, bool transposed, int[] output_padding, int groups, bool benchmark, bool deterministic, bool cudnn_enabled, bool allow_tf32) -> Tensor",
+#endif
+      "aten::_convolution(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, bool transposed, int[] output_padding, int groups, bool benchmark, bool deterministic, bool cudnn_enabled) -> Tensor",
       "aten::adaptive_avg_pool1d(Tensor self, int[] output_size) -> Tensor",
       "aten::adaptive_avg_pool2d(Tensor self, int[] output_size) -> Tensor",
       "aten::adaptive_avg_pool3d(Tensor self, int[] output_size) -> Tensor",
@@ -56,16 +62,15 @@ std::shared_ptr<OperatorSet> nn_ops_first_input_preserving() {
       "aten::max_pool3d(Tensor self, int[] kernel_size, int[] stride, int[] padding, int[] dilation, bool ceil_mode) -> Tensor",
       "aten::max_unpool2d(Tensor self, Tensor indices, int[] output_size) -> Tensor",
       "aten::max_unpool3d(Tensor self, Tensor indices, int[] output_size, int[] stride, int[] padding) -> Tensor",
-      "aten::native_dropout_backward(Tensor grad_output, Tensor mask, float scale) -> Tensor",
       "aten::nll_loss_backward(Tensor grad_output, Tensor self, Tensor target, Tensor? weight, int reduction, int ignore_index, Tensor total_weight) -> Tensor",
       "aten::reflection_pad1d(Tensor self, int[] padding) -> Tensor",
       "aten::reflection_pad2d(Tensor self, int[] padding) -> Tensor",
+#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION >= 12
       "aten::reflection_pad3d(Tensor self, int[] padding) -> Tensor",
+#endif
       "aten::replication_pad1d(Tensor self, int[] padding) -> Tensor",
       "aten::replication_pad2d(Tensor self, int[] padding) -> Tensor",
-      "aten::replication_pad3d(Tensor self, int[] padding) -> Tensor",
-      "aten::tanh_backward(Tensor grad_output, Tensor output) -> Tensor",
-      "aten::upsample_bilinear2d(Tensor self, int[] output_size, bool align_corners, float? scales_h, float? scales_w) -> Tensor",
+      "aten::replication_pad3d(Tensor self, int[] padding) -> Tensor", "aten::tanh_backward(Tensor grad_output, Tensor output) -> Tensor", "aten::upsample_bilinear2d(Tensor self, int[] output_size, bool align_corners, float? scales_h, float? scales_w) -> Tensor",
       "aten::upsample_linear1d(Tensor self, int[] output_size, bool align_corners, float? scales) -> Tensor",
       "aten::upsample_nearest1d(Tensor self, int[] output_size, float? scales) -> Tensor",
       "aten::upsample_nearest2d(Tensor self, int[] output_size, float? scales_h, float? scales_w) -> Tensor",
@@ -73,10 +78,12 @@ std::shared_ptr<OperatorSet> nn_ops_first_input_preserving() {
       "aten::upsample_trilinear3d(Tensor self, int[] output_size, bool align_corners, float? scales_d, float? scales_h, float? scales_w) -> Tensor",
       "aten::prelu(Tensor self, Tensor weight) -> Tensor",
 
-      // Added because Hardswish is really hard to convert to metatensors
+      // Added because Hardswish is really hard to convert to
+      // metatensors
       "aten::hardswish(Tensor self) -> Tensor",
       "aten::hardswish_(Tensor self) -> Tensor",
   });
+  // clang-format on
   return ops;
 };
 
diff --git a/pytorch_blade/src/compiler/jit/torch/shape_analysis.cpp b/pytorch_blade/src/compiler/jit/torch/shape_analysis.cpp
index 705e26a97c6..1af25ec5d53 100644
--- a/pytorch_blade/src/compiler/jit/torch/shape_analysis.cpp
+++ b/pytorch_blade/src/compiler/jit/torch/shape_analysis.cpp
@@ -20,6 +20,7 @@
 #include <torch/csrc/jit/ir/constants.h>
 #include <torch/csrc/jit/ir/ir.h>
 #include <torch/csrc/jit/ir/ir_views.h>
+#include <torch/csrc/jit/jit_log.h>
 #include <torch/csrc/jit/runtime/exception_message.h>
 #include <torch/csrc/jit/runtime/operator.h>
 
@@ -293,7 +294,8 @@ class ShapePropagator : public PropertyPropBase {
     static std::unordered_set<Symbol> resize_ops{
         aten::resize_,
         aten::resize_as_,
-        aten::copy_,
+        // inplaced copy would never changing the shape
+        // aten::copy_,
         aten::set_,
         aten::unsqueeze_,
         aten::t_,
@@ -751,7 +753,6 @@ class ShapePropagator : public PropertyPropBase {
     if (node->hasSideEffects()) {
       return;
     }
-
     if (node->matches("aten::cat(Tensor[] tensors, int dim) -> Tensor") ||
         node->kind() == prim::FusedConcat) {
       return PropagateCatShape(node);
@@ -856,7 +857,9 @@ class ShapePropagator : public PropertyPropBase {
             "aten::neg(Tensor self) -> Tensor",
             "aten::t(Tensor self) -> Tensor",
             "aten::sigmoid(Tensor self) -> Tensor",
+#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION > 6
             "aten::logit(Tensor self, float? eps=None) -> Tensor",
+#endif
             "aten::tanh(Tensor self) -> Tensor",
             "aten::relu(Tensor self) -> Tensor",
             "aten::asin(Tensor self) -> Tensor",
@@ -901,7 +904,11 @@ class ShapePropagator : public PropertyPropBase {
             "aten::normal(float mean, Tensor std, *, Generator? generator) -> Tensor",
             "aten::normal(Tensor mean, float std, *, Generator? generator) -> Tensor",
             "aten::permute(Tensor self, int[] dims) -> Tensor",
+#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION >= 12
             "aten::pin_memory(Tensor(a) self, Device? device=None) -> Tensor(a)",
+            "aten::gelu(Tensor self, *, str approximate='none') -> Tensor",
+            "aten::gelu_backward(Tensor grad_output, Tensor self, *, str approximate='none') -> Tensor",
+#endif
             "aten::pinverse(Tensor self, float rcond) -> Tensor",
             "aten::reciprocal(Tensor self) -> Tensor",
             "aten::relu(Tensor self) -> Tensor",
@@ -909,7 +916,6 @@ class ShapePropagator : public PropertyPropBase {
             "aten::rrelu(Tensor self, Scalar lower, Scalar upper, bool training, Generator? generator) -> Tensor",
             "aten::rsqrt(Tensor self) -> Tensor",
             "aten::selu(Tensor self) -> Tensor",
-            "aten::gelu(Tensor self, *, str approximate='none') -> Tensor",
             "aten::sigmoid(Tensor self) -> Tensor",
             "aten::sign(Tensor self) -> Tensor",
             "aten::sin(Tensor self) -> Tensor",
@@ -926,9 +932,13 @@ class ShapePropagator : public PropertyPropBase {
             "aten::trunc(Tensor self) -> Tensor",
             "aten::rot90(Tensor self, int k, int[] dims) -> Tensor",
             "aten::narrow(Tensor self, int dim, int start, int length) -> Tensor",
+#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION > 7
             "aten::slice(Tensor self, int dim, int? start=None, int? end=None, int step=1) -> Tensor",
+#else
+            "aten::slice(Tensor self, int dim, int start, int end, int step) -> Tensor",
+#endif
             "aten::alias(Tensor self) -> Tensor",
-            "aten::zero(Tensor self) -> Tensor",
+            "aten::zero_(Tensor self) -> Tensor",
         },
         [](Node* node) -> type_vec_t {
           auto input_type = node->input(0)->type()->cast<TensorType>();
@@ -1209,7 +1219,9 @@ class ShapePropagator : public PropertyPropBase {
             "aten::max(Tensor self) -> Tensor",
             "aten::min(Tensor self) -> Tensor",
             "aten::median(Tensor self) -> Tensor",
+#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION > 7
             "aten::nanmedian(Tensor self) -> Tensor",
+#endif
             "aten::norm(Tensor self, Scalar p) -> Tensor",
             "aten::std(Tensor self, bool unbiased) -> Tensor",
             "aten::trace(Tensor self) -> Tensor",
@@ -1388,7 +1400,9 @@ class ShapePropagator : public PropertyPropBase {
             "aten::max(Tensor self, int dim, bool keepdim) -> (Tensor, Tensor)",
             "aten::min(Tensor self, int dim, bool keepdim) -> (Tensor, Tensor)",
             "aten::median(Tensor self, int dim, bool keepdim) -> (Tensor, Tensor)",
+#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION > 7
             "aten::nanmedian(Tensor self, int dim, bool keepdim) -> (Tensor, Tensor)",
+#endif
             "aten::mode(Tensor self, int dim, bool keepdim) -> (Tensor, Tensor)",
         },
         [](Node* node) -> type_vec_t {
@@ -1466,6 +1480,29 @@ class ShapePropagator : public PropertyPropBase {
           return {};
         }};
 
+    // Requirements:
+    //   dims           : preserved
+    //   scalar type    : half to float if specified, otherwise preserved
+    //   device         : preserved
+    //   tensor inputs  : 1
+#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION >= 12
+    static const register_formula_for register__softmax_backward{
+        {"aten::_softmax_backward_data(Tensor grad_output, Tensor output, int dim, ScalarType input_dtype) -> Tensor"},
+        [](Node* node) -> type_vec_t {
+          if (auto type = node->input(0)->type()->cast<TensorType>()) {
+            auto scalar_type = type->scalarType();
+            auto input_dtype = node->get(attr::input_dtype)->toScalarType();
+            auto half_to_float = scalar_type != input_dtype;
+            if (half_to_float && scalar_type == at::kFloat &&
+                input_dtype == at::kHalf) {
+              return {type->withScalarType(at::kHalf)};
+            }
+            return {type};
+          }
+          return {};
+        }};
+#endif
+
     // Requirements:
     //   dims           : preserved
     //   scalar type    : dtype if specified, otherwise preserved
@@ -1543,15 +1580,19 @@ class ShapePropagator : public PropertyPropBase {
     // Additionally:
     //   - has ScalarType dtype, Layeout layout and Device device arguments
     static const register_formula_for like_factories_with_options{
-        {"aten::empty_like(Tensor self, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
-         "aten::full_like(Tensor self, Scalar fill_value, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
-         "aten::ones_like(Tensor self, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
-         "aten::rand_like(Tensor self, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
-         "aten::randint_like(Tensor self, int high, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
-         "aten::randint_like(Tensor self, int low, int high, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
-         "aten::randn_like(Tensor self, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
-         "aten::zeros_like(Tensor self, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
-         "aten::_to_copy(Tensor self, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor"},
+        {
+            "aten::empty_like(Tensor self, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
+            "aten::full_like(Tensor self, Scalar fill_value, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
+            "aten::ones_like(Tensor self, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
+            "aten::rand_like(Tensor self, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
+            "aten::randint_like(Tensor self, int high, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
+            "aten::randint_like(Tensor self, int low, int high, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
+            "aten::randn_like(Tensor self, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
+            "aten::zeros_like(Tensor self, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor",
+#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION >= 12
+            "aten::_to_copy(Tensor self, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None, bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor",
+#endif
+        },
         [](Node* node) -> type_vec_t {
           if (auto type =
                   node->namedInput(attr::self)->type()->cast<TensorType>()) {
@@ -1685,8 +1726,12 @@ class ShapePropagator : public PropertyPropBase {
         return true;
       }
     } else if (
+#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION > 7
+        node->matches(
+            "aten::native_layer_norm(Tensor input, int[] normalized_shape, Tensor? weight, Tensor? bias, float eps) -> (Tensor, Tensor, Tensor)") ||
+#endif
         node->matches(
-            "aten::native_layer_norm(Tensor input, int[] normalized_shape, Tensor? weight, Tensor? bias, float eps) -> (Tensor, Tensor, Tensor)")) {
+            "aten::native_layer_norm_backward(Tensor grad_out, Tensor input, int[] normalized_shape, Tensor mean, Tensor rstd, Tensor? weight, Tensor? bias, bool[3] output_mask) -> (Tensor, Tensor, Tensor)")) {
       if (auto type = input_type(0)) {
         node->outputs()[0]->setType(type);
         node->outputs()[1]->setType(type);
@@ -1796,6 +1841,19 @@ class ShapePropagator : public PropertyPropBase {
         node->output()->setType(weight_type->withDim(*indices_type->dim() + 1));
         return true;
       }
+    } else if (
+        node->matches(
+            "aten::embedding_dense_backward(Tensor grad_output, Tensor indices, int num_weights, int padding_idx, bool scale_grad_by_freq) -> Tensor")) {
+      auto grad_output_type = input_type(0);
+      auto maybe_num_weights = node->get<int>(attr::num_weights);
+      if (grad_output_type && maybe_num_weights &&
+          grad_output_type->sizes().size()) {
+        // TODO(yancey.yx): sizes = {num_weights, grad_output.size(-1)}
+        // ref:
+        // aten/src/ATen/native/Embedding.cpp::embedding_dense_backward_cpu.
+        node->output()->setType(grad_output_type->withDim(2));
+        return true;
+      }
     } else if (
         node->matches(
             "aten::bilinear(Tensor input1, Tensor input2, Tensor weight, Tensor? bias) -> Tensor")) {
diff --git a/pytorch_blade/src/compiler/jit/torch/shape_analysis.h b/pytorch_blade/src/compiler/jit/torch/shape_analysis.h
index fd3f87c6c70..dd8604dc722 100644
--- a/pytorch_blade/src/compiler/jit/torch/shape_analysis.h
+++ b/pytorch_blade/src/compiler/jit/torch/shape_analysis.h
@@ -28,7 +28,12 @@ struct Value;
 
 namespace torch {
 namespace blade {
-using namespace torch::jit;
+
+using torch::jit::Block;
+using torch::jit::Graph;
+using torch::jit::Node;
+using torch::jit::Value;
+
 struct propagation_error : std::exception {};
 
 class PropertyPropBase {
@@ -38,21 +43,23 @@ class PropertyPropBase {
       : graph_(std::move(graph)) {}
   virtual ~PropertyPropBase() = default;
 
-  void propagateBlock(Block* block, bool insert_expands = true);
+  void propagateBlock(torch::jit::Block* block, bool insert_expands = true);
   // insert_expands is used for shape inference
 
-  void processIf(Node* node);
-  void processLoop(Node* node);
+  void processIf(torch::jit::Node* node);
+  void processLoop(torch::jit::Node* node);
 
  protected:
-  virtual void propagateNode(Node* node, bool insert_expands = true) = 0;
+  virtual void propagateNode(
+      torch::jit::Node* node,
+      bool insert_expands = true) = 0;
   void setUnshapedType(Value* o);
-  void setUnshapedType(Node* node);
-  std::shared_ptr<Graph> graph_;
+  void setUnshapedType(torch::jit::Node* node);
+  std::shared_ptr<torch::jit::Graph> graph_;
 };
 
-void EraseShapeInformation(const std::shared_ptr<Graph>& graph);
-void PropagateInputShapes(const std::shared_ptr<Graph>& graph);
+void EraseShapeInformation(const std::shared_ptr<torch::jit::Graph>& graph);
+void PropagateInputShapes(const std::shared_ptr<torch::jit::Graph>& graph);
 
 bool mergeTypes(
     c10::ArrayRef<Value*> lhs,
diff --git a/pytorch_blade/src/ltc/disc_compiler/disc_compiler.cpp b/pytorch_blade/src/ltc/disc_compiler/disc_compiler.cpp
index cb0b5bba35c..f0b432adb3c 100644
--- a/pytorch_blade/src/ltc/disc_compiler/disc_compiler.cpp
+++ b/pytorch_blade/src/ltc/disc_compiler/disc_compiler.cpp
@@ -21,6 +21,7 @@
 #include <torch/csrc/lazy/ts_backend/ts_lowering_context.h>
 #include <torch/script.h>
 #include "common_utils/utils.h"
+#include "compiler/jit/torch/shape_analysis.h"
 #include "ltc/disc_compiler/passes/disc_fuser.h"
 #include "ltc/disc_compiler/passes/register_disc_class.h"
 
@@ -91,11 +92,12 @@ ExecutablePtr CompileToDiscExecutable(
     auto disc_inputs = std::vector<c10::IValue>{};
     return std::make_shared<Executable>(graph, disc_inputs);
   }
+  GRAPH_DEBUG("before PropagateInputShapes\n", *graph);
   EnhancementInputShape(graph, arguments);
   // Inference shape
-  torch::jit::PropagateInputShapes(graph);
+  torch::blade::PropagateInputShapes(graph);
   // cluster disc compitable nodes into a sub-graph
-  GRAPH_DEBUG("before DiscFusion\n ", *graph);
+  GRAPH_DEBUG("after PropagateInputShapes, before DiscFusion\n ", *graph);
   DiscFusion(graph);
   GRAPH_DEBUG("after DiscFusion, before EliminateDeadCode\n", *graph)
   torch::jit::EliminateDeadCode(graph);


Test patch:
diff --git a/pytorch_blade/src/compiler/jit/torch/shape_analysis_test.cpp b/pytorch_blade/src/compiler/jit/torch/shape_analysis_test.cpp
new file mode 100644
index 00000000000..46f0aa23fd0
--- /dev/null
+++ b/pytorch_blade/src/compiler/jit/torch/shape_analysis_test.cpp
@@ -0,0 +1,55 @@
+// Copyright 2022 The BladeDISC Authors. All rights reserved.
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+#include <torch/csrc/jit/ir/irparser.h>
+#include <torch/csrc/jit/testing/file_check.h>
+#include <torch/script.h>
+#include "compiler/jit/torch/shape_analysis.h"
+
+namespace torch {
+namespace blade {
+
+TEST(PropagateInputShapesTest, SimpleUnary) {
+#if PYTORCH_MAJOR_VERSOIN == 1 && PYTORCH_MINOR_VERSION >= 8
+  const std::string graph_str = R"IR(
+graph(%p1 : Float(*, *, *, device=cuda:0)):
+  %1 : Tensor = aten::relu(%p1)
+  return (%1)
+)IR";
+#elif PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION > 6
+  const std::string graph_str = R"IR(
+graph(%p1 : Float(*, *, *, device=cuda:0)):
+  %1 : Tensor = aten::relu(%p1)
+  return (%1)
+)IR";
+#else
+  const std::string graph_str = R"IR(
+graph(%p1 : Float(*, *, *)):
+  %1 : Tensor = aten::relu(%p1)
+  return (%1)
+)IR";
+#endif
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph_str, g.get());
+  torch::blade::PropagateInputShapes(g);
+  std::cout << g->toString() << std::endl;
+  torch::jit::testing::FileCheck()
+      .check(
+#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION <= 6
+          "Float(*, *, *) = aten::relu(%p1)")
+#else
+          "Float(*, *, *, device=cuda:0) = aten::relu(%p1)")
+#endif
+      ->run(*g);
+}
+} //  namespace blade
+} //  namespace torch
\ No newline at end of file