From 504e3c4946e6488a87fec485f4a476a0e1663042 Mon Sep 17 00:00:00 2001 From: Stella Laurenzo Date: Sun, 5 Jul 2020 17:45:45 -0700 Subject: [PATCH] Fixup local ndarray<->tensor transforms to preserve shape. * Preserving shape across the copy ops makes more thing shaped by default. * Inference of ndarray types will now preserve the shape when specializing the dtype. --- lib/Dialect/Numpy/IR/NumpyDialect.cpp | 11 +++++++---- lib/Python/NpcompDialect.cpp | 15 +++++++++++++-- pytest/NumpyCompiler/array_basics.py | 2 +- pytest/NumpyCompiler/ndarray_inference.py | 2 +- pytest/NumpyCompiler/ufunc.py | 2 +- 5 files changed, 23 insertions(+), 9 deletions(-) diff --git a/lib/Dialect/Numpy/IR/NumpyDialect.cpp b/lib/Dialect/Numpy/IR/NumpyDialect.cpp index f5c395afb..605527971 100644 --- a/lib/Dialect/Numpy/IR/NumpyDialect.cpp +++ b/lib/Dialect/Numpy/IR/NumpyDialect.cpp @@ -219,11 +219,14 @@ NdArrayType::mapToCPAType(Typing::CPA::Context &context) { // anyway. dtype = context.getIRValueType(getDtype()); } - auto irCtor = [](Typing::CPA::ObjectValueType *ovt, - llvm::ArrayRef fieldTypes, - MLIRContext *mlirContext, llvm::Optional) { + // Safe to capture an ArrayRef backed by type storage since it is uniqued. + auto optionalShape = getOptionalShape(); + auto irCtor = [optionalShape](Typing::CPA::ObjectValueType *ovt, + llvm::ArrayRef fieldTypes, + MLIRContext *mlirContext, + llvm::Optional) { assert(fieldTypes.size() == 1); - return NdArrayType::get(fieldTypes.front()); + return NdArrayType::get(fieldTypes.front(), optionalShape); }; return Typing::CPA::newArrayType(context, irCtor, context.getIdentifier("!NdArray"), dtype); diff --git a/lib/Python/NpcompDialect.cpp b/lib/Python/NpcompDialect.cpp index c03cfdbf3..e054b4ea6 100644 --- a/lib/Python/NpcompDialect.cpp +++ b/lib/Python/NpcompDialect.cpp @@ -121,7 +121,13 @@ public: "numpy_copy_to_tensor_op"); } auto dtype = sourceType.getDtype(); - auto tensorType = UnrankedTensorType::get(dtype); + auto optionalShape = sourceType.getOptionalShape(); + TensorType tensorType; + if (optionalShape) { + tensorType = RankedTensorType::get(*optionalShape, dtype); + } else { + tensorType = UnrankedTensorType::get(dtype); + } OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(true); Location loc = self.pyOpBuilder.getCurrentLoc(); auto op = opBuilder.create( @@ -136,7 +142,12 @@ public: "numpy_create_array_from_tensor_op"); } auto dtype = sourceType.getElementType(); - auto ndarrayType = Numpy::NdArrayType::get(dtype); + llvm::Optional> optionalShape; + if (auto rankedTensorType = + sourceType.dyn_cast()) { + optionalShape = rankedTensorType.getShape(); + } + auto ndarrayType = Numpy::NdArrayType::get(dtype, optionalShape); OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(true); Location loc = self.pyOpBuilder.getCurrentLoc(); auto op = opBuilder.create( diff --git a/pytest/NumpyCompiler/array_basics.py b/pytest/NumpyCompiler/array_basics.py index fe9303e1c..1a6ad850a 100644 --- a/pytest/NumpyCompiler/array_basics.py +++ b/pytest/NumpyCompiler/array_basics.py @@ -13,6 +13,6 @@ global_data = (np.zeros((2, 3)) + [1.0, 2.0, 3.0] * np.reshape([1.0, 2.0], @import_global def global_array_to_const(): # CHECK: %[[CST:.*]] = constant dense<{{\[\[}}1.000000e+00, 2.000000e+00, 3.000000e+00], [2.000000e+00, 4.000000e+00, 6.000000e+00]]> : tensor<2x3xf64> - # CHECK: numpy.create_array_from_tensor %[[CST]] : (tensor<2x3xf64>) -> !numpy.ndarray<*:f64> + # CHECK: numpy.create_array_from_tensor %[[CST]] : (tensor<2x3xf64>) -> !numpy.ndarray<[2,3]:f64> local_data = global_data return local_data diff --git a/pytest/NumpyCompiler/ndarray_inference.py b/pytest/NumpyCompiler/ndarray_inference.py index 3bff22a26..139ce1601 100644 --- a/pytest/NumpyCompiler/ndarray_inference.py +++ b/pytest/NumpyCompiler/ndarray_inference.py @@ -20,6 +20,6 @@ b = np.asarray([3.0, 4.0]) @import_global def global_add(): # CHECK-NOT: UnknownType - # CHECK: numpy.builtin_ufunc_call<"numpy.add"> ({{.*}}, {{.*}}) : (tensor<*xf64>, tensor<*xf64>) -> tensor<*xf64> + # CHECK: numpy.builtin_ufunc_call<"numpy.add"> ({{.*}}, {{.*}}) : (tensor<2xf64>, tensor<2xf64>) -> tensor<*xf64> # CHECK-NOT: UnknownType return np.add(a, b) diff --git a/pytest/NumpyCompiler/ufunc.py b/pytest/NumpyCompiler/ufunc.py index 14334bf65..503d654a6 100644 --- a/pytest/NumpyCompiler/ufunc.py +++ b/pytest/NumpyCompiler/ufunc.py @@ -24,7 +24,7 @@ def global_add(): # CHECK-DAG: %[[B_ARRAY:.*]] = numpy.create_array_from_tensor %[[CST_B_TENSOR]] # CHECK-DAG: %[[A:.*]] = numpy.copy_to_tensor %[[A_ARRAY]] # CHECK-DAG: %[[B:.*]] = numpy.copy_to_tensor %[[B_ARRAY]] - # CHECK: %[[R_TENSOR:.*]] = numpy.builtin_ufunc_call<"numpy.add"> (%[[A]], %[[B]]) : (tensor<*xf64>, tensor<*xf64>) -> tensor<*x!basicpy.UnknownType> + # CHECK: %[[R_TENSOR:.*]] = numpy.builtin_ufunc_call<"numpy.add"> (%[[A]], %[[B]]) : (tensor<2xf64>, tensor<2xf64>) -> tensor<*x!basicpy.UnknownType> # CHECK: numpy.create_array_from_tensor %[[R_TENSOR]] : (tensor<*x!basicpy.UnknownType>) -> !numpy.ndarray<*:?> return np.add(a, b)