Merge remote-tracking branch 'upstream/main'

Author: leshikus

include/imex/Dialect/Dist/Utils/Utils.h

@@ -206,7 +206,7 @@ inline ::mlir::ValueRange createLocalOffsetsOf(const ::mlir::Location &loc,
   return builder.create<::imex::dist::LocalOffsetsOfOp>(loc, ary).getLOffsets();
 }
 
-// create operation returning halo of distributed array
+// create operation returning all parts (owned + halos) of distributed array
 inline ::mlir::ValueRange createPartsOf(const ::mlir::Location &loc,
                                         ::mlir::OpBuilder &builder,
                                         ::mlir::Value ary) {

include/imex/Dialect/DistRuntime/IR/DistRuntimeOps.td

@@ -191,6 +191,37 @@ def GetHaloOp : DistRuntime_Op<"get_halo",
   let hasCanonicalizer = 1;
 }
 
+def CopyReshapeOp : DistRuntime_Op<"copy_reshape",
+    [Pure, DeclareOpInterfaceMethods<AsyncOpInterface>, AttrSizedOperandSegments]> {
+  let summary = "Copy adequate data from input to a new reshaped output";
+  let description = [{
+    Copy the necessary data from the distributed input array to the locally owned part of the output array.
+    The shape of the output array is assumed to be a reshaped version of the input's shape.
+
+    The local data is not modified.
+
+    Arguments:
+
+    - `team`: the distributed team owning the distributed array
+    - `gShape`: the global shape of the distributed array
+    - `lOffsets`: the offset of the local data within the global array
+    - `lArray`: the locally owned data
+    - `ngShape`: the global shape of the distributed output array
+    - `nlOffsets`: the offsets of the locally owned output array
+
+    `gShape`, `lOffsets` are variadic arguments with same size `ri` where `ri` is the rank of the global input array (e.g., one number for each dimension of the global input array).
+    `ngShape`, `nlOffsets` are variadic arguments with same size `ro` where `ro` is the rank of the global output array (e.g., one number for each dimension of the global output array).
+  }];
+  let arguments = (ins AnyAttr:$team,
+                       AnyType:$lArray, Variadic<Index>:$gShape, Variadic<Index>:$lOffsets,
+                       Variadic<Index>:$ngShape, Variadic<Index>:$nlOffsets, Variadic<Index>:$nlShape);
+  let results = (outs DistRuntime_AsyncHandle:$handle, AnyType:$nlArray);
+  let assemblyFormat = [{
+    $lArray `g_shape` $gShape `l_offs` $lOffsets `to` `n_g_shape` $ngShape `n_offs` $nlOffsets `n_shape` $nlShape attr-dict `:` `(` type(operands) `)` `->` `(` qualified(type(results)) `)`
+  }];
+  let hasCanonicalizer = 1;
+}
+
 def WaitOp : DistRuntime_Op<"wait", []> {
   let summary = "Wait for asynchronous operation to finish.";
   let description = [{

include/imex/Dialect/NDArray/IR/NDArrayOps.td

@@ -875,11 +875,11 @@ def ReshapeOp : NDArray_Op<"reshape", []> {
       See Array API.
   }];
 
-  let arguments = (ins AnyType:$src, Variadic<Index>:$shape, OptionalAttr<I1Attr>:$copy);
+  let arguments = (ins AnyType:$source, Variadic<Index>:$shape, OptionalAttr<I1Attr>:$copy);
   let results = (outs AnyType);
 
   let assemblyFormat = [{
-    $src $shape attr-dict `:` qualified(type($src)) `->` qualified(type(results))
+    $source $shape attr-dict `:` qualified(type($source)) `->` qualified(type(results))
   }];
 }
 

lib/Conversion/DistToStandard/DistToStandard.cpp

@@ -591,9 +591,7 @@ struct ReshapeOpConverter
                   ::imex::ndarray::ReshapeOp::Adaptor adaptor,
                   ::mlir::ConversionPatternRewriter &rewriter) const override {
 
-    return ::mlir::failure();
-#ifdef FIXME_RESHAPE
-    auto src = op.getArray();
+    auto src = op.getSource();
     auto srcDistType = src.getType().dyn_cast<::imex::ndarray::NDArrayType>();
     auto retDistType =
         op.getResult().getType().dyn_cast<::imex::ndarray::NDArrayType>();
@@ -603,10 +601,13 @@ struct ReshapeOpConverter
     }
 
     auto loc = op.getLoc();
+    auto elType = srcDistType.getElementType();
     auto dEnv = getDistEnv(srcDistType);
-    auto nShape = adaptor.getShape();
+    auto ngShape = adaptor.getShape();
     auto gShape = createGlobalShapeOf(loc, rewriter, src);
-    auto elType = srcDistType.getElementType();
+    auto lParts = createPartsOf(loc, rewriter, src);
+    auto lArray = lParts.size() == 1 ? lParts[0] : lParts[1];
+    auto lOffs = createLocalOffsetsOf(loc, rewriter, src);
 
     // Repartitioning is needed if any of the partitions' size is not a multiple
     // of the new chunksize.
@@ -615,37 +616,30 @@ struct ReshapeOpConverter
     assert(adaptor.getCopy().value_or(1) != 0 ||
            !"Distributed reshape currently requires copying");
 
-    auto team = dEnv.getTeam();
-    // get function args
-    auto lPart = createDefaultPartition(loc, rewriter, team, nShape);
-    auto lOffs = lPart.getLOffsets();
-    auto lShape = lPart.getLShape();
+    // FIXME: Check return type: Check that static sizes are the same as the
+    // default part sizes
 
-    // create output array with target size
-    auto outArray = rewriter.create<::imex::ndarray::CreateOp>(
-        loc, lShape, ::imex::ndarray::fromMLIR(elType), nullptr,
-        getNonDistEnvs(srcDistType));
-    auto lUMR = ::imex::ndarray::mkURMemRef(loc, rewriter, outArray);
-
-    auto idxType = rewriter.getIndexType();
-    auto gShapeMR = createURMemRefFromElements(rewriter, loc, idxType, gShape);
-    auto lOffsMR = createURMemRefFromElements(rewriter, loc, idxType, lOffs);
-    auto nShapePtr = createURMemRefFromElements(rewriter, loc, idxType, nShape);
+    auto team = dEnv.getTeam();
+    auto nPart = createDefaultPartition(loc, rewriter, team, ngShape);
+    auto nlOffs = nPart.getLOffsets();
+    auto nlShape = nPart.getLShape();
+    auto shp = getShapeFromValues(nlShape);
+    auto lRetType = ::imex::ndarray::NDArrayType::get(
+        shp, elType, getNonDistEnvs(retDistType));
 
     // call the idt runtime
-    auto fun = rewriter.getStringAttr(mkTypedFunc("_idtr_reshape", elType));
-    auto teamC = rewriter.create<::mlir::arith::ConstantOp>(
-        loc, team.cast<::mlir::IntegerAttr>());
-    (void)rewriter.create<::mlir::func::CallOp>(
-        loc, fun, ::mlir::TypeRange(),
-        ::mlir::ValueRange{teamC, gShapeMR, lOffsMR, lUMR, nShapePtr});
-
+    auto htype = ::imex::distruntime::AsyncHandleType::get(getContext());
+    auto nlArray = rewriter.create<::imex::distruntime::CopyReshapeOp>(
+        loc, ::mlir::TypeRange{htype, lRetType}, team, lArray, gShape, lOffs,
+        ngShape, nlOffs, nlShape);
+    (void)rewriter.create<::imex::distruntime::WaitOp>(loc,
+                                                       nlArray.getHandle());
     // finally init dist array
     rewriter.replaceOp(
-        op, createDistArray(loc, rewriter, team, nShape, lOffs, {outArray}));
+        op, createDistArray(loc, rewriter, team, ngShape, nlOffs,
+                            ::mlir::ValueRange{nlArray.getNlArray()}));
 
     return ::mlir::success();
-#endif
   }
 };
 

lib/Conversion/NDArrayToLinalg/NDArrayToLinalg.cpp

@@ -646,13 +646,13 @@ struct ReshapeLowering
     // check output type and get operands
     auto retArTyp = op.getType().dyn_cast<::imex::ndarray::NDArrayType>();
     auto srcArTyp =
-        op.getSrc().getType().dyn_cast<::imex::ndarray::NDArrayType>();
+        op.getSource().getType().dyn_cast<::imex::ndarray::NDArrayType>();
     if (!(retArTyp && srcArTyp)) {
       return ::mlir::failure();
     }
 
     auto loc = op.getLoc();
-    auto src = adaptor.getSrc();
+    auto src = adaptor.getSource();
     auto srcTnsr = src.getType().cast<::mlir::TensorType>();
     auto shape = adaptor.getShape();
     auto elTyp = srcTnsr.getElementType();

lib/Dialect/DistRuntime/IR/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_imex_dialect_library(IMEXDistRuntimeDialect
   DistRuntimeOps.cpp
   GetHaloOp.cpp
+  CopyReshapeOp.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/include/imex/Dialect/DistRuntime

lib/Dialect/DistRuntime/IR/CopyReshapeOp.cpp

@@ -0,0 +1,91 @@
+//===- CopyReshapeOp.cpp - distruntime dialect  -----------------*- C++ -*-===//
+//
+// Copyright 2023 Intel Corporation
+// Part of the IMEX Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file implements the CopyReshapeOp of the DistRuntime dialect.
+///
+//===----------------------------------------------------------------------===//
+
+#include <imex/Dialect/DistRuntime/IR/DistRuntimeOps.h>
+#include <imex/Dialect/NDArray/IR/NDArrayOps.h>
+#include <imex/Utils/PassUtils.h>
+
+namespace imex {
+namespace distruntime {
+
+::mlir::SmallVector<::mlir::Value> CopyReshapeOp::getDependent() {
+  return {getNlArray()};
+}
+
+} // namespace distruntime
+} // namespace imex
+
+namespace {
+
+/// Pattern to replace dynamically shaped result types
+/// by statically shaped result types.
+class CopyReshapeOpResultCanonicalizer final
+    : public mlir::OpRewritePattern<::imex::distruntime::CopyReshapeOp> {
+public:
+  using mlir::OpRewritePattern<
+      ::imex::distruntime::CopyReshapeOp>::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(::imex::distruntime::CopyReshapeOp op,
+                  ::mlir::PatternRewriter &rewriter) const override {
+
+    // check input type
+    auto nlArray = op.getNlArray();
+    auto nlType = nlArray.getType().dyn_cast<::imex::ndarray::NDArrayType>();
+    if (!nlType) {
+      return ::mlir::failure();
+    }
+    auto resShape = nlType.getShape();
+    if (!::mlir::ShapedType::isDynamicShape(resShape)) {
+      return ::mlir::failure();
+    }
+
+    // we compare result shape with expected shape and bail out if no new
+    // static info found
+    auto nlShape = ::imex::getShapeFromValues(op.getNlShape());
+    bool found = false;
+    for (auto i = 0u; i < nlShape.size(); ++i) {
+      if (::mlir::ShapedType::isDynamic(resShape[i]) &&
+          !::mlir::ShapedType::isDynamic(nlShape[i])) {
+        found = true;
+      }
+    }
+    if (!found) {
+      return ::mlir::failure();
+    }
+
+    auto elType = nlType.getElementType();
+    auto nType = nlType.cloneWith(nlShape, elType);
+    auto hType = ::imex::distruntime::AsyncHandleType::get(getContext());
+
+    auto newOp = rewriter.create<::imex::distruntime::CopyReshapeOp>(
+        op.getLoc(), ::mlir::TypeRange{hType, nType}, op.getTeamAttr(),
+        op.getLArray(), op.getGShape(), op.getLOffsets(), op.getNgShape(),
+        op.getNlOffsets(), op.getNlShape());
+
+    // cast to original types and replace op
+    auto res = rewriter.create<imex::ndarray::CastOp>(op.getLoc(), nlType,
+                                                      newOp.getNlArray());
+    rewriter.replaceOp(op, {newOp.getHandle(), res});
+
+    return ::mlir::success();
+  }
+};
+
+} // namespace
+
+void imex::distruntime::CopyReshapeOp::getCanonicalizationPatterns(
+    mlir::RewritePatternSet &results, mlir::MLIRContext *context) {
+  results.add<CopyReshapeOpResultCanonicalizer>(context);
+}

lib/Dialect/DistRuntime/IR/GetHaloOp.cpp

@@ -1,4 +1,4 @@
-//===- GetHaloOp.cpp - NDArray dialect  ----------------------*- C++ -*-===//
+//===- GetHaloOp.cpp - distruntime dialect  ---------------------*- C++ -*-===//
 //
 // Copyright 2023 Intel Corporation
 // Part of the IMEX Project, under the Apache License v2.0 with LLVM Exceptions.

lib/Dialect/DistRuntime/Transforms/DistRuntimeToIDTR.cpp

@@ -96,18 +96,20 @@ struct RuntimePrototypes {
     requireFunc(loc, builder, module, "_idtr_prank", {i64Type}, {indexType});
     requireFunc(loc, builder, module, "_idtr_reduce_all", {dataMRType, opType},
                 {});
-    requireFunc(loc, builder, module, "_idtr_reshape",
-                // team, gshape, loffs, ownOutpart, nshape
-                {i64Type, idxMRType, idxMRType, dataMRType, idxMRType}, {});
+    requireFunc(loc, builder, module, "_idtr_copy_reshape",
+                // team, gshape, loffs, lPart, ngshape, nloffs, nPart
+                {i64Type, idxMRType, idxMRType, dataMRType, idxMRType,
+                 idxMRType, dataMRType},
+                {i64Type});
     requireFunc(loc, builder, module, "_idtr_update_halo",
                 // team,    gshape,    loffs,     lPart,      bbOffset, bbShape,
                 // lHalo,   rHalo, key
                 {i64Type, idxMRType, idxMRType, dataMRType, idxMRType,
                  idxMRType, dataMRType, dataMRType, i64Type},
                 {i64Type});
     requireFunc(loc, builder, module, "_idtr_wait",
-                // handle, lHalo, rHalo
-                {i64Type, dataMRType, dataMRType}, {});
+                // handle
+                {i64Type}, {});
   }
 };
 
@@ -151,6 +153,58 @@ struct TeamMemberOpPattern
   }
 };
 
+struct CopyReshapeOpPattern
+    : public ::mlir::OpRewritePattern<::imex::distruntime::CopyReshapeOp> {
+  using ::mlir::OpRewritePattern<
+      ::imex::distruntime::CopyReshapeOp>::OpRewritePattern;
+
+  ::mlir::LogicalResult
+  matchAndRewrite(::imex::distruntime::CopyReshapeOp op,
+                  ::mlir::PatternRewriter &rewriter) const override {
+    auto lArray = op.getLArray();
+    auto arType = lArray.getType().dyn_cast<::imex::ndarray::NDArrayType>();
+    auto resType =
+        op.getNlArray().getType().dyn_cast<::imex::ndarray::NDArrayType>();
+    if (!arType || !resType) {
+      return ::mlir::failure();
+    }
+
+    auto loc = op.getLoc();
+    auto elType = resType.getElementType();
+    auto team = op.getTeam();
+    auto gShape = op.getGShape();
+    auto lOffs = op.getLOffsets();
+    auto ngShape = op.getNgShape();
+    auto nlOffs = op.getNlOffsets();
+    auto nlShape = op.getNlShape();
+
+    // create output array with target size
+    auto nlArray = rewriter.create<::imex::ndarray::CreateOp>(
+        loc, nlShape, ::imex::ndarray::fromMLIR(elType), nullptr,
+        resType.getEnvironments());
+
+    auto idxType = rewriter.getIndexType();
+    auto teamC = rewriter.create<::mlir::arith::ConstantOp>(
+        loc, team.cast<::mlir::IntegerAttr>());
+    auto gShapeMR = createURMemRefFromElements(rewriter, loc, idxType, gShape);
+    auto lOffsMR = createURMemRefFromElements(rewriter, loc, idxType, lOffs);
+    auto lArrayMR = ::imex::ndarray::mkURMemRef(loc, rewriter, lArray);
+    auto ngShapeMR =
+        createURMemRefFromElements(rewriter, loc, idxType, ngShape);
+    auto nlOffsMR = createURMemRefFromElements(rewriter, loc, idxType, nlOffs);
+    auto nlArrayMR = ::imex::ndarray::mkURMemRef(loc, rewriter, nlArray);
+
+    auto fun =
+        rewriter.getStringAttr(mkTypedFunc("_idtr_copy_reshape", elType));
+    auto handle = rewriter.create<::mlir::func::CallOp>(
+        loc, fun, rewriter.getI64Type(),
+        ::mlir::ValueRange{teamC, gShapeMR, lOffsMR, lArrayMR, ngShapeMR,
+                           nlOffsMR, nlArrayMR});
+    rewriter.replaceOp(op, {handle.getResult(0), nlArray});
+    return ::mlir::success();
+  }
+};
+
 /// @brief  lower GetHaloOp
 /// Determine sizes of halos, alloc halos and call idtr.
 /// Before accessing/reading from returned halos, the caller must
@@ -291,24 +345,9 @@ struct WaitOpPattern
   ::mlir::LogicalResult
   matchAndRewrite(::imex::distruntime::WaitOp op,
                   ::mlir::PatternRewriter &rewriter) const override {
-    auto loc = op.getLoc();
-    auto handle = op.getHandle();
-    auto uhOp = handle.getDefiningOp<::imex::distruntime::GetHaloOp>();
-    assert(uhOp);
-    auto lHalo = uhOp.getLHalo();
-    auto rHalo = uhOp.getRHalo();
-
-    auto arTyp = lHalo.getType().dyn_cast<::imex::ndarray::NDArrayType>();
-    if (!arTyp)
-      return ::mlir::failure();
-    auto elTyp = arTyp.getElementType();
-    auto lHaloMR = ::imex::ndarray::mkURMemRef(loc, rewriter, lHalo);
-    auto rHaloMR = ::imex::ndarray::mkURMemRef(loc, rewriter, rHalo);
-
-    auto fsa = rewriter.getStringAttr(mkTypedFunc("_idtr_wait", elTyp));
+    auto fsa = rewriter.getStringAttr("_idtr_wait");
     rewriter.replaceOpWithNewOp<::mlir::func::CallOp>(
-        op, fsa, ::mlir::TypeRange(),
-        ::mlir::ValueRange{handle, lHaloMR, rHaloMR});
+        op, fsa, ::mlir::TypeRange(), ::mlir::ValueRange{op.getHandle()});
 
     return ::mlir::success();
   }
@@ -325,8 +364,9 @@ struct DistRuntimeToIDTRPass
     RuntimePrototypes::add_prototypes(builder, this->getOperation());
 
     ::mlir::FrozenRewritePatternSet patterns;
-    insertPatterns<TeamSizeOpPattern, TeamMemberOpPattern, GetHaloOpPattern,
-                   AllReduceOpPattern, WaitOpPattern>(getContext(), patterns);
+    insertPatterns<CopyReshapeOpPattern, TeamSizeOpPattern, TeamMemberOpPattern,
+                   GetHaloOpPattern, AllReduceOpPattern, WaitOpPattern>(
+        getContext(), patterns);
     (void)::mlir::applyPatternsAndFoldGreedily(this->getOperation(), patterns);
   }; // runOnOperation()
 

test/CMakeLists.txt

@@ -82,7 +82,7 @@ add_lit_testsuite(check-static "Running the IMEX regression tests"
 set_target_properties(check-static PROPERTIES FOLDER "Tests")
 
 add_custom_target(check-imex
-        DEPENDS check-gen check-static
+        DEPENDS check-static
         )
 
 add_lit_testsuites(IMEX ${CMAKE_CURRENT_SOURCE_DIR} DEPENDS ${IMEX_TEST_DEPENDS})