collaborating with separate dist/ptensor passes in imex

Author: fschlimb

CMakeLists.txt

@@ -104,6 +104,7 @@ set(DDPTSrcs
 set(IDTRSrcs
     ${PROJECT_SOURCE_DIR}/src/idtr.cpp
     ${PROJECT_SOURCE_DIR}/src/CollComm.cpp
+    ${PROJECT_SOURCE_DIR}/src/DDPTensorImpl.cpp
     ${PROJECT_SOURCE_DIR}/src/Deferred.cpp
     ${PROJECT_SOURCE_DIR}/src/Factory.cpp
     ${PROJECT_SOURCE_DIR}/src/Mediator.cpp

src/Creator.cpp

@@ -8,7 +8,7 @@
 #include <imex/internal/PassUtils.h>
 
 #include <mlir/IR/Builders.h>
-#include <mlir/Dialect/Arithmetic/IR/Arithmetic.h>
+#include <mlir/Dialect/Arith/IR/Arith.h>
 #include <mlir/Dialect/Shape/IR/Shape.h>
 #include <mlir/Dialect/Tensor/IR/Tensor.h>
 #include <mlir/Dialect/Linalg/IR/Linalg.h>
@@ -171,15 +171,17 @@ struct DeferredArange : public Deferred
         auto stop = ::imex::createInt(loc, builder, _end);
         auto step = ::imex::createInt(loc, builder, _step);
         auto dtype = builder.getI64Type(); // FIXME
-        auto artype = ::imex::ptensor::PTensorType::get(builder.getContext(), ::mlir::RankedTensorType::get({-1}, dtype), false, true);
+        auto artype = ::imex::ptensor::PTensorType::get(builder.getContext(), ::mlir::RankedTensorType::get({-1}, dtype), false);
         auto dmy = ::imex::createInt<1>(loc, builder, 0);
-        auto team = ::imex::createInt(loc, builder, 1);
+        auto team = ::imex::createInt(loc, builder, reinterpret_cast<uint64_t>(getTransceiver()));
         dm.addVal(this->guid(),
                   builder.create<::imex::ptensor::ARangeOp>(loc, artype, start, stop, step, dmy, team),
-                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides) {
+                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
+                         uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned) {
             assert(rank == 1);
             assert(strides[0] == 1);
-            this->set_value(std::move(mk_tnsr(_dtype, rank, allocated, aligned, offset, sizes, strides)));
+            this->set_value(std::move(mk_tnsr(_dtype, rank, allocated, aligned, offset, sizes, strides,
+                                              gs_allocated, gs_aligned, lo_allocated, lo_aligned)));
         });
         return false;
     }

src/EWBinOp.cpp

@@ -464,8 +464,10 @@ struct DeferredEWBinOp : public Deferred
         auto bv = dm.getDependent(builder, _b);
         dm.addVal(this->guid(),
                   builder.create<::imex::ptensor::EWBinOp>(loc, av.getType(), builder.getI32IntegerAttr(ddpt2mlir(_op)), av, bv),
-                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides) {
-            this->set_value(std::move(mk_tnsr(_dtype, rank, allocated, aligned, offset, sizes, strides)));
+                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
+                         uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned) {
+            this->set_value(std::move(mk_tnsr(_dtype, rank, allocated, aligned, offset, sizes, strides,
+                                              gs_allocated, gs_aligned, lo_allocated, lo_aligned)));
         });
         return false;
     }

src/ReduceOp.cpp

@@ -7,6 +7,7 @@
 #include "ddptensor/Factory.hpp"
 
 #include <imex/Dialect/PTensor/IR/PTensorOps.h>
+#include <imex/Dialect/Dist/IR/DistOps.h>
 #include <mlir/IR/Builders.h>
 #include <mlir/Dialect/Shape/IR/Shape.h>
 
@@ -121,18 +122,25 @@ struct DeferredReduceOp : public Deferred
     {
         // FIXME reduction over individual dimensions is not supported
         auto av = dm.getDependent(builder, _a);
-        auto aPtTyp = av.getType().dyn_cast<::imex::ptensor::PTensorType>();
-        assert(aPtTyp);
+        auto aDtTyp = av.getType().dyn_cast<::imex::dist::DistTensorType>();
+        ::mlir::Type dtype;
+        if(aDtTyp) {
+            dtype = aDtTyp.getPTensorType().getRtensor().getElementType();
+        } else {
+            auto aPtTyp = av.getType().dyn_cast<::imex::ptensor::PTensorType>();
+            dtype = aPtTyp.getRtensor().getElementType();
+        }
         // return type 0d with same dtype as input
-        auto dtype = aPtTyp.getRtensor().getElementType();
-        auto retPtTyp = ::imex::ptensor::PTensorType::get(builder.getContext(), ::mlir::RankedTensorType::get({}, dtype), false, true);
+        auto retPtTyp = ::imex::ptensor::PTensorType::get(builder.getContext(), ::mlir::RankedTensorType::get({}, dtype), false);
         // reduction op
         auto mop = ddpt2mlir(_op);
         auto op = builder.getIntegerAttr(builder.getIntegerType(sizeof(mop)*8), mop);
         dm.addVal(this->guid(),
                   builder.create<::imex::ptensor::ReductionOp>(loc, retPtTyp, op, av),
-                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides) {
-            this->set_value(std::move(mk_tnsr(_dtype, rank, allocated, aligned, offset, sizes, strides)));
+                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
+                         uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned) {
+            this->set_value(std::move(mk_tnsr(_dtype, rank, allocated, aligned, offset, sizes, strides,
+                                              gs_allocated, gs_aligned, lo_allocated, lo_aligned)));
         });
         return false;
     }

src/idtr.cpp

@@ -29,6 +29,20 @@ T * mr_to_ptr(void * ptr, intptr_t offset)
 
 extern "C" {
 
+// Return number of ranks/processes in given team/communicator
+uint64_t idtr_nprocs(int64_t team)
+{
+    return getTransceiver()->nranks();
+}
+#pragma weak _idtr_nprocs = idtr_nprocs
+
+// Return rank in given team/communicator
+uint64_t idtr_prank(int64_t team)
+{
+    return getTransceiver()->rank();
+}
+#pragma weak _idtr_prank = idtr_prank
+
 // Register a global tensor of given shape.
 // Returns guid.
 // The runtime does not own or manage any memory.

src/include/ddptensor/CppTypes.hpp

@@ -148,3 +148,7 @@ enum FactoryId : int {
     F_TONUMPY,
     FACTORY_LAST
 };
+
+// size of memreftype in number of intptr_t's
+inline uint64_t memref_sz(int rank) { return 3 + 2 * rank; }
+inline uint64_t dtensor_sz(int rank) { return 2 * memref_sz(1) + memref_sz(rank) + 1; };