Repository: alibaba/BladeDISC (language: python)

Issue:
invalid output number of `torch.eq(t1, t2).sum()`
From the MNIST example of TorchDisc, the program outputs an invalid test acc number:

``` text
Test set: Average loss: 0.3399, Accuracy: 281447840135434/10000 (2814478401354%)
``` 

A tiny program to reproduce this bug:

```python
import torch
import _torch_disc as disc
disc._ltc_init_disc_backend()

def func():
    t1 = torch.tensor([[1, 2], [3, 4]])
    t2 = torch.tensor([[1, 2], [5, 6]])
    t1, t2 = t1.to('lazy'), t2.to('lazy')
    t3 = torch.eq(t1, t2)
    return t3.sum()

print(func())
```

outputs:

``` text
tensor(3328288142057042689, device='lazy:0')
```


Discussion / hints:

Fix patch:
diff --git a/docker/dev/Dockerfile.torch-disc b/docker/dev/Dockerfile.torch-disc
index b8790798368..843d1d4bdc9 100644
--- a/docker/dev/Dockerfile.torch-disc
+++ b/docker/dev/Dockerfile.torch-disc
@@ -8,5 +8,8 @@ RUN apt-get update -y && \
     rm /usr/bin/python && \
     ln -s /usr/bin/python3.8 /usr/bin/python && \
     python -m pip install --upgrade pip && \
-    python -m pip install cpython pyyaml typing_extensions virtualenv numpy
-RUN python -m pip install https://bladedisc-ci.oss-cn-hongkong.aliyuncs.com/download/torch-ltc/torch-1.12.0a0%2Bgitd9896b8-cp38-cp38-linux_x86_64.whl
\ No newline at end of file
+    python -m pip install cpython pyyaml typing_extensions virtualenv numpy torchvision
+# the torch whlee build from https://github.com/pytorch/pytorch/tree/lazy_tensor_staging, here we should
+# uninstall pytorch comes from torchvision firstly.
+# TODO(Yancey): build torchvision from the same git sha of pytorch
+RUN python -m pip uninstall pytorch && python -m pip install https://bladedisc-ci.oss-cn-hongkong.aliyuncs.com/download/torch-ltc/torch-1.12.0a0%2Bgitd9896b8-cp38-cp38-linux_x86_64.whl
\ No newline at end of file
diff --git a/torch_disc/mnist.py b/torch_disc/mnist.py
index 79bc8bd4524..f041e5e2646 100644
--- a/torch_disc/mnist.py
+++ b/torch_disc/mnist.py
@@ -72,7 +72,7 @@ def test(model, device, test_loader):
             output = model(data)
             test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
             pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
-            correct += pred.eq(target.view_as(pred)).sum().item()
+            correct += pred.eq(target.view_as(pred)).sum(dtype=torch.int).item()
             disc._step_marker()
 
     test_loss /= len(test_loader.dataset)
@@ -147,4 +147,4 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
\ No newline at end of file
+    main()
diff --git a/torch_disc/torch_disc/csrc/disc_compiler/disc_compiler.cpp b/torch_disc/torch_disc/csrc/disc_compiler/disc_compiler.cpp
index 493cdafd378..6ae074211b8 100644
--- a/torch_disc/torch_disc/csrc/disc_compiler/disc_compiler.cpp
+++ b/torch_disc/torch_disc/csrc/disc_compiler/disc_compiler.cpp
@@ -79,7 +79,6 @@ ExecutablePtr CompileToDiscExecutable(
   EnhancementInputShape(graph, arguments);
   // Inference shape
   torch::jit::PropagateInputShapes(graph);
-
   // cluster disc compitable nodes into a sub-graph
   ClusterDiscNodes(graph);
   torch::jit::EliminateDeadCode(graph);
diff --git a/torch_disc/torch_disc/csrc/disc_compiler/passes/cluster.cpp b/torch_disc/torch_disc/csrc/disc_compiler/passes/cluster.cpp
index a3aac562380..a513a1c3f3e 100644
--- a/torch_disc/torch_disc/csrc/disc_compiler/passes/cluster.cpp
+++ b/torch_disc/torch_disc/csrc/disc_compiler/passes/cluster.cpp
@@ -24,40 +24,63 @@ using namespace ::torch::jit;
 // conversion module into a group. We should re-implement this.
 std::vector<Node*> FakeCluster(const std::shared_ptr<Graph>& graph) {
   std::vector<Node*> nodes;
-  for (auto node : graph->nodes()) {
+  auto is_disc_compilable = [](torch::jit::Node* node) {
     if (torch::blade::IsMlirMhloSupported(*node) &&
         node->kind() != prim::Constant) {
-      if (node->kind() == aten::addmm) continue;
-      nodes.push_back(node);
+      for (auto& input : node->inputs()) {
+        // input should be Tensor or Scalar with explict type
+        auto typ = input->type();
+        if (!typ->cast<c10::TensorType>() &&
+            c10::tryScalarTypeFromJitType(*typ) == c10::nullopt)
+          return false;
+      }
+      return true;
     }
-  }
+    return false;
+  };
+
+  std::copy_if(graph->nodes().begin(), graph->nodes().end(),
+               std::back_inserter(nodes), is_disc_compilable);
   return nodes;
 }
 
+c10::TypePtr getScalarTypePtr(at::ScalarType& typ) {
+  if (c10::isFloatingType(typ)) {
+    return c10::FloatType::get();
+  } else if (c10::isIntegralType(typ)) {
+    return c10::IntType::get();
+  } else if (typ == c10::ScalarType::Bool) {
+    return c10::BoolType::get();
+  }
+  TORCH_CHECK(false, "unsupported scalar type: ", typ);
+}
+
 // Give:
 //  with prim::FusionGroup(
-//      %1: Scalar,
-//      %2: Scalar):
-//    %3 Tensor = aten::add(%1, %2)
-//  return %3
+//      %p0: Tensor,
+//      %p1: Tensor,
+//      %p2: Scalar):
+//    %1 Tensor = aten::add(%p0, %p1, %p2)
+//  return %1
 //
 // Execute: CastToTensorInputs(sub_graph)
 //
 // After:
 //  with prim::FusionGroup(
-//      %1.1: Tensor,
-//      %2.1: Tensor):
-//    %4 = aten::item(%1.1, 1)
-//    %5 = aten::item(%2.1, 1)
-//    %3 Tensor = aten::add(%4, %5)
-//    return %3
-void CastBoundaryScalarToTensor(Graph* disc_graph, size_t i) {
+//      %p0.1: Tensor,
+//      %p1.1: Tensor,
+//      %p2.1: Tensor):
+//    %1 int = aten::item(%p2.1)
+//    %2 Tensor = aten::add(%p0.1, %p1.1, %1)
+//    return %2
+void CastBoundaryScalarToTensor(Graph* disc_graph, size_t i,
+                                at::ScalarType& typ) {
   auto new_input = disc_graph->insertInput(
       i, c10::string(disc_graph->inputs()[i]->debugName() + ".1"));
+  new_input->setType(TensorType::create(typ, c10::nullopt, 0, false));
   auto orig_input = disc_graph->inputs()[i + 1];
   auto item_node = disc_graph->create(aten::item, {new_input});
-  // TODO(Yancey1989): supports more types
-  item_node->output()->setType(c10::IntType::get());
+  item_node->output()->setType(getScalarTypePtr(typ));
   disc_graph->appendNode(item_node);
   orig_input->replaceAllUsesWith(item_node->output());
   item_node->moveBefore(item_node->output()->uses()[0].user);
@@ -77,7 +100,8 @@ void CastGraphInputsToTensor(const std::shared_ptr<Graph>& graph,
       sub_graph->replaceInput(i, cast_tensor->output());
 
       // TODO(Yancey1989): cast output
-      CastBoundaryScalarToTensor(disc_graph, i);
+      auto scalar_type = c10::scalarTypeFromJitType(*input->type());
+      CastBoundaryScalarToTensor(disc_graph, i, scalar_type);
     }
   }
 }
diff --git a/torch_disc/torch_disc/csrc/disc_compiler/passes/register_disc_class.cpp b/torch_disc/torch_disc/csrc/disc_compiler/passes/register_disc_class.cpp
index 3bf580a2c12..2930e30a498 100644
--- a/torch_disc/torch_disc/csrc/disc_compiler/passes/register_disc_class.cpp
+++ b/torch_disc/torch_disc/csrc/disc_compiler/passes/register_disc_class.cpp
@@ -12,6 +12,7 @@
 #include "torch_disc/csrc/disc_compiler/passes/register_disc_class.h"
 
 #include "compiler/mlir/converters/mhlo_conversion.h"
+#include "lazy_tensors/computation_client/sys_util.h"
 #include "torch_disc/csrc/disc_compiler/passes/disc_class.h"
 #include "torch_disc/csrc/disc_compiler/passes/io.h"
 
@@ -23,8 +24,10 @@ std::string DiscCMD(const std::string& mlir_fname,
                     const std::string& out_fname) {
   std::stringstream ss;
   std::string logf = mlir_fname + ".log";
-  ss << "./disc_compiler_main " << mlir_fname << " " << out_fname << " > "
-     << logf << " 2>&1 ";
+  auto binary_path = lazy_tensors::sys_util::GetEnvString(
+      "DISC_COMPILER_BINARY_PATH", "./disc_compiler_main");
+  ss << binary_path << " " << mlir_fname << " " << out_fname << " > " << logf
+     << " 2>&1 ";
   return ss.str();
 }
 


Test patch:
diff --git a/scripts/ci/test_torch_disc.sh b/scripts/ci/test_torch_disc.sh
index d2b9f084037..0028b8c9e63 100644
--- a/scripts/ci/test_torch_disc.sh
+++ b/scripts/ci/test_torch_disc.sh
@@ -26,6 +26,7 @@ python setup.py develop
 # 5. test a e2e demo
 ln -s bazel-bin/torch_disc/_torch_disc.so ./_torch_disc.so
 ln -s bazel-out/k8-opt/bin/external/org_tensorflow/tensorflow/compiler/mlir/xla/ral/libral_base_context.so ./libral_base_context.so
+python -m unittest discover torch_disc/python/tests
 python disc_demo.py
 
 deactivate
diff --git a/torch_disc/torch_disc/python/tests/test_mnist.py b/torch_disc/torch_disc/python/tests/test_mnist.py
new file mode 100644
index 00000000000..3abf6dcdf61
--- /dev/null
+++ b/torch_disc/torch_disc/python/tests/test_mnist.py
@@ -0,0 +1,129 @@
+# 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.
+
+from __future__ import print_function
+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.testing import assert_allclose
+from torchvision import datasets, transforms
+from torch.optim.lr_scheduler import StepLR
+import _torch_disc as disc
+disc._ltc_init_disc_backend()
+import unittest
+
+LOG_INTERVAL=100
+
+## Define the NN architecture
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.fc1 = nn.Linear(28*28, 50)
+        self.fc1_drop = nn.Dropout(0.2)
+        self.fc2 = nn.Linear(50, 50)
+        self.fc2_drop = nn.Dropout(0.2)
+        self.fc3 = nn.Linear(50, 10)
+
+    def forward(self, x):
+        x = x.view(-1, 28*28)
+        x = F.relu(self.fc1(x))
+        x = self.fc1_drop(x)
+        x = F.relu(self.fc2(x))
+        x = self.fc2_drop(x)
+        return F.log_softmax(self.fc3(x), dim=1)
+
+
+def train(model, device, train_loader, optimizer, epoch):
+    model.train()
+    train_loss = 0
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        disc._step_marker()
+        train_loss += loss.item() * len(data)
+        if batch_idx % LOG_INTERVAL == 0:
+            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
+                epoch, batch_idx * len(data), len(train_loader.dataset),
+                100. * batch_idx / len(train_loader), loss.item()))
+    train_loss /= len(train_loader.dataset)
+    print('Train Epoch:{} Average loss: {:.4f}'.format(epoch, train_loss))
+
+
+def test(model, device, test_loader):
+    model.eval()
+    test_loss = 0
+    correct = 0
+    with torch.no_grad():
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
+            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
+            correct += pred.eq(target.view_as(pred)).sum(dtype=torch.int).item()
+            disc._step_marker()
+
+    test_loss /= len(test_loader.dataset)
+
+    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
+        test_loss, correct, len(test_loader.dataset),
+        100. * correct / len(test_loader.dataset)))
+    return correct * 1.0 / len(test_loader.dataset)
+
+class TestMnist(unittest.TestCase):
+    def mnit(self, device):
+        torch.manual_seed(2)
+        epochs = 2
+        lr = 1.0
+        gamma = 0.7
+        train_kwargs = {'batch_size': 64}
+        test_kwargs = {'batch_size': 1000}
+
+        transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.1307,), (0.3081,))
+        ])
+        dataset1 = datasets.MNIST('~/.cache/data', train=True, download=True,
+                       transform=transform)
+        dataset2 = datasets.MNIST('~/.cache/data', train=False,
+                       transform=transform)
+        train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs)
+        test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
+
+        model = Net().to(device)
+        optimizer = optim.Adadelta(model.parameters(), lr=lr)
+
+        scheduler = StepLR(optimizer, step_size=1, gamma=gamma)
+        test_acc = None
+        for epoch in range(1, epochs + 1):
+            train(model, device, train_loader, optimizer, epoch)
+            test_acc = test(model, device, test_loader)
+            scheduler.step()
+
+        return test_acc
+
+    def test_acc(self):
+        lazy_device = torch.device('lazy')
+        device = torch.device('cpu')
+        actual_acc = self.mnit(lazy_device)
+        # test on CUDA device
+        expect_acc = self.mnit(device)
+        print(expect_acc)
+        print(actual_acc)
+        assert_allclose(expect_acc, actual_acc, rtol=1e-3, atol=0)
+
+if __name__ == '__main__':
+    unittest.main()