fix bug and warning

Author: AllanZyne

src/EWBinOp.cpp

@@ -120,7 +120,7 @@ struct DeferredEWBinOp : public Deferred {
     auto av = dm.getDependent(builder, Registry::get(_a));
     auto bv = dm.getDependent(builder, Registry::get(_b));
 
-    auto aTyp = av.getType().cast<::imex::ndarray::NDArrayType>();
+    auto aTyp = ::mlir::cast<::imex::ndarray::NDArrayType>(av.getType());
     auto outElemType =
         ::imex::ndarray::toMLIR(builder, SHARPY::jit::getPTDType(_dtype));
     auto outTyp = aTyp.cloneWith(shape(), outElemType);

src/EWUnyOp.cpp

@@ -105,7 +105,7 @@ struct DeferredEWUnyOp : public Deferred {
                      jit::DepManager &dm) override {
     auto av = dm.getDependent(builder, Registry::get(_a));
 
-    auto aTyp = av.getType().cast<::imex::ndarray::NDArrayType>();
+    auto aTyp = ::mlir::cast<::imex::ndarray::NDArrayType>(av.getType());
     auto outTyp = aTyp.cloneWith(shape(), aTyp.getElementType());
 
     auto ndOpId = sharpy(_op);

src/IEWBinOp.cpp

@@ -71,7 +71,7 @@ struct DeferredIEWBinOp : public Deferred {
     auto av = dm.getDependent(builder, Registry::get(_a));
     auto bv = dm.getDependent(builder, Registry::get(_b));
 
-    auto aTyp = av.getType().cast<::imex::ndarray::NDArrayType>();
+    auto aTyp = ::mlir::cast<::imex::ndarray::NDArrayType>(av.getType());
     auto outTyp = aTyp.cloneWith(shape(), aTyp.getElementType());
 
     auto binop = builder.create<::imex::ndarray::EWBinOp>(

src/ManipOp.cpp

@@ -41,7 +41,7 @@ struct DeferredReshape : public Deferred {
             ? ::mlir::IntegerAttr()
             : ::imex::getIntAttr(builder, COPY_ALWAYS ? true : false, 1);
 
-    auto aTyp = av.getType().cast<::imex::ndarray::NDArrayType>();
+    auto aTyp = ::mlir::cast<::imex::ndarray::NDArrayType>(av.getType());
     auto outTyp = imex::dist::cloneWithShape(aTyp, shape());
 
     auto op =
@@ -106,7 +106,7 @@ struct DeferredAsType : public Deferred {
     // construct NDArrayType with same shape and given dtype
     ::imex::ndarray::DType ndDType = dispatch<convDType>(dtype);
     auto mlirElType = ::imex::ndarray::toMLIR(builder, ndDType);
-    auto arType = av.getType().dyn_cast<::imex::ndarray::NDArrayType>();
+    auto arType = ::mlir::dyn_cast<::imex::ndarray::NDArrayType>(av.getType());
     if (!arType) {
       throw std::invalid_argument(
           "Encountered unexpected ndarray type in astype.");
@@ -157,7 +157,7 @@ struct DeferredToDevice : public Deferred {
                      jit::DepManager &dm) override {
     auto av = dm.getDependent(builder, Registry::get(_a));
 
-    auto srcType = av.getType().dyn_cast<::imex::ndarray::NDArrayType>();
+    auto srcType = ::mlir::dyn_cast<::imex::ndarray::NDArrayType>(av.getType());
     if (!srcType) {
       throw std::invalid_argument(
           "Encountered unexpected ndarray type in to_device.");

src/NDArray.cpp

@@ -113,7 +113,9 @@ void NDArray::NDADeleter::operator()(NDArray *a) const {
             std::cerr << "sharpy fini: detected possible memory leak\n";
           } else {
             auto av = dm.addDependent(builder, a);
-            builder.create<::imex::ndarray::DeleteOp>(loc, av);
+            auto deleteOp = builder.create<::imex::ndarray::DeleteOp>(loc, av);
+            deleteOp->setAttr("bufferization.manual_deallocation",
+                              builder.getUnitAttr());
             dm.drop(a->guid());
           }
           return false;

src/ReduceOp.cpp

@@ -61,7 +61,7 @@ struct DeferredReduceOp : public Deferred {
     // FIXME reduction over individual dimensions is not supported
     auto av = dm.getDependent(builder, Registry::get(_a));
     // return type 0d with same dtype as input
-    auto aTyp = av.getType().cast<::imex::ndarray::NDArrayType>();
+    auto aTyp = ::mlir::cast<::imex::ndarray::NDArrayType>(av.getType());
     auto outTyp = ::imex::dist::cloneWithShape(aTyp, shape());
     // reduction op
     auto mop = sharpy2mlir(_op);

src/SetGetItem.cpp

@@ -277,7 +277,7 @@ struct DeferredGetItem : public Deferred {
     const auto &offs = _slc.offsets();
     const auto &sizes = shape();
     const auto &strides = _slc.strides();
-    auto aTyp = av.getType().cast<::imex::ndarray::NDArrayType>();
+    auto aTyp = ::mlir::cast<::imex::ndarray::NDArrayType>(av.getType());
     auto outTyp = ::imex::dist::cloneWithShape(aTyp, shape());
 
     // now we can create the NDArray op using the above Values

src/idtr.cpp

@@ -641,7 +641,6 @@ template <typename T> class ndarray {
     size_t size = lSize();
     id idx = firstLocalIndex();
     while (size--) {
-      std::cout << "idx: " << idx[0] << ", " << idx[1] << std::endl;
       callback(idx);
       idx.next(_gShape);
     }
@@ -713,12 +712,13 @@ template <typename T> class WaitPermute {
 public:
   WaitPermute(SHARPY::Transceiver *tc, SHARPY::Transceiver::WaitHandle hdl,
               SHARPY::rank_type nRanks, std::vector<Parts> &&parts,
-              std::vector<int64_t> &&axes, ndarray<T> &&output,
-              std::vector<T> &&receiveBuffer, std::vector<int> &&receiveOffsets,
+              std::vector<int64_t> &&axes, std::vector<int64_t> oGShape,
+              ndarray<T> &&output, std::vector<T> &&receiveBuffer,
+              std::vector<int> &&receiveOffsets,
               std::vector<int> &&receiveSizes)
       : tc(tc), hdl(hdl), nRanks(nRanks), parts(std::move(parts)),
-        axes(std::move(axes)), output(std::move(output)),
-        receiveBuffer(std::move(receiveBuffer)),
+        axes(std::move(axes)), oGShape(std::move(oGShape)),
+        output(std::move(output)), receiveBuffer(std::move(receiveBuffer)),
         receiveOffsets(std::move(receiveOffsets)),
         receiveSizes(std::move(receiveSizes)) {}
 
@@ -735,8 +735,6 @@ template <typename T> class WaitPermute {
     std::vector<size_t> receiveRankBufferCount(nRanks, 0);
     output.localIndices([&](const id &outputIndex) {
       id inputIndex = outputIndex.permute(axes);
-      std::cout << "inputIndex: " << inputIndex[0] << ", " << inputIndex[1]
-                << std::endl;
       auto rank = getInputRank(parts, inputIndex[0]);
       auto &count = receiveRankBufferCount[rank];
       output[outputIndex] = receiveRankBuffer[rank][count++];
@@ -749,6 +747,7 @@ template <typename T> class WaitPermute {
   SHARPY::rank_type nRanks;
   std::vector<Parts> parts;
   std::vector<int64_t> axes;
+  std::vector<int64_t> oGShape;
   ndarray<T> output;
   std::vector<T> receiveBuffer;
   std::vector<int> receiveOffsets;
@@ -893,8 +892,9 @@ WaitHandleBase *_idtr_copy_permute(SHARPY::DTypeId sharpytype,
                           receiveSizes.data(), receiveOffsets.data());
 
   auto wait = WaitPermute(tc, hdl, nRanks, std::move(parts), std::move(axes),
-                          std::move(output), std::move(receiveBuffer),
-                          std::move(receiveOffsets), std::move(receiveSizes));
+                          std::move(oGShape), std::move(output),
+                          std::move(receiveBuffer), std::move(receiveOffsets),
+                          std::move(receiveSizes));
 
   assert(parts.empty() && axes.empty() && receiveBuffer.empty() &&
          receiveOffsets.empty() && receiveSizes.empty());

src/jit/mlir.cpp

@@ -692,6 +692,7 @@ static const std::string cpu_pipeline =
     "canonicalize,"
     "imex-remove-temporaries,"
     "buffer-deallocation-pipeline,"
+    "convert-bufferization-to-memref,"
     "func.func(convert-linalg-to-parallel-loops),"
     "func.func(scf-parallel-loop-fusion),"
     "drop-regions,"