Add a trivial copy elision canonicalization on ndarray->tensor.

* This elides the very common code the compiler adds for chaining otherwise tensor-related numpy ops together.
* More aggressive canonicalizations would require more advanced analysis.

include/npcomp/Dialect/Numpy/IR/NumpyOps.td

@@ -76,6 +76,7 @@ def Numpy_CopyToTensorOp : Numpy_Op<"copy_to_tensor", [
   let assemblyFormat = [{
     $source attr-dict `:` functional-type($source, $dest)
   }];
+  let hasCanonicalizer = 1;
 }
 
 //----------------------------------------------------------------------------//

lib/Dialect/Numpy/IR/NumpyOps.cpp

@@ -10,6 +10,7 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/PatternMatch.h"
 #include "npcomp/Dialect/Basicpy/IR/BasicpyDialect.h"
 #include "npcomp/Dialect/Numpy/IR/NumpyDialect.h"
 
@@ -57,6 +58,34 @@ void BuiltinUfuncCallOp::addCPAConstraints(Typing::CPA::Context &context) {
   }
 }
 
+//----------------------------------------------------------------------------//
+// CreateArrayFromTensorOp
+//----------------------------------------------------------------------------//
+
+namespace {
+/// Match create_array_from_tensor -> copy_to_tensor and elide in favor
+/// of the original tensor.
+class ElideCreateRedundantArrayFromTensor
+    : public OpRewritePattern<CopyToTensorOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(CopyToTensorOp op,
+                                PatternRewriter &rewriter) const {
+    auto createArrayOp =
+        dyn_cast_or_null<CreateArrayFromTensorOp>(op.source().getDefiningOp());
+    if (createArrayOp && createArrayOp.dest().hasOneUse()) {
+      rewriter.replaceOp(op, createArrayOp.source());
+    }
+    return success();
+  }
+};
+} // namespace
+
+void CopyToTensorOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<ElideCreateRedundantArrayFromTensor>(context);
+}
+
 namespace mlir {
 namespace NPCOMP {
 namespace Numpy {

pytest/NumpyCompiler/ndarray_inference.py

@@ -1,4 +1,4 @@
-# RUN: %PYTHON %s | npcomp-opt -split-input-file -npcomp-cpa-type-inference | FileCheck %s --dump-input=fail
+# RUN: %PYTHON %s | npcomp-opt -split-input-file -npcomp-cpa-type-inference -canonicalize | FileCheck %s --dump-input=fail
 
 import numpy as np
 from npcomp.compiler import test_config
@@ -20,6 +20,7 @@ b = np.asarray([3.0, 4.0])
 @import_global
 def global_add():
   # CHECK-NOT: UnknownType
-  # CHECK: numpy.builtin_ufunc_call<"numpy.add"> ({{.*}}, {{.*}}) : (tensor<2xf64>, tensor<2xf64>) -> tensor<*xf64>
+  # CHECK: numpy.builtin_ufunc_call<"numpy.multiply"> ({{.*}}, {{.*}}) : (tensor<2xf64>, tensor<2xf64>) -> tensor<*xf64>
+  # CHECK: numpy.builtin_ufunc_call<"numpy.add"> ({{.*}}, {{.*}}) : (tensor<2xf64>, tensor<*xf64>) -> tensor<*xf64>
   # CHECK-NOT: UnknownType
-  return np.add(a, b)
+  return np.add(a, np.multiply(a, b))

test/Dialect/Numpy/canonicalizers.mlir

@@ -0,0 +1,27 @@
+// RUN: npcomp-opt -split-input-file %s -canonicalize | FileCheck --dump-input=fail %s
+
+// CHECK-LABEL: func @elideCreateRedundantArrayFromTensor
+func @elideCreateRedundantArrayFromTensor() -> tensor<2xf64> {
+  // CHECK: %[[CST:.*]] = constant
+  // CHECK-NOT: numpy.create_array_from_tensor
+  // CHECK-NOT: numpy.copy_to_tensor
+  %cst = constant dense<[1.000000e+00, 2.000000e+00]> : tensor<2xf64>
+  %0 = numpy.create_array_from_tensor %cst : (tensor<2xf64>) -> !numpy.ndarray<[2]:f64>
+  %1 = numpy.copy_to_tensor %0 : (!numpy.ndarray<[2]:f64>) -> tensor<2xf64>
+  // CHECK: return %[[CST]]
+  return %1 : tensor<2xf64>
+}
+
+// This test verifies that the very trivial elision is not overly aggressive.
+// Note that in this example, it is still safe to remove the copy, but the
+// analysis has not yet been written to do that safely.
+// CHECK-LABEL: func @elideCreateRedundantArrayFromTensorNonTrivial
+func @elideCreateRedundantArrayFromTensorNonTrivial() -> (tensor<2xf64>, tensor<2xf64>) {
+  // CHECK: numpy.create_array_from_tensor
+  // CHECK: numpy.copy_to_tensor
+  %cst = constant dense<[1.000000e+00, 2.000000e+00]> : tensor<2xf64>
+  %0 = numpy.create_array_from_tensor %cst : (tensor<2xf64>) -> !numpy.ndarray<[2]:f64>
+  %1 = numpy.copy_to_tensor %0 : (!numpy.ndarray<[2]:f64>) -> tensor<2xf64>
+  %2 = numpy.copy_to_tensor %0 : (!numpy.ndarray<[2]:f64>) -> tensor<2xf64>
+  return %1, %2 : tensor<2xf64>, tensor<2xf64>
+}