Minor fix in UST level order for DIA

examples/sparse_tensor.cu

@@ -200,6 +200,7 @@ int main([[maybe_unused]] int argc, [[maybe_unused]] char **argv) {
   auto AdiaJ = experimental::make_tensor_dia<experimental::DIA_INDEX_J>(
       dvals, doffsets, {6, 6});
   print(AdiaJ);
+  printf("DiaJ(3,3)=%f\n", AdiaJ(3, 3));
 
   //
   // Perform a direct SpMV. This is also the correct way of performing

include/matx/core/make_sparse_tensor.h

@@ -222,16 +222,16 @@ auto make_tensor_dia(ValTensor &val, CrdTensor &off,
                          matxInvalidParameter, "data arrays should be rank-1");
   // Note that the DIA API typically does not involve positions.
   // However, under the formal DSL specifications, the top level
-  // compression should set up pos[0] = {0, nse}. This is done
+  // compression should set up pos[0] = {0, #diags}. This is done
   // here, using the same memory space as the other data.
   matxMemorySpace_t space = GetPointerKind(val.GetStorage().data());
   auto tp = makeDefaultNonOwningZeroStorage<POS>(2, space);
-  setVal(tp.data() + 1, static_cast<POS>(val.Size(0)), space);
+  setVal(tp.data() + 1, static_cast<POS>(off.Size(0)), space);
   // Construct DIA-I/J.
   using DIA = std::conditional_t<std::is_same_v<IDX, DIA_INDEX_I>, DIAI, DIAJ>;
   return sparse_tensor_t<VAL, CRD, POS, DIA>(
       shape, val.GetStorage(),
-      {makeDefaultNonOwningEmptyStorage<CRD>(), off.GetStorage()},
+      {off.GetStorage(), makeDefaultNonOwningEmptyStorage<CRD>()},
       {tp, makeDefaultNonOwningEmptyStorage<POS>()});
 }
 

include/matx/transforms/matmul/matvec_cusparse.h

@@ -318,8 +318,8 @@ void sparse_matvec_impl(TensorTypeC &C, const TensorTypeA &a,
     TA *AD = a.Data();
     TA *BD = b.Data();
     TA *CD = c.Data();
-    CRD *diags = a.CRDData(1);
-    uint64_t numD = a.crdSize(1);
+    CRD *diags = a.CRDData(0);
+    uint64_t numD = a.crdSize(0);
     uint64_t m = a.Size(0);
     uint64_t n = a.Size(1);
     uint32_t THREADS = static_cast<uint32_t>(std::min(m, 1024LU));

include/matx/transforms/solve/solve_cusparse.h

@@ -159,8 +159,8 @@ void sparse_dia_solve_impl(TensorTypeC &C, const TensorTypeA &a,
       MATX_THROW(matxNotSupported, "Tridiagonal solve overwrites rhs");
     }
     using CRD = typename atype::crd_type;
-    CRD *diags = a.CRDData(1);
-    const index_t numD = a.crdSize(1);
+    CRD *diags = a.CRDData(0);
+    const index_t numD = a.crdSize(0);
     if (numD != 3 || diags[0] != -1 || diags[1] != 0 || diags[2] != 1) {
       MATX_THROW(matxNotSupported, "Only tridiagonal solve supported");
     }