Repository: alibaba/BladeDISC (language: python)

Issue:
Add support for diffusion models
We are currently focus on support diffusion models with comparable performance.
dev branch is https://github.com/alibaba/BladeDISC/tree/stable_diffusion_for_unibuild.
Include work items: 
- [x] add support for rank 0 arith.const https://github.com/alibaba/BladeDISC/pull/868
- [x] add shape inference for some op https://github.com/alibaba/BladeDISC/pull/871
- [x] performance optimize for attention kernel https://github.com/alibaba/BladeDISC/pull/869
- [x]   add support for group_norm op https://github.com/alibaba/BladeDISC/pull/874
   - [x] lowering group_norm to disc codegen 
   - [x] with bladnn tuning support for better performance


Discussion / hints:

Fix patch:
diff --git a/pytorch_blade/pytorch_blade/compiler/jit/torch/op_registry.cpp b/pytorch_blade/pytorch_blade/compiler/jit/torch/op_registry.cpp
index 2fd8d4c8325..06fe1619b8b 100644
--- a/pytorch_blade/pytorch_blade/compiler/jit/torch/op_registry.cpp
+++ b/pytorch_blade/pytorch_blade/compiler/jit/torch/op_registry.cpp
@@ -46,6 +46,7 @@ std::shared_ptr<SchemaSet> nn_ops_first_input_preserving() {
 #if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION > 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) -> 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",
       "aten::adaptive_avg_pool1d(Tensor self, int[] output_size) -> Tensor",
       "aten::adaptive_avg_pool2d(Tensor self, int[] output_size) -> Tensor",
diff --git a/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp b/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp
index cc137a4eda1..694f931560c 100644
--- a/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp
+++ b/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp
@@ -870,6 +870,7 @@ class ShapePropagator : public PropertyPropBase {
             "aten::hardtanh(Tensor self, Scalar min_val, Scalar max_val) -> Tensor",
             "aten::glu(Tensor self, int dim) -> Tensor",
             "aten::inverse(Tensor self) -> Tensor",
+            "aten::group_norm(Tensor input, int num_groups, Tensor? weight, Tensor? bias, float eps, bool cudnn_enabled) -> Tensor",
             "aten::leaky_relu(Tensor self, Scalar negative_slope) -> Tensor",
             "aten::leaky_relu_(Tensor self, Scalar negative_slope) -> Tensor",
             "aten::lgamma(Tensor self) -> Tensor",
@@ -1278,6 +1279,37 @@ class ShapePropagator : public PropertyPropBase {
           return {};
         }};
 
+    // Requirements:
+    //   device         : Device
+    //   tensor inputs  : 1
+    //   tensor outputs : 1
+    // Additionally:
+    //   - First input should be the only tensor input
+    static const register_formula_for aten_to_device{
+        {"aten::to.device(Tensor self, Device device, ScalarType dtype, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor"},
+        [](Node* node) -> type_vec_t {
+          at::optional<IValue> maybe_device_option = node->get(attr::device);
+          if (auto type = node->input(0)->type()->cast<TensorType>()) {
+            auto ret = type;
+            if (maybe_device_option && !maybe_device_option->isNone()) {
+              auto device = maybe_device_option->toDevice();
+#if PYTORCH_VERSION_GE(1, 11)
+              return {ret->withDevice(device)};
+#else
+              return {TensorType::create(
+                  ret->scalarType(),
+                  device,
+                  ret->sizes(),
+                  ret->strides(),
+                  /*requires_grad=*/c10::nullopt)};
+#endif
+            } else {
+              return {ret};
+            }
+          }
+          return {};
+        }};
+
     // Requirements:
     //   dims           : 0
     //   scalar type    : preserved
diff --git a/pytorch_blade/pytorch_blade/compiler/mlir/converters/torch_mlir_op_filter.cpp b/pytorch_blade/pytorch_blade/compiler/mlir/converters/torch_mlir_op_filter.cpp
index f6941bebe71..85a3dda0f4f 100644
--- a/pytorch_blade/pytorch_blade/compiler/mlir/converters/torch_mlir_op_filter.cpp
+++ b/pytorch_blade/pytorch_blade/compiler/mlir/converters/torch_mlir_op_filter.cpp
@@ -46,19 +46,21 @@ const std::unordered_set<std::string> &GetTorchMlirWhiteList() {
       "aten::addmm",
       "aten::arange",
       "aten::unbind",
+      "aten::baddmm",
       "aten::batch_norm",
       "aten::bitwise_not",
       "aten::bmm",
       "aten::cat",
-      "aten::chunk",
       "aten::contiguous",
       "aten::_convolution",
       "aten::convolution",
       "aten::conv1d",
       "aten::conv2d",
+      "aten::cos",
       "aten::div",
       "aten::einsum",
       "aten::embedding",
+      "aten::empty",
       "aten::eq",
       "aten::gt",
       "aten::ge",
@@ -107,6 +109,7 @@ const std::unordered_set<std::string> &GetTorchMlirWhiteList() {
       "aten::select",
       "aten::sigmoid",
       "aten::silu",
+      "aten::sin",
       "aten::size",
       "aten::slice",
       "aten::softmax",
@@ -120,6 +123,7 @@ const std::unordered_set<std::string> &GetTorchMlirWhiteList() {
       "aten::tensor",
       "aten::to",
       "aten::to.dtype",
+      "aten::to.device",
       "aten::transpose",
       "aten::type_as",
       "aten::unsqueeze",
@@ -128,6 +132,7 @@ const std::unordered_set<std::string> &GetTorchMlirWhiteList() {
       "prim::Constant",
       "prim::ListConstruct",
       "prim::ListUnpack",
+      "prim::NumToTensor",
       // Torch Blade custom ops follows:
       "aten::add_inplace", // use aten namespace to work with PyTorch mutation pass
       "aten::sub_inplace", // use aten namespace to work with PyTorch mutation pass


Test patch:
diff --git a/pytorch_blade/tests/torchscript/basics.graph b/pytorch_blade/tests/torchscript/basics.graph
index bc0ed71cdb2..664b67061a6 100644
--- a/pytorch_blade/tests/torchscript/basics.graph
+++ b/pytorch_blade/tests/torchscript/basics.graph
@@ -61,3 +61,30 @@ graph(%p1 : Float(*, *, *, device=cuda:0),
 // CHECK: Float(*, *, *, device=cuda:0) = aten::tanh_backward(%p1, %p2)
   %1 : Tensor = aten::tanh_backward(%p1, %p2)
   return (%1)
+
+
+// CHECK-LABEL: graph
+graph(%p1 : Float(*, *, *, *, device=cuda:0),
+  %p2 : Float(*, *, *, *, device=cuda:0),
+  %p3 : Float(*, device=cuda:0)
+):
+  %3 : int[] = prim::Constant[value=[1, 1]]()
+  %4 : bool = prim::Constant[value=0]()
+  %5 : int[] = prim::Constant[value=[0, 0]]()
+  %6 : int = prim::Constant[value=1]()
+  // CHECK: Float(*, *, *, *, device=cuda:0) = aten::_convolution(%p1, %p2, %p3, %3, %3, %3, %4, %5, %6, %4, %4, %4, %4)
+  %7 : Tensor = aten::_convolution(%p1, %p2, %p3, %3, %3, %3, %4, %5, %6, %4, %4, %4, %4)
+  return (%7)
+
+// CHECK-LABEL: graph
+graph(%p1 : Float(*, *, *, *, device=cuda:0),
+  %p2 : Float(*, *, *, *, device=cuda:0),
+  %p3 : Float(*, device=cuda:0)
+):
+  %3 : int[] = prim::Constant[value=[1, 1]]()
+  %4 : bool = prim::Constant[value=0]()
+  %5 : int[] = prim::Constant[value=[0, 0]]()
+  %6 : int = prim::Constant[value=1]()
+  // CHECK: Float(*, *, *, *, device=cuda:0) = aten::convolution(%p1, %p2, %p3, %3, %3, %3, %4, %5, %6)
+  %7 : Tensor = aten::convolution(%p1, %p2, %p3, %3, %3, %3, %4, %5, %6)
+  return (%7)
diff --git a/pytorch_blade/tests/torchscript/since_1_10.graph b/pytorch_blade/tests/torchscript/since_1_10.graph
index 0910a1d38de..f821ddd6553 100644
--- a/pytorch_blade/tests/torchscript/since_1_10.graph
+++ b/pytorch_blade/tests/torchscript/since_1_10.graph
@@ -23,3 +23,22 @@ graph(%p1 : Half(*, *)):
   %1 : Tensor = aten::as_strided(%p1, %size, %strides, %none)
   return (%1)
 
+// CHECK-LABEL: graph
+graph(%p1 : Float(*, *, *, device=cpu)):
+  %1 : Device = prim::Constant[value="cuda:0"]()
+  %2 : int = prim::Constant[value=5]()
+  %3 : bool = prim::Constant[value=0]()
+  %4 : NoneType = prim::Constant()
+  // CHECK: Float(*, *, *, device=cuda:0) = aten::to(%p1, %1, %2, %3, %3, %4)
+  %3 : Tensor = aten::to(%p1, %1, %2, %3, %3, %4)
+  return (%3)
+
+// CHECK-LABEL: graph
+graph(%p1 : Float(20, 30, 40, device=cpu)):
+  %1 : Device = prim::Constant[value="cuda:0"]()
+  %2 : int = prim::Constant[value=5]()
+  %3 : bool = prim::Constant[value=0]()
+  %4 : NoneType = prim::Constant()
+  // CHECK: Float(20, 30, 40, device=cuda:0) = aten::to(%p1, %1, %2, %3, %3, %4)
+  %3 : Tensor = aten::to(%p1, %1, %2, %3, %3, %4)
+  return (%3)
\ No newline at end of file
diff --git a/pytorch_blade/tests/torchscript/slice_like.graph b/pytorch_blade/tests/torchscript/slice_like.graph
index ef2837db4cf..67330103fb6 100644
--- a/pytorch_blade/tests/torchscript/slice_like.graph
+++ b/pytorch_blade/tests/torchscript/slice_like.graph
@@ -19,3 +19,11 @@ graph(%p1 : Float(8, 768, 512)):
   %slice : Tensor = aten::slice(%p1, %cst_1, %cst0, %int_max, %cst1) 
   return (%slice)
 
+// CHECK-LABEL: graph
+graph(%p1 : Float(1, 512, requires_grad=0, device=cpu),
+      %p2 : int):
+  %2 : int = prim::Constant[value=0]()
+  %3 : int = prim::Constant[value=1]()
+// CHECK: Float(1, *, requires_grad=0, device=cpu) = aten::slice(%p1, %3, %2, %p2, %3)
+  %4 : Tensor = aten::slice(%p1, %3, %2, %p2, %3)
+  return (%4)