[NDArray] Add "permute_dims" operation (#841)

Adding "permute_dims" operation to ndarray

Author: AllanZyne

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

@@ -231,4 +231,39 @@ def WaitOp : DistRuntime_Op<"wait", []> {
   let arguments = (ins DistRuntime_AsyncHandle:$handle);
 }
 
+def CopyPermuteOp : DistRuntime_Op<"copy_permute",
+    [Pure, DeclareOpInterfaceMethods<AsyncOpInterface>, AttrSizedOperandSegments]> {
+  let summary = "Copy adequate data from input to a new permuted output";
+  let description = [{
+    Copy the necessary data from the distributed input array to the locally owned part of the new output array.
+    The shape of the output array is assumed to be a permuted 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
+    - `nlShape`: the local 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>:$nlShape,
+                       Variadic<Index>:$nlOffsets,
+                       DenseI64ArrayAttr:$axes);
+  let results = (outs DistRuntime_AsyncHandle:$handle, AnyType:$nlArray);
+  let assemblyFormat = [{
+    $lArray `g_shape` $gShape `l_offs` $lOffsets `to` `n_offs` $nlOffsets `n_shape` $nlShape `axes` $axes attr-dict `:` `(` type(operands) `)` `->` `(` qualified(type(results)) `)`
+  }];
+  let hasCanonicalizer = 1;
+}
+
 #endif // _DistRuntime_OPS_TD_INCLUDED_

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

@@ -948,4 +948,25 @@ def CastElemTypeOp: NDArray_Op<"cast_elemtype", [Pure]> {
     let hasCanonicalizer = 1;
 }
 
+def PermuteDimsOp : NDArray_Op<"permute_dims", []> {
+  let summary = "Permutes the axes (dimensions) of an array to a new array.";
+  let description = [{
+      Permutes the axes (dimensions) of an array.
+      The output array is a new array.
+  }];
+
+  let arguments = (ins
+    NDArray_NDArray:$source,
+    DenseI64ArrayAttr:$axes
+  );
+  let results = (outs NDArray_NDArray);
+
+  let assemblyFormat = [{
+    $source $axes attr-dict `:` qualified(type($source)) `->` qualified(type(results))
+  }];
+
+  let hasCanonicalizer = 1;
+  let hasVerifier = 1;
+}
+
 #endif // _NDARRAY_OPS_TD_INCLUDED_

lib/Conversion/DistToStandard/DistToStandard.cpp

@@ -1602,6 +1602,70 @@ struct RePartitionOpConverter
   }
 };
 
+/// Convert a global ndarray::PermuteDimsOp on a distributed array
+/// to ndarray::PermuteDimsOp on the local data.
+/// If needed, adds a repartition op.
+/// The local partition (e.g. a RankedTensor) is wrapped in a
+/// non-distributed NDArray and re-applied to PermuteDimsOp.
+/// op gets replaced with global distributed array
+struct PermuteDimsOpConverter
+    : public ::mlir::OpConversionPattern<::imex::ndarray::PermuteDimsOp> {
+  using ::mlir::OpConversionPattern<
+      ::imex::ndarray::PermuteDimsOp>::OpConversionPattern;
+
+  /// Initialize the pattern.
+  void initialize() {
+    /// Signal that this pattern safely handles recursive application.
+    setHasBoundedRewriteRecursion();
+  }
+
+  ::mlir::LogicalResult
+  matchAndRewrite(::imex::ndarray::PermuteDimsOp op,
+                  ::imex::ndarray::PermuteDimsOp::Adaptor adaptor,
+                  ::mlir::ConversionPatternRewriter &rewriter) const override {
+
+    auto src = op.getSource();
+    auto dst = op.getResult();
+    auto srcType = mlir::dyn_cast<::imex::ndarray::NDArrayType>(src.getType());
+    auto dstType = mlir::dyn_cast<::imex::ndarray::NDArrayType>(dst.getType());
+    if (!(srcType && isDist(srcType) && dstType && isDist(dstType))) {
+      return ::mlir::failure();
+    }
+
+    auto loc = op.getLoc();
+    auto srcEnv = getDistEnv(srcType);
+    auto team = srcEnv.getTeam();
+    auto elementType = srcType.getElementType();
+
+    auto srcGShape = createGlobalShapeOf(loc, rewriter, src);
+    auto srcLParts = createPartsOf(loc, rewriter, src);
+    auto srcLArray = srcLParts.size() == 1 ? srcLParts[0] : srcLParts[1];
+    auto srcLOffsets = createLocalOffsetsOf(loc, rewriter, src);
+
+    auto dstGShape = createGlobalShapeOf(loc, rewriter, dst);
+    auto dstLPart = createDefaultPartition(loc, rewriter, team, dstGShape);
+    auto dstLOffsets = dstLPart.getLOffsets();
+    auto dstLShape = dstLPart.getLShape();
+    auto dstLShapeIndex = getShapeFromValues(dstLShape);
+    auto dstLType = ::imex::ndarray::NDArrayType::get(
+        dstLShapeIndex, elementType, getNonDistEnvs(dstType));
+
+    // call the dist runtime
+    auto handleType = ::imex::distruntime::AsyncHandleType::get(getContext());
+    auto distLArray = rewriter.create<::imex::distruntime::CopyPermuteOp>(
+        loc, ::mlir::TypeRange{handleType, dstLType}, team, srcLArray,
+        srcGShape, srcLOffsets, dstLOffsets, dstLShape, adaptor.getAxes());
+    (void)rewriter.create<::imex::distruntime::WaitOp>(loc,
+                                                       distLArray.getHandle());
+    // finally init dist array
+    rewriter.replaceOp(
+        op, createDistArray(loc, rewriter, team, srcGShape, dstLOffsets,
+                            ::mlir::ValueRange{distLArray.getNlArray()}));
+
+    return ::mlir::success();
+  }
+};
+
 // *******************************
 // ***** Pass infrastructure *****
 // *******************************
@@ -1696,23 +1760,23 @@ struct ConvertDistToStandardPass
         ::imex::ndarray::ReductionOp, ::imex::ndarray::ToTensorOp,
         ::imex::ndarray::DeleteOp, ::imex::ndarray::CastElemTypeOp,
         ::imex::region::EnvironmentRegionOp,
-        ::imex::region::EnvironmentRegionYieldOp>(
+        ::imex::region::EnvironmentRegionYieldOp,
+        ::imex::ndarray::PermuteDimsOp>(
         [&](::mlir::Operation *op) { return typeConverter.isLegal(op); });
     target.addLegalOp<::imex::dist::InitDistArrayOp>();
 
     // All the dist conversion patterns/rewriter
     ::mlir::RewritePatternSet patterns(&ctxt);
     // all these patterns are converted
-    patterns
-        .insert<LinSpaceOpConverter, CreateOpConverter, CopyOpConverter,
-                ReductionOpConverter, ToTensorOpConverter,
-                InsertSliceOpConverter, SubviewOpConverter, EWBinOpConverter,
-                EWUnyOpConverter, LocalBoundingBoxOpConverter,
-                LocalCoreOpConverter, RePartitionOpConverter,
-                ReshapeOpConverter, LocalTargetOfSliceOpConverter,
-                DefaultPartitionOpConverter, LocalOffsetsOfOpConverter,
-                PartsOfOpConverter, DeleteOpConverter, CastElemTypeOpConverter>(
-            typeConverter, &ctxt);
+    patterns.insert<
+        LinSpaceOpConverter, CreateOpConverter, CopyOpConverter,
+        ReductionOpConverter, ToTensorOpConverter, InsertSliceOpConverter,
+        SubviewOpConverter, EWBinOpConverter, EWUnyOpConverter,
+        LocalBoundingBoxOpConverter, LocalCoreOpConverter,
+        RePartitionOpConverter, ReshapeOpConverter,
+        LocalTargetOfSliceOpConverter, DefaultPartitionOpConverter,
+        LocalOffsetsOfOpConverter, PartsOfOpConverter, DeleteOpConverter,
+        CastElemTypeOpConverter, PermuteDimsOpConverter>(typeConverter, &ctxt);
     mlir::scf::populateSCFStructuralTypeConversionsAndLegality(
         typeConverter, patterns, target);
     ::imex::populateRegionTypeConversionPatterns(patterns, typeConverter);

lib/Conversion/NDArrayToLinalg/NDArrayToLinalg.cpp

@@ -1217,6 +1217,39 @@ struct ReductionOpLowering
   }
 };
 
+/// Convert NDArray's permute_dims operations and their return type to
+/// Linalg/tensor.
+struct PermuteDimsOpLowering
+    : public ::mlir::OpConversionPattern<::imex::ndarray::PermuteDimsOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  ::mlir::LogicalResult
+  matchAndRewrite(::imex::ndarray::PermuteDimsOp op,
+                  ::imex::ndarray::PermuteDimsOp::Adaptor adaptor,
+                  ::mlir::ConversionPatternRewriter &rewriter) const override {
+
+    auto loc = op->getLoc();
+    auto srcTnsr = adaptor.getSource();
+
+    // convert src array to memref
+    auto srcArType = mlir::dyn_cast_or_null<::imex::ndarray::NDArrayType>(
+        op.getSource().getType());
+    if (!srcArType)
+      return mlir::failure();
+    auto srcMRType = srcArType.getMemRefType(srcTnsr);
+    auto srcMR = createToMemRef(loc, rewriter, srcTnsr, srcMRType);
+
+    auto perm = ::mlir::AffineMapAttr::get(::mlir::AffineMap::getPermutationMap(
+        adaptor.getAxes(), rewriter.getContext()));
+    mlir::memref::TransposeOp transposeOp =
+        rewriter.create<mlir::memref::TransposeOp>(loc, srcMR, perm);
+
+    rewriter.replaceOp(op, transposeOp.getResult());
+
+    return ::mlir::success();
+  }
+};
+
 // *******************************
 // ***** Pass infrastructure *****
 // *******************************
@@ -1320,13 +1353,14 @@ struct ConvertNDArrayToLinalgPass
         [&](mlir::Operation *op) { return typeConverter.isLegal(op); });
 
     ::mlir::RewritePatternSet patterns(&ctxt);
-    patterns.insert<
-        ToTensorLowering, SubviewLowering, ExtractSliceLowering,
-        InsertSliceLowering, ImmutableInsertSliceLowering, LinSpaceLowering,
-        LoadOpLowering, CreateLowering, EWBinOpLowering, DimOpLowering,
-        EWUnyOpLowering, ReductionOpLowering, ReshapeLowering, CastLowering,
-        CopyLowering, DeleteLowering, CastElemTypeLowering, FromMemRefLowering>(
-        typeConverter, &ctxt);
+    patterns.insert<ToTensorLowering, SubviewLowering, ExtractSliceLowering,
+                    InsertSliceLowering, ImmutableInsertSliceLowering,
+                    LinSpaceLowering, LoadOpLowering, CreateLowering,
+                    EWBinOpLowering, DimOpLowering, EWUnyOpLowering,
+                    ReductionOpLowering, ReshapeLowering, CastLowering,
+                    CopyLowering, DeleteLowering, CastElemTypeLowering,
+                    FromMemRefLowering, PermuteDimsOpLowering>(typeConverter,
+                                                               &ctxt);
     ::imex::populateRegionTypeConversionPatterns(patterns, typeConverter);
 
     // populate function boundaries using our special type converter

lib/Dialect/Dist/Transforms/DistCoalesce.cpp

@@ -99,6 +99,7 @@ struct DistCoalescePass : public ::imex::DistCoalesceBase<DistCoalescePass> {
     if (auto op =
             isDefByAnyOf<::imex::dist::InitDistArrayOp, ::imex::dist::EWBinOp,
                          ::imex::dist::EWUnyOp, ::imex::ndarray::ReshapeOp,
+                         ::imex::ndarray::PermuteDimsOp,
                          ::mlir::UnrealizedConversionCastOp,
                          ::imex::ndarray::CopyOp>(val)) {
       return op;

lib/Dialect/DistRuntime/IR/CMakeLists.txt

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

lib/Dialect/DistRuntime/IR/CopyPermuteOp.cpp

@@ -0,0 +1,127 @@
+//===- CopyPermuteOp.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 CopyPermuteOp 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> CopyPermuteOp::getDependent() {
+  return {getNlArray()};
+}
+
+} // namespace distruntime
+} // namespace imex
+
+namespace {
+
+/// Pattern to replace dynamically shaped result types
+/// by statically shaped result types.
+class CopyPermuteOpResultCanonicalizer final
+    : public mlir::OpRewritePattern<::imex::distruntime::CopyPermuteOp> {
+public:
+  using mlir::OpRewritePattern<
+      ::imex::distruntime::CopyPermuteOp>::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(::imex::distruntime::CopyPermuteOp op,
+                  ::mlir::PatternRewriter &rewriter) const override {
+
+    // check input type
+    auto dstArray = op.getNlArray();
+    auto dstType =
+        mlir::dyn_cast<::imex::ndarray::NDArrayType>(dstArray.getType());
+    if (!dstType) {
+      return ::mlir::failure();
+    }
+    auto dstShape = dstType.getShape();
+    if (!::mlir::ShapedType::isDynamicShape(dstShape)) {
+      return ::mlir::failure();
+    }
+
+    auto src = op.getLArray();
+    auto srcType = mlir::dyn_cast<::imex::ndarray::NDArrayType>(src.getType());
+    if (!srcType) {
+      return ::mlir::failure();
+    }
+    auto srcShape = srcType.getShape();
+    if (::mlir::ShapedType::isDynamicShape(srcShape)) {
+      return ::mlir::failure();
+    }
+
+    auto axes = op.getAxes();
+    if (axes.size() != dstShape.size()) {
+      return ::mlir::failure();
+    }
+
+    ::mlir::SmallVector<int64_t> permutedShape;
+    for (auto i = 0u; i < srcShape.size(); ++i) {
+      permutedShape.push_back(srcShape[axes[i]]);
+    }
+
+    auto elType = dstType.getElementType();
+    auto nType = dstType.cloneWith(permutedShape, elType);
+    auto hType = ::imex::distruntime::AsyncHandleType::get(getContext());
+
+    auto newOp = rewriter.create<::imex::distruntime::CopyPermuteOp>(
+        op.getLoc(), ::mlir::TypeRange{hType, nType}, op.getTeamAttr(),
+        op.getLArray(), op.getGShape(), op.getLOffsets(), op.getNlOffsets(),
+        op.getNlShape(), op.getAxes());
+
+    // cast to original types and replace op
+    auto res = rewriter.create<imex::ndarray::CastOp>(op.getLoc(), dstType,
+                                                      newOp.getNlArray());
+    rewriter.replaceOp(op, {newOp.getHandle(), res});
+
+    return ::mlir::success();
+  }
+};
+
+/// Pattern to rewrite a subview op with CastOp arguments.
+/// Ported from mlir::tensor::ExtractSliceOp
+class CopyPermuteCastFolder final
+    : public mlir::OpRewritePattern<::imex::distruntime::CopyPermuteOp> {
+public:
+  using mlir::OpRewritePattern<
+      ::imex::distruntime::CopyPermuteOp>::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(::imex::distruntime::CopyPermuteOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    auto src = op.getLArray();
+    auto castOp = mlir::dyn_cast<imex::ndarray::CastOp>(src.getDefiningOp());
+    if (!castOp)
+      return mlir::failure();
+
+    if (!imex::ndarray::canFoldIntoConsumerOp(castOp))
+      return mlir::failure();
+
+    auto newOp = rewriter.create<::imex::distruntime::CopyPermuteOp>(
+        op.getLoc(), op->getResultTypes(), op.getTeamAttr(), castOp.getSource(),
+        op.getGShape(), op.getLOffsets(), op.getNlOffsets(), op.getNlShape(),
+        op.getAxes());
+    rewriter.replaceOp(op, newOp);
+
+    return mlir::success();
+  }
+};
+
+} // namespace
+
+void imex::distruntime::CopyPermuteOp::getCanonicalizationPatterns(
+    mlir::RewritePatternSet &results, mlir::MLIRContext *context) {
+  results.add<CopyPermuteOpResultCanonicalizer, CopyPermuteCastFolder>(context);
+}