issue/21 - removed descriptor overwrite

src/models/cache_manager.hpp

@@ -21,23 +21,14 @@ inline void hash_combine(size_t &seed, T value, typename std::enable_if<std::is_
     hash_combine(seed, static_cast<size_t>(value));
 }
 
-// Specialization for float to handle potential precision issues
-inline void hash_combine(size_t &seed, float value) {
-    // Treat float bits as uint32_t for consistent hashing
-    uint32_t int_value;
-    static_assert(sizeof(value) == sizeof(int_value), "Size mismatch");
-    std::memcpy(&int_value, &value, sizeof(value));
-    hash_combine(seed, static_cast<size_t>(int_value));
-}
-
 // Helper function to compute hash for tensor descriptors
-inline size_t computeTensorDescHash(std::shared_ptr<TensorDesc> desc) {
+inline size_t computeTensorDescHash(std::shared_ptr<Tensor> tensor) {
     size_t seed = 0;
-    hash_combine(seed, desc->dtype());
-    for (auto dim : desc->shape()) {
+    hash_combine(seed, tensor->dtype());
+    for (auto dim : tensor->shape()) {
         hash_combine(seed, dim);
     }
-    for (auto stride : desc->strides()) {
+    for (auto stride : tensor->strides()) {
         hash_combine(seed, static_cast<size_t>(stride));
     }
     return seed;
@@ -185,7 +176,7 @@ public:
 
 class CacheManager {
 private:
-    const size_t DEFAULT_CACHE_CAPACITY = 100;
+    const size_t DEFAULT_CACHE_CAPACITY = 128;
 
     LRUDescriptorCache<infiniopRMSNormDescriptor_t> rms_norm_cache;
     LRUDescriptorCache<infiniopGemmDescriptor_t> gemm_cache;
@@ -267,11 +258,11 @@ public:
         random_sample_cache.put(key, desc);
     }
 
-    static size_t createDescriptorKey(std::shared_ptr<TensorDesc> desc0,
-                                      std::shared_ptr<TensorDesc> desc1,
-                                      std::shared_ptr<TensorDesc> desc2,
-                                      std::shared_ptr<TensorDesc> desc3,
-                                      std::shared_ptr<TensorDesc> desc4) {
+    static size_t createDescriptorKey(std::shared_ptr<Tensor> desc0,
+                                      std::shared_ptr<Tensor> desc1,
+                                      std::shared_ptr<Tensor> desc2,
+                                      std::shared_ptr<Tensor> desc3,
+                                      std::shared_ptr<Tensor> desc4) {
         size_t seed = 0;
         if (desc0) {
             hash_combine(seed, computeTensorDescHash(desc0));

src/models/inference_context.cpp

@@ -16,7 +16,7 @@ void InferenceContext::rmsnorm(std::shared_ptr<Tensor> y,
                                std::shared_ptr<Tensor> x,
                                std::shared_ptr<Tensor> w,
                                float epsilon) {
-    size_t key = CacheManager::createDescriptorKey(y->tdesc(), x->tdesc(), w->tdesc(), nullptr, nullptr);
+    size_t key = CacheManager::createDescriptorKey(y, x, w, nullptr, nullptr);
 
     infiniopRMSNormDescriptor_t desc;
     if (!cache_manager->getRMSNormDescriptor(key, desc)) {
@@ -35,23 +35,16 @@ void InferenceContext::rmsnorm(std::shared_ptr<Tensor> y,
         y->data(), x->data(), w->data(), stream));
 }
 
-void InferenceContext::gemm(std::shared_ptr<Tensor> c, std::shared_ptr<TensorDesc> c_desc_overwrite,
-                            std::shared_ptr<Tensor> a, std::shared_ptr<TensorDesc> a_desc_overwrite,
-                            std::shared_ptr<Tensor> b, std::shared_ptr<TensorDesc> b_desc_overwrite,
+void InferenceContext::gemm(std::shared_ptr<Tensor> c,
+                            std::shared_ptr<Tensor> a,
+                            std::shared_ptr<Tensor> b,
                             float alpha, float beta) {
-    size_t key = CacheManager::createDescriptorKey(
-        c_desc_overwrite ? c_desc_overwrite : c->tdesc(),
-        a_desc_overwrite ? a_desc_overwrite : a->tdesc(),
-        b_desc_overwrite ? b_desc_overwrite : b->tdesc(),
-        nullptr, nullptr);
+    size_t key = CacheManager::createDescriptorKey(c, a, b,
+                                                   nullptr, nullptr);
 
     infiniopGemmDescriptor_t desc;
     if (!cache_manager->getGemmDescriptor(key, desc)) {
-        RUN_INFINI(infiniopCreateGemmDescriptor(
-            rsrc->handle, &desc,
-            c_desc_overwrite ? c_desc_overwrite->desc() : c->desc(),
-            a_desc_overwrite ? a_desc_overwrite->desc() : a->desc(),
-            b_desc_overwrite ? b_desc_overwrite->desc() : b->desc()));
+        RUN_INFINI(infiniopCreateGemmDescriptor(rsrc->handle, &desc, c->desc(), a->desc(), b->desc()));
         cache_manager->putGemmDescriptor(key, desc);
     }
 
@@ -65,19 +58,13 @@ void InferenceContext::gemm(std::shared_ptr<Tensor> c, std::shared_ptr<TensorDes
         c->data(), a->data(), b->data(), alpha, beta, stream));
 }
 
-void InferenceContext::rearrange(std::shared_ptr<Tensor> dst, std::shared_ptr<TensorDesc> dst_desc_overwrite,
-                                 std::shared_ptr<Tensor> src, std::shared_ptr<TensorDesc> src_desc_overwrite) {
-    size_t key = CacheManager::createDescriptorKey(
-        dst_desc_overwrite ? dst_desc_overwrite : dst->tdesc(),
-        src_desc_overwrite ? src_desc_overwrite : src->tdesc(),
-        nullptr, nullptr, nullptr);
+void InferenceContext::rearrange(std::shared_ptr<Tensor> dst,
+                                 std::shared_ptr<Tensor> src) {
+    size_t key = CacheManager::createDescriptorKey(dst, src, nullptr, nullptr, nullptr);
 
     infiniopRearrangeDescriptor_t desc;
     if (!cache_manager->getRearrangeDescriptor(key, desc)) {
-        RUN_INFINI(infiniopCreateRearrangeDescriptor(
-            rsrc->handle, &desc,
-            dst_desc_overwrite ? dst_desc_overwrite->desc() : dst->desc(),
-            src_desc_overwrite ? src_desc_overwrite->desc() : src->desc()));
+        RUN_INFINI(infiniopCreateRearrangeDescriptor(rsrc->handle, &desc, dst->desc(), src->desc()));
         cache_manager->putRearrangeDescriptor(key, desc);
     }
 
@@ -93,7 +80,7 @@ void InferenceContext::rope(std::shared_ptr<Tensor> q,
                             std::shared_ptr<Tensor> pos,
                             std::shared_ptr<Tensor> sin,
                             std::shared_ptr<Tensor> cos) {
-    size_t key = CacheManager::createDescriptorKey(q->tdesc(), k->tdesc(), pos->tdesc(), sin->tdesc(), cos->tdesc());
+    size_t key = CacheManager::createDescriptorKey(q, k, pos, sin, cos);
 
     infiniopRoPEDescriptor_t desc;
     if (!cache_manager->getRoPEDescriptor(key, desc)) {
@@ -114,19 +101,14 @@ void InferenceContext::rope(std::shared_ptr<Tensor> q,
         sin->data(), cos->data(), stream));
 }
 
-void InferenceContext::causalSoftmax(std::shared_ptr<Tensor> y, std::shared_ptr<TensorDesc> y_desc_overwrite,
-                                     std::shared_ptr<Tensor> x, std::shared_ptr<TensorDesc> x_desc_overwrite) {
-    size_t key = CacheManager::createDescriptorKey(
-        y_desc_overwrite ? y_desc_overwrite : y->tdesc(),
-        x_desc_overwrite ? x_desc_overwrite : x->tdesc(),
-        nullptr, nullptr, nullptr);
+void InferenceContext::causalSoftmax(std::shared_ptr<Tensor> y,
+                                     std::shared_ptr<Tensor> x) {
+    size_t key = CacheManager::createDescriptorKey(y, x, nullptr, nullptr, nullptr);
 
     infiniopCausalSoftmaxDescriptor_t desc;
     if (!cache_manager->getCausalSoftmaxDescriptor(key, desc)) {
         RUN_INFINI(infiniopCreateCausalSoftmaxDescriptor(
-            rsrc->handle, &desc,
-            y_desc_overwrite ? y_desc_overwrite->desc() : y->desc(),
-            x_desc_overwrite ? x_desc_overwrite->desc() : x->desc()));
+            rsrc->handle, &desc, y->desc(), x->desc()));
         cache_manager->putCausalSoftmaxDescriptor(key, desc);
     }
 
@@ -139,8 +121,10 @@ void InferenceContext::causalSoftmax(std::shared_ptr<Tensor> y, std::shared_ptr<
                                      y->data(), x->data(), stream));
 }
 
-void InferenceContext::swiglu(std::shared_ptr<Tensor> out, std::shared_ptr<Tensor> up, std::shared_ptr<Tensor> gate) {
-    size_t key = CacheManager::createDescriptorKey(out->tdesc(), up->tdesc(), gate->tdesc(), nullptr, nullptr);
+void InferenceContext::swiglu(std::shared_ptr<Tensor> out,
+                              std::shared_ptr<Tensor> up,
+                              std::shared_ptr<Tensor> gate) {
+    size_t key = CacheManager::createDescriptorKey(out, up, gate, nullptr, nullptr);
 
     infiniopSwiGLUDescriptor_t desc;
     if (!cache_manager->getSwiGLUDescriptor(key, desc)) {
@@ -158,20 +142,15 @@ void InferenceContext::swiglu(std::shared_ptr<Tensor> out, std::shared_ptr<Tenso
                               out->data(), up->data(), gate->data(), stream));
 }
 
-void InferenceContext::randomSample(std::shared_ptr<Tensor> out, std::shared_ptr<TensorDesc> out_desc_overwrite,
-                                    std::shared_ptr<Tensor> prob, std::shared_ptr<TensorDesc> prob_desc_overwrite,
+void InferenceContext::randomSample(std::shared_ptr<Tensor> out,
+                                    std::shared_ptr<Tensor> prob,
                                     float random_val, float top_p, uint32_t top_k, float temperature) {
-    size_t key = CacheManager::createDescriptorKey(
-        out_desc_overwrite ? out_desc_overwrite : out->tdesc(),
-        prob_desc_overwrite ? prob_desc_overwrite : prob->tdesc(),
-        nullptr, nullptr, nullptr);
+    size_t key = CacheManager::createDescriptorKey(out, prob, nullptr, nullptr, nullptr);
 
     infiniopRandomSampleDescriptor_t desc;
     if (!cache_manager->getRandomSampleDescriptor(key, desc)) {
         RUN_INFINI(infiniopCreateRandomSampleDescriptor(
-            rsrc->handle, &desc,
-            out_desc_overwrite ? out_desc_overwrite->desc() : out->desc(),
-            prob_desc_overwrite ? prob_desc_overwrite->desc() : prob->desc()));
+            rsrc->handle, &desc, out->desc(), prob->desc()));
         cache_manager->putRandomSampleDescriptor(key, desc);
     }
 

src/models/inference_context.hpp

@@ -15,25 +15,28 @@ struct InferenceContext {
     InferenceContext(DeviceResource *rsrc, CacheManager *cache_manager, infinirtStream_t stream);
 
     void ensure_workspace(size_t required_size);
+
     void rmsnorm(std::shared_ptr<Tensor> y,
                  std::shared_ptr<Tensor> x,
                  std::shared_ptr<Tensor> w,
                  float epsilon);
-    void gemm(std::shared_ptr<Tensor> c, std::shared_ptr<TensorDesc> c_desc_overwrite,
-              std::shared_ptr<Tensor> a, std::shared_ptr<TensorDesc> a_desc_overwrite,
-              std::shared_ptr<Tensor> b, std::shared_ptr<TensorDesc> b_desc_overwrite,
+    void gemm(std::shared_ptr<Tensor> c,
+              std::shared_ptr<Tensor> a,
+              std::shared_ptr<Tensor> b,
               float alpha, float beta);
-    void rearrange(std::shared_ptr<Tensor> dst, std::shared_ptr<TensorDesc> dst_desc_overwrite,
-                   std::shared_ptr<Tensor> src, std::shared_ptr<TensorDesc> src_desc_overwrite);
+    void rearrange(std::shared_ptr<Tensor> dst,
+                   std::shared_ptr<Tensor> src);
     void rope(std::shared_ptr<Tensor> q,
               std::shared_ptr<Tensor> k,
               std::shared_ptr<Tensor> pos,
               std::shared_ptr<Tensor> sin,
               std::shared_ptr<Tensor> cos);
-    void causalSoftmax(std::shared_ptr<Tensor> y, std::shared_ptr<TensorDesc> y_desc_overwrite,
-                       std::shared_ptr<Tensor> x, std::shared_ptr<TensorDesc> x_desc_overwrite);
-    void swiglu(std::shared_ptr<Tensor> out, std::shared_ptr<Tensor> up, std::shared_ptr<Tensor> gate);
-    void randomSample(std::shared_ptr<Tensor> out, std::shared_ptr<TensorDesc> out_desc_overwrite,
-                      std::shared_ptr<Tensor> prob, std::shared_ptr<TensorDesc> prob_desc_overwrite,
+    void causalSoftmax(std::shared_ptr<Tensor> y,
+                       std::shared_ptr<Tensor> x);
+    void swiglu(std::shared_ptr<Tensor> out,
+                std::shared_ptr<Tensor> up,
+                std::shared_ptr<Tensor> gate);
+    void randomSample(std::shared_ptr<Tensor> out,
+                      std::shared_ptr<Tensor> prob,
                       float random_val, float top_p, uint32_t top_k, float temperature);
 };

src/models/jiuge/jiuge.cpp

@@ -166,15 +166,10 @@ void inferDeviceBatch(const JiugeMeta &meta, DeviceResource &rsrc,
     }
 
     // Attention
-    auto qkv_desc = TensorDesc::create(dt_logits, qkv_buf->shape(), qkv_buf->strides());
-    auto b_attn_qkv_desc = TensorDesc::create(dt_logits, {ntok, (nh + nkvh * 2) * dh}, {0, 1});
-    auto o_desc = TensorDesc::create(dt_logits, o_buf->shape(), o_buf->strides());
-
     qkv_buf->dimSplit(1, {nh + nkvh * 2, dh}); // (ntok, nh + 2 * nkvh, dh)
     // attention inner
     size_t max_qk_size = 0;
     size_t max_seq_len = 0;
-    o_buf->dimSplit(1, {nh, dh});
 
     for (uint32_t req = 0; req < nreq; req++) {
         auto past_len = req_pos[req];
@@ -193,24 +188,19 @@ void inferDeviceBatch(const JiugeMeta &meta, DeviceResource &rsrc,
     auto gate_buf = gate_up_buf->slice(1, 0, di);
     auto up_buf = gate_up_buf->slice(1, di, di);
 
-    // Output and sample
-    auto result_desc = TensorDesc::create(INFINI_DTYPE_I64, {}, {});
-    auto prob_desc = TensorDesc::create(dt_logits, {dvoc}, {1});
-
     // Compute
     for (uint32_t layer = 0; layer < nlayer; layer++) {
         // 1. Attention
         // rms norm
         ctx.rmsnorm(logits_out, logits_in, rsrc.w_attn_norm[layer], meta.epsilon);
         // qkv_proj
+        qkv_buf->dimMerge(1, 2);
         if (has_qkv_bias) {
-            ctx.rearrange(qkv_buf, qkv_desc, rsrc.b_attn_qkv[layer], b_attn_qkv_desc);
+            ctx.rearrange(qkv_buf, rsrc.b_attn_qkv[layer]->reDesc({ntok, (nh + nkvh * 2) * dh}, {0, 1}));
         }
-        ctx.gemm(qkv_buf, qkv_desc,
-                 logits_out, nullptr,
-                 rsrc.w_attn_qkv[layer], nullptr,
-                 1.0, has_qkv_bias ? 1.0 : 0.0);
+        ctx.gemm(qkv_buf, logits_out, rsrc.w_attn_qkv[layer], 1.0, has_qkv_bias ? 1.0 : 0.0);
         // rope
+        qkv_buf->dimSplit(1, {nh + nkvh * 2, dh});
         ctx.rope(qkv_buf->slice(1, 0, nh), qkv_buf->slice(1, 0, nh), pos_ids_buf, rsrc.sin_table, rsrc.cos_table);
         ctx.rope(qkv_buf->slice(1, nh, nkvh), qkv_buf->slice(1, nh, nkvh), pos_ids_buf, rsrc.sin_table, rsrc.cos_table);
 
@@ -219,43 +209,41 @@ void inferDeviceBatch(const JiugeMeta &meta, DeviceResource &rsrc,
             auto past_len = req_pos[req];
             auto seq_len = req_lens[req];
             auto total_len = past_len + seq_len;
-            auto o = o_buf->slice({{0, token_offset, seq_len}});
+            auto o = o_buf->dimSplit(1, {nh, dh})->slice({{0, token_offset, seq_len}});
             auto q = qkv_buf->slice({{0, token_offset, seq_len}, {1, 0, nh}});
             auto k = qkv_buf->slice({{0, token_offset, seq_len}, {1, nh, nkvh}});
             auto v = qkv_buf->slice({{0, token_offset, seq_len}, {1, nh + nkvh, nkvh}});
 
-            auto qt_rearrange_desc = TensorDesc::create(dt_logits, {nkvh, ngroup, seq_len, dh});
-            auto qt_gemm_desc = TensorDesc::create(dt_logits, {nkvh, ngroup * seq_len, dh});
-            auto qk_gemm_desc = TensorDesc::create(dt_logits, {nkvh, ngroup * seq_len, total_len});
             // self attention
             // concat
-            ctx.rearrange(kv_caches[req]->k[idev][layer]->slice(0, past_len, seq_len), nullptr, k, nullptr);
-            ctx.rearrange(kv_caches[req]->v[idev][layer]->slice(0, past_len, seq_len), nullptr, v, nullptr);
+            ctx.rearrange(kv_caches[req]->k[idev][layer]->slice(0, past_len, seq_len), k);
+            ctx.rearrange(kv_caches[req]->v[idev][layer]->slice(0, past_len, seq_len), v);
             // qk
-            ctx.rearrange(rearrange_q_buf, qt_rearrange_desc,
-                          q->dimSplit(1, {nkvh, ngroup})->permute({1, 2, 0, 3}), nullptr);
-            ctx.gemm(qk_buf, qk_gemm_desc,
-                     rearrange_q_buf, qt_gemm_desc,
-                     kv_caches[req]->k[idev][layer]->slice(0, 0, total_len)->permute({1, 2, 0}), nullptr,
-                     1. / sqrt(dh), 0.0);
+            ctx.rearrange(rearrange_q_buf->dimSplit(1, {ngroup, seq_len}),
+                          q->dimSplit(1, {nkvh, ngroup})->permute({1, 2, 0, 3}));
+            qk_buf->dimSplit(1, {seq_len, total_len});
+            qk_buf->dimSplit(0, {nkvh, ngroup});
+            qk_buf->dimMerge(1, 2);
+            ctx.gemm(qk_buf, rearrange_q_buf->dimMerge(1, 2), kv_caches[req]->k[idev][layer]->slice(0, 0, total_len)->permute({1, 2, 0}), 1. / sqrt(dh), 0.0);
             // softmax
-            auto qk_desc = TensorDesc::create(dt_logits, {nkvh * ngroup, seq_len, total_len});
-            ctx.causalSoftmax(qk_buf, qk_desc, qk_buf, qk_desc);
-            ctx.gemm(attn_val_buf, qt_gemm_desc,
-                     qk_buf, qk_gemm_desc,
-                     kv_caches[req]->v[idev][layer]->slice(0, 0, total_len)->permute({1, 0, 2}), nullptr,
-                     1.0, 0.0);
+            qk_buf->dimSplit(1, {ngroup, seq_len});
+            qk_buf->dimMerge(0, 1);
+            ctx.causalSoftmax(qk_buf, qk_buf);
+            qk_buf->dimSplit(0, {nkvh, ngroup});
+            qk_buf->dimMerge(1, 2);
+            ctx.gemm(attn_val_buf, qk_buf, kv_caches[req]->v[idev][layer]->slice(0, 0, total_len)->permute({1, 0, 2}), 1.0, 0.0);
+            qk_buf->dimSplit(1, {ngroup, seq_len});
+            qk_buf->dimMerge(2, 3);
+            qk_buf->dimMerge(0, 1);
             // rearrange attn val
-            ctx.rearrange(o, TensorDesc::createWithOrder(dt_logits, {nkvh, ngroup, seq_len, dh}, {1, 2, 0, 3}),
-                          attn_val_buf, qt_rearrange_desc);
+            attn_val_buf->dimSplit(0, {nkvh, ngroup});
+            ctx.rearrange(o->dimSplit(1, {nkvh, ngroup})->permute({1, 2, 0, 3}), attn_val_buf);
+            attn_val_buf->dimMerge(0, 1);
 
             token_offset += seq_len;
         }
         // o_proj
-        ctx.gemm(logits_in, nullptr,
-                 o_buf, o_desc,
-                 rsrc.w_attn_out[layer], nullptr,
-                 1.0, idev == 0 ? 1.0 : 0.0); // only rank 0 adds residual
+        ctx.gemm(logits_in, o_buf->dimMerge(1, 2), rsrc.w_attn_out[layer], 1.0, idev == 0 ? 1.0 : 0.0); // only rank 0 adds residual
 
         // All_reduce if distributed
         if (rsrc.comm != nullptr) {
@@ -267,15 +255,9 @@ void inferDeviceBatch(const JiugeMeta &meta, DeviceResource &rsrc,
         // 2. FFN
         // rms_norm
         ctx.rmsnorm(logits_out, logits_in, rsrc.w_ffn_norm[layer], meta.epsilon);
-        ctx.gemm(gate_up_buf, nullptr,
-                 logits_out, nullptr,
-                 rsrc.w_ffn_gate_up[layer], nullptr,
-                 1.0, 0.0);
+        ctx.gemm(gate_up_buf, logits_out, rsrc.w_ffn_gate_up[layer], 1.0, 0.0);
         ctx.swiglu(gate_buf, up_buf, gate_buf);
-        ctx.gemm(logits_in, nullptr,
-                 gate_buf, nullptr,
-                 rsrc.w_ffn_down[layer], nullptr,
-                 1.0, idev == 0 ? 1.0 : 0.0); // only rank 0 adds residual
+        ctx.gemm(logits_in, gate_buf, rsrc.w_ffn_down[layer], 1.0, idev == 0 ? 1.0 : 0.0); // only rank 0 adds residual
 
         // All_reduce if distributed
         if (rsrc.comm != nullptr) {
@@ -296,18 +278,15 @@ void inferDeviceBatch(const JiugeMeta &meta, DeviceResource &rsrc,
                         rsrc.w_out_norm,
                         meta.epsilon);
         }
-        ctx.gemm(prob_buf, nullptr,
-                 logits_out->slice(0, 0, nreq), nullptr,
-                 rsrc.w_out_embd, nullptr,
-                 1.0, 0.0);
+        ctx.gemm(prob_buf, logits_out->slice(0, 0, nreq), rsrc.w_out_embd, 1.0, 0.0);
         std::random_device _rd;
         std::mt19937 gen(_rd());
         token_offset = 0;
         for (uint32_t req = 0; req < nreq; req++) {
             auto seq_len = req_lens[req];
             float random_val = std::uniform_real_distribution<float>(0, 1)(gen);
-            ctx.randomSample(result_buf->slice(0, req, 1), result_desc,
-                             prob_buf->slice(0, req, 1), prob_desc,
+            ctx.randomSample(result_buf->reDesc({}, {}),
+                             prob_buf->reDesc({dvoc}, {1}),
                              random_val, topp[req], topk[req], temperature[req]);
             token_offset += seq_len;
         }
@@ -354,7 +333,7 @@ inferBatch(struct JiugeModel *model,
 
 void launchDevice(const JiugeMeta &meta, const JiugeWeights *weights, DeviceResource *rsrc, InferState &state, InferRequest &req,
                   infiniDevice_t device, int idev, int ndev, int dev_id, infinicclComm_t comm) {
-    CacheManager cache_manager(100);
+    CacheManager cache_manager(256);
     InferenceContext ctx(rsrc, &cache_manager, rsrc->stream);
 
     // Create Device Resource

src/tensor.hpp

@@ -78,6 +78,7 @@ public:
     void dimMerge(size_t dim_start, size_t dim_end);
     void dimSplit(size_t dim, const std::vector<size_t> &dims);
     void permute(const std::vector<size_t> &order);
+    void reDesc(const std::vector<size_t> new_shape, const std::vector<ptrdiff_t> new_strides);
 };
 
 class Tensor : public std::enable_shared_from_this<Tensor> {
@@ -110,6 +111,7 @@ public:
     std::shared_ptr<Tensor> dimSplit(size_t dim,
                                      const std::vector<size_t> &dims);
     std::shared_ptr<Tensor> permute(const std::vector<size_t> &order);
+    std::shared_ptr<Tensor> reDesc(const std::vector<size_t> new_shape, const std::vector<ptrdiff_t> new_strides);
     void *data(ptrdiff_t offset = 0);
     void const *data(ptrdiff_t offset = 0) const;
     void copyFrom(std::shared_ptr<Tensor const> src, infiniopHandle_t handle,
@@ -120,7 +122,6 @@ public:
     infiniDtype_t dtype() const;
     bool isContigous() const;
     infiniopTensorDescriptor_t desc() const;
-    std::shared_ptr<TensorDesc> tdesc() const;
     ptrdiff_t dataOffset() const;
     infiniDevice_t deviceType() const;
     int deviceId() const;

src/tensor/tensor.cpp

@@ -108,7 +108,6 @@ ptrdiff_t Tensor::dataOffset() const {
 }
 
 infiniopTensorDescriptor_t Tensor::desc() const { return _desc->desc(); }
-std::shared_ptr<TensorDesc> Tensor::tdesc() const { return _desc; }
 
 std::shared_ptr<Tensor> Tensor::buffer(infiniDtype_t dtype,
                                        const std::vector<size_t> &shape,

src/tensor/transform.cpp

@@ -114,3 +114,14 @@ std::shared_ptr<Tensor> Tensor::permute(const std::vector<size_t> &order) {
     this->_desc->permute(order);
     return shared_from_this();
 }
+
+void TensorDesc::reDesc(const std::vector<size_t> new_shape, const std::vector<ptrdiff_t> new_strides) {
+    this->_shape = new_shape;
+    this->_strides = new_strides;
+    this->resetDesc();
+}
+
+std::shared_ptr<Tensor> Tensor::reDesc(const std::vector<size_t> new_shape, const std::vector<ptrdiff_t> new_strides) {
+    this->_desc->reDesc(new_shape, new_strides);
+    return shared_from_this();
+}