From efbcf0aa4498ecff584bc6d087fa70ffbab9746c Mon Sep 17 00:00:00 2001 From: Stella Laurenzo Date: Wed, 8 Jul 2020 22:51:54 -0700 Subject: [PATCH] Add NumpyPublicFunctionsToTensor pass. * Rewrites public function signatures to operate on tensors (vs ndarray). * Most of our backends presume immutable tensors at public function boundaries. --- .../Dialect/Basicpy/Transforms/Passes.h | 1 - include/npcomp/Dialect/Numpy/CMakeLists.txt | 1 + .../npcomp/Dialect/Numpy/IR/NumpyDialect.h | 4 + .../Dialect/Numpy/Transforms/CMakeLists.txt | 6 ++ .../npcomp/Dialect/Numpy/Transforms/Passes.h | 26 +++++ .../npcomp/Dialect/Numpy/Transforms/Passes.td | 23 +++++ lib/CMakeLists.txt | 1 + lib/Dialect/Numpy/CMakeLists.txt | 1 + lib/Dialect/Numpy/IR/NumpyDialect.cpp | 9 ++ lib/Dialect/Numpy/Transforms/CMakeLists.txt | 17 ++++ lib/Dialect/Numpy/Transforms/PassDetail.h | 25 +++++ .../Transforms/PublicFunctionToTensor.cpp | 98 +++++++++++++++++++ lib/InitAll.cpp | 3 + .../Numpy/public_functions_to_tensor.mlir | 45 +++++++++ 14 files changed, 259 insertions(+), 1 deletion(-) create mode 100644 include/npcomp/Dialect/Numpy/Transforms/CMakeLists.txt create mode 100644 include/npcomp/Dialect/Numpy/Transforms/Passes.h create mode 100644 include/npcomp/Dialect/Numpy/Transforms/Passes.td create mode 100644 lib/Dialect/Numpy/Transforms/CMakeLists.txt create mode 100644 lib/Dialect/Numpy/Transforms/PassDetail.h create mode 100644 lib/Dialect/Numpy/Transforms/PublicFunctionToTensor.cpp create mode 100644 test/Dialect/Numpy/public_functions_to_tensor.mlir diff --git a/include/npcomp/Dialect/Basicpy/Transforms/Passes.h b/include/npcomp/Dialect/Basicpy/Transforms/Passes.h index 3dbfdb9d4..f6f79b131 100644 --- a/include/npcomp/Dialect/Basicpy/Transforms/Passes.h +++ b/include/npcomp/Dialect/Basicpy/Transforms/Passes.h @@ -18,7 +18,6 @@ namespace NPCOMP { namespace Basicpy { std::unique_ptr> createFunctionTypeInferencePass(); -std::unique_ptr> createCPAFunctionTypeInferencePass(); } // namespace Basicpy } // namespace NPCOMP diff --git a/include/npcomp/Dialect/Numpy/CMakeLists.txt b/include/npcomp/Dialect/Numpy/CMakeLists.txt index f33061b2d..9f57627c3 100644 --- a/include/npcomp/Dialect/Numpy/CMakeLists.txt +++ b/include/npcomp/Dialect/Numpy/CMakeLists.txt @@ -1 +1,2 @@ add_subdirectory(IR) +add_subdirectory(Transforms) diff --git a/include/npcomp/Dialect/Numpy/IR/NumpyDialect.h b/include/npcomp/Dialect/Numpy/IR/NumpyDialect.h index cc94aafd1..db84742dd 100644 --- a/include/npcomp/Dialect/Numpy/IR/NumpyDialect.h +++ b/include/npcomp/Dialect/Numpy/IR/NumpyDialect.h @@ -10,6 +10,7 @@ #define NPCOMP_DIALECT_NUMPY_IR_NUMPY_DIALECT_H #include "mlir/IR/Dialect.h" +#include "mlir/IR/StandardTypes.h" #include "npcomp/Dialect/Common.h" #include "npcomp/Typing/Analysis/CPA/Interfaces.h" @@ -64,6 +65,9 @@ public: /// unknown dimensions are -1. llvm::Optional> getOptionalShape(); + /// Converts to an equivalent TensorType. + TensorType toTensorType(); + // CPA::TypeMapInterface methods. Typing::CPA::TypeNode *mapToCPAType(Typing::CPA::Context &context); }; diff --git a/include/npcomp/Dialect/Numpy/Transforms/CMakeLists.txt b/include/npcomp/Dialect/Numpy/Transforms/CMakeLists.txt new file mode 100644 index 000000000..de80f15bd --- /dev/null +++ b/include/npcomp/Dialect/Numpy/Transforms/CMakeLists.txt @@ -0,0 +1,6 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls) +add_public_tablegen_target(NPCOMPNumpyPassIncGen) + +add_mlir_doc(Passes -gen-pass-doc NPCOMPNumpyTransforms ./) + diff --git a/include/npcomp/Dialect/Numpy/Transforms/Passes.h b/include/npcomp/Dialect/Numpy/Transforms/Passes.h new file mode 100644 index 000000000..01ac7de9f --- /dev/null +++ b/include/npcomp/Dialect/Numpy/Transforms/Passes.h @@ -0,0 +1,26 @@ +//===------------------------------------------------------------*- C++ -*-===// +// +// Part of the LLVM 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 +// +//===----------------------------------------------------------------------===// + +#ifndef NPCOMP_DIALECT_NUMPY_TRANSFORMS_PASSES_H +#define NPCOMP_DIALECT_NUMPY_TRANSFORMS_PASSES_H + +#include "mlir/Pass/Pass.h" + +#include + +namespace mlir { +namespace NPCOMP { +namespace Numpy { + +std::unique_ptr> createPublicFunctionsToTensorPass(); + +} // namespace Numpy +} // namespace NPCOMP +} // namespace mlir + +#endif // NPCOMP_DIALECT_NUMPY_TRANSFORMS_PASSES_H diff --git a/include/npcomp/Dialect/Numpy/Transforms/Passes.td b/include/npcomp/Dialect/Numpy/Transforms/Passes.td new file mode 100644 index 000000000..ae41945f7 --- /dev/null +++ b/include/npcomp/Dialect/Numpy/Transforms/Passes.td @@ -0,0 +1,23 @@ +//===-- Passes.td - Pass definition file -------------------*- tablegen -*-===// +// +// Part of the LLVM 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 +// +//===----------------------------------------------------------------------===// + +#ifndef NPCOMP_NUMPY_PASSES +#define NPCOMP_NUMPY_PASSES + +include "mlir/Pass/PassBase.td" + +//===----------------------------------------------------------------------===// +// TypeInference +//===----------------------------------------------------------------------===// + +def NumpyPublicFunctionsToTensor : Pass<"numpy-public-functions-to-tensor", "ModuleOp"> { + let summary = "Converts public functions to operate on tensors (instead of ndarray)"; + let constructor = "mlir::NPCOMP::Numpy::createPublicFunctionsToTensorPass()"; +} + +#endif // NPCOMP_NUMPY_PASSES diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt index 012ba34bb..da3ffe120 100644 --- a/lib/CMakeLists.txt +++ b/lib/CMakeLists.txt @@ -40,6 +40,7 @@ add_mlir_library(NPCOMPInitAll NPCOMPBasicpyDialect NPCOMPBasicpyPasses NPCOMPNumpyDialect + NPCOMPNumpyPasses NPCOMPTypingPasses ${ALL_DEPENDS} diff --git a/lib/Dialect/Numpy/CMakeLists.txt b/lib/Dialect/Numpy/CMakeLists.txt index f33061b2d..9f57627c3 100644 --- a/lib/Dialect/Numpy/CMakeLists.txt +++ b/lib/Dialect/Numpy/CMakeLists.txt @@ -1 +1,2 @@ add_subdirectory(IR) +add_subdirectory(Transforms) diff --git a/lib/Dialect/Numpy/IR/NumpyDialect.cpp b/lib/Dialect/Numpy/IR/NumpyDialect.cpp index 605527971..5362e53dd 100644 --- a/lib/Dialect/Numpy/IR/NumpyDialect.cpp +++ b/lib/Dialect/Numpy/IR/NumpyDialect.cpp @@ -210,6 +210,15 @@ llvm::Optional> NdArrayType::getOptionalShape() { return getImpl()->getOptionalShape(); } +TensorType NdArrayType::toTensorType() { + auto shape = getOptionalShape(); + if (shape) { + return RankedTensorType::get(*shape, getDtype()); + } else { + return UnrankedTensorType::get(getDtype()); + } +} + Typing::CPA::TypeNode * NdArrayType::mapToCPAType(Typing::CPA::Context &context) { llvm::Optional dtype; diff --git a/lib/Dialect/Numpy/Transforms/CMakeLists.txt b/lib/Dialect/Numpy/Transforms/CMakeLists.txt new file mode 100644 index 000000000..2eff1001d --- /dev/null +++ b/lib/Dialect/Numpy/Transforms/CMakeLists.txt @@ -0,0 +1,17 @@ +add_mlir_conversion_library(NPCOMPNumpyPasses + PublicFunctionToTensor.cpp + + ADDITIONAL_HEADER_DIRS + ${PROJECT_SOURCE_DIR}/include/npcomp/Dialect/Numpy/Transforms + + DEPENDS + NPCOMPNumpyPassIncGen + + LINK_COMPONENTS + Core + + LINK_LIBS PUBLIC + MLIRIR + MLIRPass + NPCOMPNumpyDialect +) diff --git a/lib/Dialect/Numpy/Transforms/PassDetail.h b/lib/Dialect/Numpy/Transforms/PassDetail.h new file mode 100644 index 000000000..4ca3a3c5d --- /dev/null +++ b/lib/Dialect/Numpy/Transforms/PassDetail.h @@ -0,0 +1,25 @@ +//===- PassDetail.h - Pass details ------------------------------*- C++ -*-===// +// +// Part of the LLVM 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 +// +//===----------------------------------------------------------------------===// + +#ifndef NPCOMP_DIALECT_NUMPY_TRANSFORMS_PASSDETAIL_H +#define NPCOMP_DIALECT_NUMPY_TRANSFORMS_PASSDETAIL_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace NPCOMP { +namespace Numpy { + +#define GEN_PASS_CLASSES +#include "npcomp/Dialect/Numpy/Transforms/Passes.h.inc" + +} // namespace Numpy +} // namespace NPCOMP +} // end namespace mlir + +#endif // NPCOMP_DIALECT_NUMPY_TRANSFORMS_PASSDETAIL_H diff --git a/lib/Dialect/Numpy/Transforms/PublicFunctionToTensor.cpp b/lib/Dialect/Numpy/Transforms/PublicFunctionToTensor.cpp new file mode 100644 index 000000000..d9358def7 --- /dev/null +++ b/lib/Dialect/Numpy/Transforms/PublicFunctionToTensor.cpp @@ -0,0 +1,98 @@ +//===- PublicFunctionToTensor.cpp - Type inference passes --------*- C++-*-===// +// +// This file is licensed 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 +// +//===----------------------------------------------------------------------===// + +#include "PassDetail.h" + +#include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Module.h" +#include "npcomp/Dialect/Numpy/IR/NumpyDialect.h" +#include "npcomp/Dialect/Numpy/IR/NumpyOps.h" +#include "npcomp/Dialect/Numpy/Transforms/Passes.h" + +using namespace mlir; +using namespace mlir::NPCOMP::Numpy; + +namespace { + +class PublicFunctionsToTensorPass + : public NumpyPublicFunctionsToTensorBase { + void runOnOperation() override { + auto module = getOperation(); + module.walk([&](FuncOp func) { + if (func.getVisibility() != SymbolTable::Visibility::Public) + return; + if (func.isExternal()) + return; + auto uses = SymbolTable::getSymbolUses(func, module); + if (!uses || uses->begin() != uses->end()) { + func.emitWarning() << "unimplemented: cannot convert ndarray->tensor " + << "signature for public function with uses"; + return; + } + rewriteSignature(func); + }); + } + + void rewriteSignature(FuncOp func) { + auto &entryBlock = func.getBlocks().front(); + auto funcType = func.getType(); + auto loc = func.getLoc(); + + // Rewrite inputs. + auto builder = OpBuilder::atBlockBegin(&entryBlock); + auto inputTypes = llvm::to_vector<4>(funcType.getInputs()); + for (unsigned i = 0; i < inputTypes.size(); ++i) { + auto arrayType = inputTypes[i].dyn_cast(); + if (!arrayType) + continue; + Type tensorType = arrayType.toTensorType(); + BlockArgument argument = entryBlock.getArgument(i); + argument.setType(tensorType); + auto createOp = + builder.create(loc, arrayType, argument); + argument.replaceAllUsesExcept(createOp, + SmallPtrSet{createOp}); + inputTypes[i] = tensorType; + } + + // Rewrite result signature. + auto resultTypes = llvm::to_vector<4>(funcType.getResults()); + for (auto &resultType : resultTypes) { + auto arrayType = resultType.dyn_cast(); + if (arrayType) + resultType = arrayType.toTensorType(); + } + + // Update signature. + funcType = + FunctionType::get(inputTypes, resultTypes, funcType.getContext()); + func.setType(funcType); + + // Rewrite all return terminators. + func.walk([&](ReturnOp term) { + OpBuilder builder(term); + for (unsigned i = 0; i < term.getNumOperands(); ++i) { + Value operand = term.getOperand(i); + auto arrayType = operand.getType().dyn_cast(); + if (!arrayType) + continue; + Type tensorType = arrayType.toTensorType(); + auto copyOp = builder.create(loc, tensorType, operand); + term.setOperand(i, copyOp); + } + }); + } +}; + +} // namespace + +std::unique_ptr> +mlir::NPCOMP::Numpy::createPublicFunctionsToTensorPass() { + return std::make_unique(); +} diff --git a/lib/InitAll.cpp b/lib/InitAll.cpp index dce812e34..999267b4b 100644 --- a/lib/InitAll.cpp +++ b/lib/InitAll.cpp @@ -12,6 +12,7 @@ #include "npcomp/Dialect/Basicpy/Transforms/Passes.h" #include "npcomp/Dialect/Npcomprt/IR/NpcomprtDialect.h" #include "npcomp/Dialect/Numpy/IR/NumpyDialect.h" +#include "npcomp/Dialect/Numpy/Transforms/Passes.h" #include "npcomp/Dialect/TCF/IR/TCFDialect.h" #include "npcomp/Dialect/TCP/IR/TCPDialect.h" #include "npcomp/Typing/Transforms/Passes.h" @@ -93,6 +94,8 @@ void mlir::NPCOMP::registerAllPasses() { #define GEN_PASS_REGISTRATION #include "npcomp/Dialect/Basicpy/Transforms/Passes.h.inc" #define GEN_PASS_REGISTRATION +#include "npcomp/Dialect/Numpy/Transforms/Passes.h.inc" +#define GEN_PASS_REGISTRATION #include "npcomp/Typing/Transforms/Passes.h.inc" registerDependencyPasses(); } diff --git a/test/Dialect/Numpy/public_functions_to_tensor.mlir b/test/Dialect/Numpy/public_functions_to_tensor.mlir new file mode 100644 index 000000000..f714c5b7b --- /dev/null +++ b/test/Dialect/Numpy/public_functions_to_tensor.mlir @@ -0,0 +1,45 @@ +// RUN: npcomp-opt -split-input-file %s -verify-diagnostics -allow-unregistered-dialect -numpy-public-functions-to-tensor | FileCheck --dump-input=fail %s + +// CHECK-LABEL: legalConversion +module @legalConversion { + // CHECK: @f(%arg0: tensor<3x?xf32>, %arg1: i32, %arg2: tensor<*xf32>) -> (i32, tensor<3x?xf32>, tensor<*xf32>) + func @f(%arg0: !numpy.ndarray<[3,?]:f32>, %arg1: i32, %arg2: !numpy.ndarray<*:f32>) -> + (i32, !numpy.ndarray<[3,?]:f32>, !numpy.ndarray<*:f32>) { + // CHECK: %[[CREATE0:.+]] = numpy.create_array_from_tensor %arg0 + // CHECK: %[[CREATE1:.+]] = numpy.create_array_from_tensor %arg2 + // CHECK: %[[R0:.+]] = "unfoldable_arg0"(%[[CREATE0]]) : (!numpy.ndarray<[3,?]:f32>) -> !numpy.ndarray<[3,?]:f32> + // CHECK: %[[R1:.+]] = "unfoldable_arg1"(%[[CREATE1]]) : (!numpy.ndarray<*:f32>) -> !numpy.ndarray<*:f32> + %0 = "unfoldable_arg0"(%arg0) : (!numpy.ndarray<[3,?]:f32>) -> !numpy.ndarray<[3,?]:f32> + %1 = "unfoldable_arg1"(%arg2) : (!numpy.ndarray<*:f32>) -> !numpy.ndarray<*:f32> + // CHECK: %[[COPY0:.+]] = numpy.copy_to_tensor %[[R0]] + // CHECK: %[[COPY1:.+]] = numpy.copy_to_tensor %[[R1]] + // CHECK: return %arg1, %[[COPY0]], %[[COPY1]] : i32, tensor<3x?xf32>, tensor<*xf32> + return %arg1, %0, %1 : i32, !numpy.ndarray<[3,?]:f32>, !numpy.ndarray<*:f32> + } +} + +// ----- +// CHECK-LABEL: @nonPublic +module @nonPublic { + // CHECK: @f(%arg0: !numpy.ndarray<[3,?]:f32>) -> !numpy.ndarray<[3,?]:f32> + func @f(%arg0: !numpy.ndarray<[3,?]:f32>) -> (!numpy.ndarray<[3,?]:f32>) + attributes { sym_visibility = "private" } { + return %arg0 : !numpy.ndarray<[3,?]:f32> + } +} + +// ----- +// CHECK-LABEL: @called +module @called { + // CHECK: @f(%arg0: !numpy.ndarray<*:f32>) -> !numpy.ndarray<*:f32> + // expected-warning @+1 {{unimplemented: cannot convert}} + func @f(%arg0: !numpy.ndarray<*:f32>) -> !numpy.ndarray<*:f32> { + return %arg0 : !numpy.ndarray<*:f32> + } + + func @caller(%arg0: !numpy.ndarray<*:f32>) -> !numpy.ndarray<*:f32> + attributes { sym_visibility = "private" } { + %0 = call @f(%arg0) : (!numpy.ndarray<*:f32>) -> !numpy.ndarray<*:f32> + return %0 : !numpy.ndarray<*:f32> + } +}