Support Kokkos::submdspan for slicing Kokkos::mdspan arrays

integration_tests/CMakeLists.txt

@@ -188,6 +188,7 @@ RUN(NAME array_02.cpp LABELS gcc llvm)
 RUN(NAME array_03.cpp LABELS gcc llvm)
 RUN(NAME array_03_mdspan.cpp LABELS gcc llvm)
 RUN(NAME array_04.cpp LABELS gcc llvm)
+RUN(NAME array_04_mdspan.cpp LABELS gcc llvm)
 RUN(NAME array_05.cpp LABELS gcc llvm)
 RUN(NAME array_06.cpp LABELS gcc llvm)
 RUN(NAME array_07.cpp LABELS gcc llvm)

integration_tests/array_04.cpp

@@ -13,11 +13,17 @@ int main() {
 
     xt::xtensor<double, 1> arr2 {5.0, 6.0, 7.0};
     xt::xtensor<double, 1> res;
+    xt::xtensor_fixed<double, xt::xshape<3>> res_ = {7.0, 11.0, 14.0};
 
     res = xt::empty<double>({3});
     res = xt::view(arr1, 1) + arr2;
-    std::cout<< arr1 << arr2 <<std::endl;
-    std::cout << res << std::endl;
+    std::cout << arr1 << "\n";
+    std::cout << arr2 << "\n";
+    std::cout << res << "\n";
+
+    if( xt::any(xt::abs(res - res_) > 1e-8) ) {
+        exit(2);
+    }
 
     return 0;
 }

integration_tests/array_04_mdspan.cpp

@@ -0,0 +1,49 @@
+#include <iostream>
+#include <vector>
+#include <mdspan/mdspan.hpp>
+
+void print2(Kokkos::mdspan<double, Kokkos::dextents<int32_t, 2>>& arr) {
+    for( int32_t i = 0; i < arr.extent(0); i++ ) {
+        for( int32_t j = 0; j < arr.extent(1) ; j++ ) {
+            std::cout << arr(i, j) << " ";
+        }
+    }
+}
+
+void print1(Kokkos::mdspan<double, Kokkos::dextents<int32_t, 1>>& arr) {
+    for( int32_t i = 0; i < arr.size(); i++ ) {
+        std::cout << arr[i] << " ";
+    }
+    std::cout << "\n";
+}
+
+int main() {
+
+    std::vector<double> arr1_data = {1.0, 2.0, 3.0, 2.0, 5.0, 7.0, 2.0, 5.0, 7.0};
+    Kokkos::mdspan<double, Kokkos::dextents<int32_t, 2>> arr1(arr1_data.data(), 3, 3);
+
+    std::vector<double> arr2_data = {5.0, 6.0, 7.0};
+    Kokkos::mdspan<double, Kokkos::dextents<int32_t, 1>> arr2(arr2_data.data(), 3);
+    Kokkos::mdspan<double, Kokkos::dextents<int32_t, 1>> res;
+    std::vector<double> res_data_correct = {7.0, 11.0, 14.0};
+
+    std::vector<double> res_data; res_data.reserve(3);
+    res = Kokkos::mdspan<double, Kokkos::extents<int32_t, 3>>(res_data.data());
+    res = Kokkos::submdspan(arr1, 1, Kokkos::full_extent);
+    for( int32_t i = 0; i < res.size(); i++ ) {
+      res[i] = res[i] + arr2[i];
+    }
+
+    print2(arr1);
+    print1(arr2);
+    print1(res);
+
+
+    for( int32_t i = 0; i < res.size(); i++ ) {
+        if( res[i] != res_data_correct[i] ) {
+            exit(2);
+        }
+    }
+
+    return 0;
+}

src/lc/clang_ast_to_asr.cpp

@@ -43,16 +43,20 @@ enum SpecialFunc {
     PushBack,
     Clear,
     Data,
+    Reserve,
 };
 
 std::map<std::string, SpecialFunc> special_function_map = {
     {"printf", SpecialFunc::Printf},
     {"exit", SpecialFunc::Exit},
     {"view", SpecialFunc::View},
+    {"submdspan", SpecialFunc::View},
     {"shape", SpecialFunc::Shape},
+    {"extent", SpecialFunc::Shape},
     {"empty", SpecialFunc::Empty},
     {"fill", SpecialFunc::Fill},
     {"all", SpecialFunc::All},
+    {"full_extent", SpecialFunc::All},
     {"any", SpecialFunc::Any},
     {"not_equal", SpecialFunc::NotEqual},
     {"equal", SpecialFunc::Equal},
@@ -72,6 +76,7 @@ std::map<std::string, SpecialFunc> special_function_map = {
     {"push_back", SpecialFunc::PushBack},
     {"clear", SpecialFunc::Clear},
     {"data", SpecialFunc::Data},
+    {"reserve", SpecialFunc::Reserve},
 };
 
 class OneTimeUseString {
@@ -455,6 +460,7 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
                 if( array_type && is_third_party_cpp_array ) {
                     *is_third_party_cpp_array = true;
                     *array_type = ThirdPartyCPPArrayTypes::MDSpanArray;
+                    xshape_result->from_pointer_n(xtensor_fixed_dims.p, xtensor_fixed_dims.size());
                 }
                 type = ASRUtils::TYPE(ASR::make_Array_t(al, l,
                     ClangTypeToASRType(qual_type),
@@ -496,6 +502,34 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
                     xshape_result->push_back(al, dim);
                 }
                 return nullptr;
+            } else if( template_name == "extents" ) {
+                const std::vector<clang::TemplateArgument>& template_arguments = template_specialization->template_arguments();
+                if( xshape_result == nullptr ) {
+                    throw std::runtime_error("Result Vec<ASR::dimention_t>* not provided.");
+                }
+
+                const clang::QualType& qual_type = template_arguments.at(0).getAsType();
+                ASR::ttype_t* index_type = ClangTypeToASRType(qual_type);
+                if( !ASR::is_a<ASR::Integer_t>(*index_type) ||
+                    ASRUtils::extract_kind_from_ttype_t(index_type) != 4 ) {
+                    throw std::runtime_error("Only int32_t should be used for index type in Kokkos::dextents.");
+                }
+                ASR::expr_t* zero = ASRUtils::EXPR(ASR::make_IntegerConstant_t(al, l, 0,
+                    ASRUtils::TYPE(ASR::make_Integer_t(al, l, 4))));
+                for( int i = 1; i < template_arguments.size(); i++ ) {
+                    clang::Expr* clang_rank = template_arguments.at(i).getAsExpr();
+                    TraverseStmt(clang_rank);
+                    int rank = 0;
+                    if( !ASRUtils::extract_value(ASRUtils::EXPR(tmp.get()), rank) ) {
+                        throw std::runtime_error("Rank provided in the xshape must be a constant.");
+                    }
+                    ASR::dimension_t dim; dim.loc = l;
+                    dim.m_length = ASRUtils::EXPR(ASR::make_IntegerConstant_t(al, l, rank,
+                        ASRUtils::TYPE(ASR::make_Integer_t(al, l, 4))));
+                    dim.m_start = zero;
+                    xshape_result->push_back(al, dim);
+                }
+                return nullptr;
             } else if( template_name == "dextents" ) {
                 const std::vector<clang::TemplateArgument>& template_arguments = template_specialization->template_arguments();
                 if( xshape_result == nullptr ) {
@@ -553,7 +587,7 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
         } else if( clang_type->getTypeClass() == clang::Type::TypeClass::Record ) {
             const clang::CXXRecordDecl* record_type = clang_type->getAsCXXRecordDecl();
             std::string name = record_type->getNameAsString();
-            if( name == "xtensor_container" || name == "vector" ) {
+            if( name == "xtensor_container" || name == "vector" || name == "mdspan" ) {
                 return nullptr;
             }
             ASR::symbol_t* type_t = current_scope->resolve_symbol(name);
@@ -565,6 +599,8 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
             } else {
                 type = ASRUtils::TYPE(ASR::make_Struct_t(al, l, type_t));
             }
+        } else if( clang_type->getTypeClass() == clang::Type::TypeClass::SubstTemplateTypeParm ) {
+            return nullptr;
         } else {
             throw std::runtime_error("clang::QualType not yet supported " +
                 std::string(clang_type->getTypeClassName()));
@@ -995,6 +1031,13 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
                     ASRUtils::make_ArrayBroadcast_t_util(al, Lloc(x), obj, value);
                     tmp = ASR::make_Assignment_t(al, Lloc(x), obj, value, nullptr);
                     is_stmt_created = true;
+                } else {
+                    // This means that right hand side of operator=
+                    // was a CXXConstructor and there was an allocation
+                    // statement created for it. So skip and return should
+                    // done in this case.
+                    tmp = nullptr;
+                    is_stmt_created = false;
                 }
                 assignment_target = nullptr;
             } else if( ASRUtils::is_complex(*ASRUtils::expr_type(obj)) ||
@@ -1393,8 +1436,9 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
                 }
                 alloc_arg.m_dims = alloc_dims.p; alloc_arg.n_dims = alloc_dims.size();
                 alloc_args.push_back(al, alloc_arg);
-                tmp = ASR::make_Allocate_t(al, Lloc(x), alloc_args.p, alloc_args.size(),
-                    nullptr, nullptr, nullptr);
+                current_body->push_back(al, ASRUtils::STMT(ASR::make_Allocate_t(al, Lloc(x),
+                    alloc_args.p, alloc_args.size(), nullptr, nullptr, nullptr)));
+                tmp = nullptr;
                 is_stmt_created = true;
             } else {
                 throw std::runtime_error("Only {...} is allowed for supplying shape to xt::empty.");
@@ -1437,6 +1481,13 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
             tmp = ASRUtils::make_Cast_t_value(al, Lloc(x), callee, ASR::cast_kindType::ListToArray,
                 ASRUtils::TYPE(ASR::make_Array_t(al, Lloc(x), ASRUtils::get_contained_type(ASRUtils::expr_type(callee)),
                 dims.p, dims.size(), ASR::array_physical_typeType::UnboundedPointerToDataArray)));
+        } else if (sf == SpecialFunc::Reserve) {
+            if( args.size() > 1 ) {
+                throw std::runtime_error("std::vector::reserve should be called with only one argument.");
+            }
+
+            tmp = ASR::make_ListReserve_t(al, Lloc(x), callee, args[0]);
+            is_stmt_created = true;
         } else {
             throw std::runtime_error("Only printf and exit special functions supported");
         }
@@ -1579,7 +1630,37 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
         ASR::ttype_t* constructor_type = ClangTypeToASRType(x->getType(), &shape_result,
             &third_party_array_type, &is_third_party_array_type);
         if( is_third_party_array_type && third_party_array_type == ThirdPartyCPPArrayTypes::MDSpanArray ) {
-            if( x->getNumArgs() >= 1 ) {
+            if( x->getNumArgs() == 0 ) {
+                tmp = nullptr;
+                is_stmt_created = false;
+                return true;
+            }
+            if( x->getNumArgs() == 1 ) {
+                if( shape_result.size() != ASRUtils::extract_n_dims_from_ttype(constructor_type) ) {
+                    throw std::runtime_error("Allocation sizes for all dimensions not provided.");
+                }
+                if( !ASRUtils::is_fixed_size_array(shape_result.p, shape_result.size()) ) {
+                    throw std::runtime_error("Allocation size of Kokkos::mdspan array isn't specified correctly.");
+                }
+
+                Vec<ASR::alloc_arg_t> alloc_args; alloc_args.reserve(al, 1);
+                ASR::alloc_arg_t alloc_arg;
+                alloc_arg.loc = Lloc(x);
+                LCOMPILERS_ASSERT(assignment_target != nullptr);
+                alloc_arg.m_a = assignment_target;
+                alloc_arg.m_dims = shape_result.p; alloc_arg.n_dims = shape_result.size();
+                alloc_arg.m_len_expr = nullptr; alloc_arg.m_type = nullptr;
+                alloc_args.push_back(al, alloc_arg);
+                Vec<ASR::expr_t*> dealloc_args; dealloc_args.reserve(al, 1);
+                dealloc_args.push_back(al, assignment_target);
+                current_body->push_back(al, ASRUtils::STMT(ASR::make_ExplicitDeallocate_t(
+                    al, Lloc(x), dealloc_args.p, dealloc_args.size())));
+                current_body->push_back(al, ASRUtils::STMT(ASR::make_Allocate_t(al, Lloc(x),
+                    alloc_args.p, alloc_args.size(), nullptr, nullptr, nullptr)));
+                tmp = nullptr;
+                is_stmt_created = true;
+                return true;
+            } else if( x->getNumArgs() >= 1 ) {
                 clang::Expr* _data = x->getArg(0);
                 TraverseStmt(_data);
                 ASR::expr_t* list_to_array = ASRUtils::EXPR(tmp.get());
@@ -1602,7 +1683,7 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
                     TraverseStmt(argi);
                     ASR::dimension_t alloc_dim;
                     alloc_dim.loc = Lloc(x);
-                    alloc_dim.m_start = ASRUtils::EXPR(ASR::make_IntegerConstant_t(al, Lloc(x), 1,
+                    alloc_dim.m_start = ASRUtils::EXPR(ASR::make_IntegerConstant_t(al, Lloc(x), 0,
                         ASRUtils::TYPE(ASR::make_Integer_t(al, Lloc(x), 4))));
                     alloc_dim.m_length = ASRUtils::EXPR(tmp.get());
                     alloc_dims.push_back(al, alloc_dim);
@@ -1625,9 +1706,13 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
                 return true;
             } else {
                 throw std::runtime_error("Kokkos::mdspan constructor is called "
-                                         "with incorrect number of arguments, "
-                                         "expected 3 and found, " + std::to_string(x->getNumArgs()));
+                    "with incorrect number of arguments, expected " + std::to_string(
+                        ASRUtils::extract_n_dims_from_ttype(constructor_type) + 1) +
+                        " and found, " + std::to_string(x->getNumArgs()));
             }
+        } else if( constructor_type == nullptr ) {
+            clang::Expr* x_ = x->getArg(0);
+            return TraverseStmt(x_);
         } else if( x->getNumArgs() >= 0 ) {
             return clang::RecursiveASTVisitor<ClangASTtoASRVisitor>::TraverseCXXConstructExpr(x);
         }
@@ -1751,8 +1836,12 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
             case clang::APValue::Struct: {
                 ASR::ttype_t* v_type = ASRUtils::type_get_past_const(ASRUtils::symbol_type(v));
                 if( !ASR::is_a<ASR::Struct_t>(*v_type) ) {
-                    throw std::runtime_error("Expected ASR::Struct_t type found, " +
-                        ASRUtils::type_to_str(v_type));
+                    // throw std::runtime_error("Expected ASR::Struct_t type found, " +
+                    //     ASRUtils::type_to_str(v_type));
+                    // No error, just return, clang marking something as Struct
+                    // doesn't mean that it necessarily maps to ASR Struct, it can
+                    // map to ASR Array type as well
+                    return ;
                 }
                 ASR::Struct_t* struct_t = ASR::down_cast<ASR::Struct_t>(v_type);
                 ASR::StructType_t* struct_type_t = ASR::down_cast<ASR::StructType_t>(
@@ -1820,6 +1909,9 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
                     tmp = ASR::make_Assignment_t(al, Lloc(x), var, init_val, nullptr);
                     is_stmt_created = true;
                 }
+            } else {
+                is_stmt_created = false;
+                tmp = nullptr;
             }
 
             if( x->getEvaluatedValue() ) {
@@ -2376,8 +2468,8 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
         ASR::symbol_t* sym = resolve_symbol(name);
         if( name == "operator<<" || name == "cout" || name == "endl" ||
             name == "operator()" || name == "operator+" || name == "operator=" ||
-            name == "operator*" || name == "view" || name == "empty" ||
-            name == "all" || name == "any" || name == "not_equal" ||
+            name == "operator*" || name == "view" || name == "submdspan" || name == "empty" ||
+            name == "all" || name == "full_extent" || name == "any" || name == "not_equal" ||
             name == "exit" || name == "printf" || name == "exp" ||
             name == "sum" || name == "amax" || name == "abs" ||
             name == "operator-" || name == "operator/" || name == "operator>" ||
@@ -2390,7 +2482,11 @@ class ClangASTtoASRVisitor: public clang::RecursiveASTVisitor<ClangASTtoASRVisit
                 throw std::runtime_error("Special function " + name + " cannot be overshadowed yet.");
             }
             if( sym == nullptr ) {
-                cxx_operator_name_obj.set(name);
+                if( name == "full_extent" ) {
+                    is_all_called = true;
+                } else {
+                    cxx_operator_name_obj.set(name);
+                }
                 tmp = nullptr;
                 return true;
             }

src/libasr/ASR.asdl

@@ -229,6 +229,7 @@ stmt
     | BlockCall(int label, symbol m)
     | SetInsert(expr a, expr ele)
     | SetRemove(expr a, expr ele)
+    | ListReserve(expr a, expr size)
     | ListInsert(expr a, expr pos, expr ele)
     | ListRemove(expr a, expr ele)
     | ListClear(expr a)

src/libasr/codegen/asr_to_llvm.cpp

@@ -1023,6 +1023,30 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
         visit_AllocateUtil(x, x.m_stat, false);
     }
 
+    void visit_ListReserve(const ASR::ListReserve_t& x) {
+        ASR::ttype_t* el_type = ASRUtils::get_contained_type(
+                                        ASRUtils::expr_type(x.m_a));
+        int64_t ptr_loads_copy = ptr_loads;
+        ptr_loads = 0;
+        this->visit_expr(*x.m_a);
+        llvm::Value* plist = tmp;
+
+        ptr_loads = 1;
+        this->visit_expr_wrapper(x.m_size, true);
+        ptr_loads = ptr_loads_copy;
+        llvm::Value *pos = tmp;
+
+        llvm::Value* list_data_ptr = list_api->get_pointer_to_list_data(plist);
+        llvm::Value* size = list_api->len(plist);
+        llvm::DataLayout data_layout(module.get());
+        llvm::Type* llvm_el_type = llvm_utils->get_type_from_ttype_t_util(el_type, module.get());
+        size_t el_struct_size = data_layout.getTypeAllocSize(llvm_el_type);
+        size = builder->CreateMul(size, llvm::ConstantInt::get(
+            llvm::Type::getInt32Ty(context), llvm::APInt(32, el_struct_size)));
+        llvm::Value* alloc_ptr = LLVM::lfortran_malloc(context, *module, *builder, size);
+        builder->CreateStore(builder->CreateBitCast(alloc_ptr, llvm_el_type->getPointerTo()), list_data_ptr);
+    }
+
     void visit_ReAlloc(const ASR::ReAlloc_t& x) {
         LCOMPILERS_ASSERT(x.n_args == 1);
         handle_allocated(x.m_args[0].m_a);

tests/reference/asr-array_04-f95b8eb.json

@@ -2,11 +2,11 @@
     "basename": "asr-array_04-f95b8eb",
     "cmd": "lc --show-asr --no-color {infile} -o {outfile}",
     "infile": "tests/../integration_tests/array_04.cpp",
-    "infile_hash": "dd38786061d4a19743e4c32c2a3ecb08193777f81a2a0d3e87af0f0a",
+    "infile_hash": "86bab3bcf300c07e9861e92ce21215829ae2893db1d378776b8ab332",
     "outfile": null,
     "outfile_hash": null,
     "stdout": "asr-array_04-f95b8eb.stdout",
-    "stdout_hash": "ecd1c3d056e6792b7c6ec0baa0730c74d60745d34d459ffb6122aee1",
+    "stdout_hash": "b72419a0a8bd62f822377669e6e98ca64a95400f6c7ce7b7ad33a68a",
     "stderr": null,
     "stderr_hash": null,
     "returncode": 0