llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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.
pull/1/head
Stella Laurenzo 2020-07-05 17:45:45 -07:00
parent fae15ec5e7
commit 504e3c4946
5 changed files with 23 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
lib/Dialect/Numpy/IR/NumpyDialect.cpp
View File

@ -219,11 +219,14 @@ NdArrayType::mapToCPAType(Typing::CPA::Context &context) {
// anyway. // anyway.
dtype = context.getIRValueType(getDtype()); dtype = context.getIRValueType(getDtype());
} }
auto irCtor = [](Typing::CPA::ObjectValueType *ovt, // Safe to capture an ArrayRef backed by type storage since it is uniqued.
llvm::ArrayRef<mlir::Type> fieldTypes, auto optionalShape = getOptionalShape();
MLIRContext *mlirContext, llvm::Optional<Location>) { auto irCtor = [optionalShape](Typing::CPA::ObjectValueType *ovt,
llvm::ArrayRef<mlir::Type> fieldTypes,
MLIRContext *mlirContext,
llvm::Optional<Location>) {
assert(fieldTypes.size() == 1); assert(fieldTypes.size() == 1);
return NdArrayType::get(fieldTypes.front()); return NdArrayType::get(fieldTypes.front(), optionalShape);
}; };
return Typing::CPA::newArrayType(context, irCtor, return Typing::CPA::newArrayType(context, irCtor,
context.getIdentifier("!NdArray"), dtype); context.getIdentifier("!NdArray"), dtype);

15
lib/Python/NpcompDialect.cpp
View File

@ -121,7 +121,13 @@ public:
"numpy_copy_to_tensor_op"); "numpy_copy_to_tensor_op");
} }
auto dtype = sourceType.getDtype(); 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); OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(true);
Location loc = self.pyOpBuilder.getCurrentLoc(); Location loc = self.pyOpBuilder.getCurrentLoc();
auto op = opBuilder.create<Numpy::CopyToTensorOp>( auto op = opBuilder.create<Numpy::CopyToTensorOp>(
@ -136,7 +142,12 @@ public:
"numpy_create_array_from_tensor_op"); "numpy_create_array_from_tensor_op");
} }
auto dtype = sourceType.getElementType(); auto dtype = sourceType.getElementType();
auto ndarrayType = Numpy::NdArrayType::get(dtype); llvm::Optional<ArrayRef<int64_t>> optionalShape;
if (auto rankedTensorType =
sourceType.dyn_cast<RankedTensorType>()) {
optionalShape = rankedTensorType.getShape();
}
auto ndarrayType = Numpy::NdArrayType::get(dtype, optionalShape);
OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(true); OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(true);
Location loc = self.pyOpBuilder.getCurrentLoc(); Location loc = self.pyOpBuilder.getCurrentLoc();
auto op = opBuilder.create<Numpy::CreateArrayFromTensorOp>( auto op = opBuilder.create<Numpy::CreateArrayFromTensorOp>(

2
pytest/NumpyCompiler/array_basics.py
View File

@ -13,6 +13,6 @@ global_data = (np.zeros((2, 3)) + [1.0, 2.0, 3.0] * np.reshape([1.0, 2.0],
@import_global @import_global
def global_array_to_const(): 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: %[[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 local_data = global_data
return local_data return local_data

2
pytest/NumpyCompiler/ndarray_inference.py
View File

@ -20,6 +20,6 @@ b = np.asarray([3.0, 4.0])
@import_global @import_global
def global_add(): def global_add():
# CHECK-NOT: UnknownType # 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 # CHECK-NOT: UnknownType
return np.add(a, b) return np.add(a, b)

2
pytest/NumpyCompiler/ufunc.py
View File

@ -24,7 +24,7 @@ def global_add():
# CHECK-DAG: %[[B_ARRAY:.*]] = numpy.create_array_from_tensor %[[CST_B_TENSOR]] # CHECK-DAG: %[[B_ARRAY:.*]] = numpy.create_array_from_tensor %[[CST_B_TENSOR]]
# CHECK-DAG: %[[A:.*]] = numpy.copy_to_tensor %[[A_ARRAY]] # CHECK-DAG: %[[A:.*]] = numpy.copy_to_tensor %[[A_ARRAY]]
# CHECK-DAG: %[[B:.*]] = numpy.copy_to_tensor %[[B_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<*:?> # CHECK: numpy.create_array_from_tensor %[[R_TENSOR]] : (tensor<*x!basicpy.UnknownType>) -> !numpy.ndarray<*:?>
return np.add(a, b) return np.add(a, b)