[NDArray] Bugfix for "ndarray.permute_dims" operation (#854)

* fix argument order of distruntime.copy_permute
intel · Sep 9, 2024 · 1271507 · 1271507
1 parent 0c9aa46
commit 1271507
Show file tree

Hide file tree

Showing 7 changed files with 32 additions and 26 deletions.
diff --git a/include/imex/Dialect/DistRuntime/IR/DistRuntimeOps.td b/include/imex/Dialect/DistRuntime/IR/DistRuntimeOps.td
@@ -256,8 +256,8 @@ def CopyPermuteOp : DistRuntime_Op<"copy_permute",
                        AnyType:$lArray,
                        Variadic<Index>:$gShape,
                        Variadic<Index>:$lOffsets,
-                       Variadic<Index>:$nlShape,
                        Variadic<Index>:$nlOffsets,
+                       Variadic<Index>:$nlShape,
                        DenseI64ArrayAttr:$axes);
   let results = (outs DistRuntime_AsyncHandle:$handle, AnyType:$nlArray);
   let assemblyFormat = [{

diff --git a/include/imex/Dialect/NDArray/IR/NDArrayOps.td b/include/imex/Dialect/NDArray/IR/NDArrayOps.td
@@ -159,7 +159,8 @@ class NDArray_Op<string mnemonic, list<Trait> traits = []> :
     Op<NDArray_Dialect, mnemonic, traits>;
 
 
-def DeleteOp : NDArray_Op<"delete"> {
+def DeleteOp : NDArray_Op<"delete", [
+    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let summary = "Explicitly delete an NDArray, freeing its memory";
   let description = [{
       Allow explicitly deleting the memory of an NDArray. It is assumed

diff --git a/lib/Conversion/DistToStandard/DistToStandard.cpp b/lib/Conversion/DistToStandard/DistToStandard.cpp
@@ -256,7 +256,8 @@ struct DeleteOpConverter
 
     // apply DeleteOp to all parts
     for (auto p : lParts) {
-      (void)rewriter.create<::imex::ndarray::DeleteOp>(loc, p);
+      auto newOp = rewriter.create<::imex::ndarray::DeleteOp>(loc, p);
+      newOp->setAttrs(adaptor.getAttributes());
     }
 
     rewriter.eraseOp(op);
@@ -1659,7 +1660,7 @@ struct PermuteDimsOpConverter
                                                        distLArray.getHandle());
     // finally init dist array
     rewriter.replaceOp(
-        op, createDistArray(loc, rewriter, team, srcGShape, dstLOffsets,
+        op, createDistArray(loc, rewriter, team, dstGShape, dstLOffsets,
                             ::mlir::ValueRange{distLArray.getNlArray()}));
 
     return ::mlir::success();

diff --git a/lib/Conversion/NDArrayToLinalg/NDArrayToLinalg.cpp b/lib/Conversion/NDArrayToLinalg/NDArrayToLinalg.cpp
@@ -566,7 +566,9 @@ struct DeleteLowering
     auto inp = adaptor.getInput();
     auto inpMR = createToMemRef(op.getLoc(), rewriter, inp,
                                 inpArType.getMemRefType(inp));
-    rewriter.replaceOpWithNewOp<::mlir::memref::DeallocOp>(op, inpMR);
+    auto newOp =
+        rewriter.replaceOpWithNewOp<::mlir::memref::DeallocOp>(op, inpMR);
+    newOp->setAttrs(op->getAttrs());
 
     return ::mlir::success();
   }

diff --git a/lib/Dialect/DistRuntime/IR/CopyPermuteOp.cpp b/lib/Dialect/DistRuntime/IR/CopyPermuteOp.cpp
@@ -52,28 +52,9 @@ class CopyPermuteOpResultCanonicalizer final
       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 dstLShape = ::imex::getShapeFromValues(op.getNlShape());
     auto elType = dstType.getElementType();
-    auto nType = dstType.cloneWith(permutedShape, elType);
+    auto nType = dstType.cloneWith(dstLShape, elType);
     auto hType = ::imex::distruntime::AsyncHandleType::get(getContext());
 
     auto newOp = rewriter.create<::imex::distruntime::CopyPermuteOp>(

diff --git a/lib/Dialect/NDArray/IR/CMakeLists.txt b/lib/Dialect/NDArray/IR/CMakeLists.txt
@@ -9,6 +9,7 @@ add_imex_dialect_library(IMEXNDArrayDialect
   EWBinOp.cpp
   EWUnyOp.cpp
   PermuteDimsOp.cpp
+  DeleteOp.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/include/mlir/Dialect/NDArray

diff --git a/lib/Dialect/NDArray/IR/DeleteOp.cpp b/lib/Dialect/NDArray/IR/DeleteOp.cpp
@@ -0,0 +1,20 @@
+//===- DeleteOp.cpp - NDArray dialect  --------------------------*- C++ -*-===//
+//
+// Copyright 2024 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 DeleteOp of the NDArray dialect.
+///
+//===----------------------------------------------------------------------===//
+
+#include <imex/Dialect/NDArray/IR/NDArrayOps.h>
+
+void imex::ndarray::DeleteOp::getEffects(
+    ::llvm::SmallVectorImpl<::mlir::MemoryEffects::EffectInstance> &effects) {
+  effects.emplace_back(::mlir::MemoryEffects::Free().get());
+}