expose stride from disc ral to torch tensor

pytorch_blade/pytorch_blade/compiler/mlir/runtime/ral_context.cpp

@@ -221,14 +221,16 @@ at::List<at::Tensor> RalContext::CreateAndBindingOutputs(
       out_tensor = torch::from_blob(
           const_cast<void*>(out_buf->data()),
           out_buf->shape(),
+          out_buf->strides(),
           deleter,
           option);
       out_buf->release();
     } else {
-      out_tensor =
-          torch::from_blob(
-              const_cast<void*>(out_buf->data()), out_buf->shape(), option)
-              .clone();
+      out_tensor = torch::from_blob(
+          const_cast<void*>(out_buf->data()),
+          out_buf->shape(),
+          out_buf->strides(),
+          option);
     }
     outputs.push_back(out_tensor);
   }

tao_compiler/mlir/xla/ral/context/base/base_context.cc

@@ -78,8 +78,8 @@ void BaseExecutionContext::bindOutput(
     return;
   }
 
-  output->reset(
-      new BaseOutputBufferWrapper(it->second.buffer, it->second.shape));
+  output->reset(new BaseOutputBufferWrapper(it->second.buffer, it->second.shape,
+                                            it->second.strides));
   if (!output_ptr_set.insert(it->second.buffer).second) {
     // This buffer is used as output before, thus is already set a deleter.
     return;
@@ -192,6 +192,9 @@ void ral_base_cuda_send_output(ExecutionContext* ctx, int64_t output_idx,
   for (int i = 0; i < N; ++i) {
     tensor.shape.push_back(memref.sizes[i]);
   }
+  for (int i = 0; i < N; ++i) {
+    tensor.strides.push_back(memref.strides[i]);
+  }
   for (int i = 0; i < N; ++i) {
     TAO_VLOG(1) << "tensor dim size = " << tensor.shape[i];
   }

tao_compiler/mlir/xla/ral/context/base/base_context.h

@@ -43,6 +43,9 @@ class BaseOutputBufferWrapper : public OutputBufferWrapper {
  public:
   BaseOutputBufferWrapper(buffer_t data, buffer_shape_t shape)
       : data_(data), shape_(shape) {}
+  BaseOutputBufferWrapper(buffer_t data, buffer_shape_t shape,
+                          buffer_shape_t strides)
+      : data_(data), shape_(shape), strides_(strides) {}
   ~BaseOutputBufferWrapper() {
     if (deleter_) deleter_(data_);
   }
@@ -51,6 +54,7 @@ class BaseOutputBufferWrapper : public OutputBufferWrapper {
 
   const_buffer_t data() override { return data_; }
   const buffer_shape_t& shape() override { return shape_; }
+  const buffer_shape_t& strides() override { return strides_; }
   void set_deleter(Deleter deleter) { deleter_ = deleter; }
 
   // Returns true if this wrapper is the exclusive owner
@@ -69,6 +73,7 @@ class BaseOutputBufferWrapper : public OutputBufferWrapper {
  private:
   buffer_t data_;
   buffer_shape_t shape_;
+  buffer_shape_t strides_;
   Deleter deleter_;
   bool owned_ = false;
 };
@@ -92,6 +97,7 @@ class InternalAllocator : public Allocator {
 struct Tensor {
   void* buffer;
   std::vector<int64_t> shape;
+  std::vector<int64_t> strides;
 };
 
 struct BaseExecutionContext : public tao::ral::ExecutionContext {

tao_compiler/mlir/xla/ral/context/base/cpu/cpu_context_impl.cc

@@ -100,7 +100,6 @@ void BaseCpuExecutionContext::setOutputDeleter(OutputBufferWrapper& output) {
     // This buffer is a pesistent buffer, thus no need to set a deleter.
     return;
   }
-
   auto hid = host_ptr_map.find(buffer);
   if (hid != host_ptr_map.end()) {
     if (--hid->second == 0) {

tao_compiler/mlir/xla/ral/context/context_util.h

@@ -66,6 +66,28 @@ tao::ral::MemRefType<T, N> assignMemRef(void* ptr, const ShapeTy& shape) {
   return memref;
 }
 
+template <typename T, int N, typename ShapeTy>
+tao::ral::MemRefType<T, N> assignMemRef(void* ptr, const ShapeTy& shape,
+                                        const ShapeTy& strides) {
+  tao::ral::MemRefType<T, N> memref;
+  memref.basePtr = reinterpret_cast<T*>(ptr);
+  memref.data = reinterpret_cast<T*>(ptr);
+  memref.offset = 0;
+  for (int i = 0; i < N; ++i) {
+    memref.sizes[i] = shape[i];
+  }
+
+  for (int i = 0; i < N; ++i) {
+    memref.strides[i] = strides[i];
+  }
+
+  if (TAO_VLOG_IS_ON(1)) {
+    print_memref(memref, "assigned");
+  }
+
+  return memref;
+}
+
 template <typename T>
 tao::ral::MemRefType<T, 0> assignMemRef_0d(void* ptr) {
   tao::ral::MemRefType<T, 0> memref;

tao_compiler/mlir/xla/ral/ral_context.h

@@ -140,6 +140,7 @@ class OutputBufferWrapper {
   virtual ~OutputBufferWrapper() {}
   virtual const_buffer_t data() = 0;
   virtual const buffer_shape_t& shape() = 0;
+  virtual const buffer_shape_t& strides() = 0;
   // Returns true if this wrapper is the exclusive owner
   virtual bool owned() const = 0;
   // mark that this wrapper exclusively owns the underlying buffer.