[tensorexpr] Fix the buf size of discontiguous tensors (pytorch#69657)

Summary: Pull Request resolved: pytorch#69657

Test Plan: Imported from OSS

Reviewed By: ZolotukhinM

Differential Revision: D32974473

Pulled By: huiguoo

fbshipit-source-id: 52dcd13d0ad7f7e4f1beb69dcaabc8ceb386ffca

Author: huiguoo

torch/csrc/jit/tensorexpr/kernel.cpp

@@ -937,9 +937,19 @@ Tensor TensorExprKernel::bindInput(const torch::jit::Value* input) {
         bufferArgs_.emplace_back(inBuffer);
         break;
       }
+      ExprHandle flat_size = 1;
+      for (size_t i = 0; i < *tt->sizes().size(); i++) {
+        auto size = *tt->sizes()[i];
+        if (size == 0) {
+          flat_size = 0;
+          break;
+        }
+        flat_size = flat_size + (size - 1) * *tt->strides()[i];
+      }
+      flat_size = IRSimplifier::simplify(flat_size);
       BufHandle inBuffer(
           "t" + input_name_map_[input],
-          {0},
+          {flat_size},
           ToDtype(static_cast<ScalarType>(*tt->scalarType())));
       std::vector<DimArg> inputTensorDims;
       for (size_t i = 0; i < *tt->sizes().size(); i++) {