diff --git a/mlir/include/mlir/Transforms/Bufferize.h b/mlir/include/mlir/Transforms/Bufferize.h --- a/mlir/include/mlir/Transforms/Bufferize.h +++ b/mlir/include/mlir/Transforms/Bufferize.h @@ -44,12 +44,6 @@ /// except for the ranked-tensor types which is converted to memref types. class BufferizeTypeConverter : public TypeConverter { public: - /// This enum is for showing how buffer placement operation converters should - /// conduct with certain result type after type conversion. This value can be - /// set/get for each specific type using setResultConversionKind or - /// getResultConversionKind. - enum ResultConversionKind { AppendToArgumentsList, KeepAsFunctionResult }; - BufferizeTypeConverter(); /// This method tries to decompose a value of a certain type using provided @@ -82,26 +76,6 @@ addConversion(std::forward(callback)); } - /// This method returns ResultConversionKind for the mapping from `origin` - /// type to `input` type. - ResultConversionKind getResultConversionKind(Type origin, Type input); - - /// This method registers ResultConversionKind for the mapping from type 'T' - /// to type 'U'. - template - void setResultConversionKind(ResultConversionKind kind) { - assert((kind != AppendToArgumentsList || - llvm::is_one_of::value) && - "Only the memref typed values can be set to be appended to the " - "function argument list at the moment"); - resultTypeConversions.emplace_back( - [=](Type origin, Type input) -> Optional { - if (origin.template isa() && input.template isa()) - return kind; - return llvm::None; - }); - } - private: using DecomposeValueConversionCallFn = std::function( OpBuilder &, Location, Type, Value, SmallVectorImpl &)>; @@ -109,9 +83,6 @@ using DecomposeTypeConversionCallFn = std::function(Type, SmallVectorImpl &)>; - using ResultConversionKindFn = - std::function(Type, Type)>; - /// Generate a wrapper for the given decompose value conversion callback. template DecomposeValueConversionCallFn @@ -139,7 +110,6 @@ }; } - SmallVector resultTypeConversions; SmallVector decomposeValueConversions; SmallVector decomposeTypeConversions; }; @@ -221,48 +191,10 @@ LogicalResult matchAndRewrite(ReturnOpSourceTy returnOp, ArrayRef operands, ConversionPatternRewriter &rewriter) const final { - Location loc = returnOp.getLoc(); - - // Split the operands depending on whether they need a copy operation or - // they remain as operands of the return operation. If an operand is - // decomposable and a decompose callback function has been provided by the - // user, it will be unpacked. - SmallVector newOperands, needCopyOperands; - OpBuilder builder(returnOp); - for (auto operand : llvm::enumerate(operands)) { - SmallVector values; - this->converter.tryDecomposeValue(builder, loc, operand.value().getType(), - operand.value(), values); - Type type = returnOp.getOperand(operand.index()).getType(); - SmallVector originTypes; - this->converter.tryDecomposeType(type, originTypes); - for (auto value : llvm::enumerate(values)) { - Type origin = originTypes[value.index()]; - Type converted = value.value().getType(); - auto kind = this->converter.getResultConversionKind(origin, converted); - if (kind == BufferizeTypeConverter::KeepAsFunctionResult) - newOperands.push_back(value.value()); - else - // kind = BufferizeTypeConverter::AppendToArgumentsList - needCopyOperands.push_back(value.value()); - } - } - - // Insert Copy operations instead for the operands that have been removed - // from operand list and appended to the function arguments list. - Block &entryBlock = returnOp.getParentRegion()->front(); - unsigned numFuncArgs = entryBlock.getNumArguments(); - if (needCopyOperands.size() > numFuncArgs) - return returnOp.emitError( - "The number of operands that need Copy operations is more " - "than the number of target function arguments."); - unsigned destArgNum = numFuncArgs - needCopyOperands.size(); - rewriter.setInsertionPoint(returnOp); - for (Value operand : needCopyOperands) { - rewriter.create(loc, operand, - entryBlock.getArgument(destArgNum)); - ++destArgNum; - } + SmallVector newOperands; + for (auto operand : operands) + this->converter.tryDecomposeValue( + rewriter, returnOp.getLoc(), operand.getType(), operand, newOperands); rewriter.replaceOpWithNewOp(returnOp, newOperands); return success(); } diff --git a/mlir/lib/Transforms/Bufferize.cpp b/mlir/lib/Transforms/Bufferize.cpp --- a/mlir/lib/Transforms/Bufferize.cpp +++ b/mlir/lib/Transforms/Bufferize.cpp @@ -63,15 +63,6 @@ types.push_back(type); } -/// This method returns ResultConversionKind for the input type. -BufferizeTypeConverter::ResultConversionKind -BufferizeTypeConverter::getResultConversionKind(Type origin, Type converted) { - for (auto &conversion : resultTypeConversions) - if (auto res = conversion(origin, converted)) - return res.getValue(); - return KeepAsFunctionResult; -} - void mlir::populateBufferizeMaterializationLegality(ConversionTarget &target) { target.addLegalOp(); } @@ -140,16 +131,8 @@ for (Type resultType : funcType.getResults()) { SmallVector originTypes; converter.tryDecomposeType(resultType, originTypes); - for (auto origin : originTypes) { - Type converted = converter.convertType(origin); - auto kind = converter.getResultConversionKind(origin, converted); - if (kind == BufferizeTypeConverter::AppendToArgumentsList) { - conversion.addInputs(converted); - } else { - assert(kind == BufferizeTypeConverter::KeepAsFunctionResult); - newResultTypes.push_back(converted); - } - } + for (auto origin : originTypes) + newResultTypes.push_back(converter.convertType(origin)); } if (failed(rewriter.convertRegionTypes(&funcOp.getBody(), converter, @@ -168,66 +151,12 @@ // BufferizeCallOpConverter //===----------------------------------------------------------------------===// -namespace { -// This class represents a mapping from a result to a list of values and some -// results that have not yet constructed. Instead, the indices of these -// results in the operation that will be constructed are known. They will be -// replaced with the actual values when they are available. The order of -// adding to this mapping is important. -class CallOpResultMapping { -public: - CallOpResultMapping() { order = 0; }; - - /// Add an available value to the mapping. - void addMapping(Value value) { toValuesMapping.push_back({order++, value}); } - - /// Add the index of unavailble result value to the mapping. - void addMapping(unsigned index) { - toIndicesMapping.push_back({order++, index}); - } - - /// This method returns the mapping values list. The unknown result values - /// that only their indices are available are replaced with their values. - void getMappingValues(ValueRange valuesToReplaceIndices, - SmallVectorImpl &values) { - // Append available values to the list. - SmallVector, 2> res(toValuesMapping.begin(), - toValuesMapping.end()); - // Replace the indices with the actual values. - for (const std::pair &entry : toIndicesMapping) { - assert(entry.second < valuesToReplaceIndices.size() && - "The value index is out of range."); - res.push_back({entry.first, valuesToReplaceIndices[entry.second]}); - } - // Sort the values based on their adding orders. - llvm::sort(res, [](const std::pair &v1, - const std::pair &v2) { - return v1.first < v2.first; - }); - // Fill the values. - for (const std::pair &entry : res) - values.push_back(entry.second); - } - -private: - /// Keeping the inserting order of mapping values. - int order; - - /// Containing the mapping values with their inserting orders. - SmallVector, 2> toValuesMapping; - - /// Containing the indices of result values with their inserting orders. - SmallVector, 2> toIndicesMapping; -}; -} // namespace - /// Performs the actual rewriting step. LogicalResult BufferizeCallOpConverter::matchAndRewrite( CallOp callOp, ArrayRef operands, ConversionPatternRewriter &rewriter) const { Location loc = callOp.getLoc(); - OpBuilder builder(callOp); SmallVector newOperands; // TODO: if the CallOp references a FuncOp that only has a declaration (e.g. @@ -237,39 +166,25 @@ // Create the operands list of the new `CallOp`. It unpacks the decomposable // values if a decompose callback function has been provided by the user. - for (auto operand : operands) { - SmallVector values; - converter.tryDecomposeValue(builder, loc, operand.getType(), operand, - values); - newOperands.append(values.begin(), values.end()); - } + for (auto operand : operands) + converter.tryDecomposeValue(rewriter, loc, operand.getType(), operand, + newOperands); - // Create the new result types for the new `CallOp` and a mapping from the old - // result to new value(s). + // Create the new result types for the new `CallOp` and track the indices in + // the new call op's results that correspond to the old call op's results. SmallVector newResultTypes; - SmallVector mappings; - mappings.resize(callOp.getNumResults()); + SmallVector, 4> expandedResultIndices; + expandedResultIndices.resize(callOp.getNumResults()); for (auto result : llvm::enumerate(callOp.getResults())) { SmallVector originTypes; converter.tryDecomposeType(result.value().getType(), originTypes); - auto &resultMapping = mappings[result.index()]; + auto &resultMapping = expandedResultIndices[result.index()]; for (Type origin : originTypes) { Type converted = converter.convertType(origin); - auto kind = converter.getResultConversionKind(origin, converted); - if (kind == BufferizeTypeConverter::KeepAsFunctionResult) { - newResultTypes.push_back(converted); - // The result value is not yet available. Its index is kept and it is - // replaced with the actual value of the new `CallOp` later. - resultMapping.addMapping(newResultTypes.size() - 1); - } else { - // kind = BufferizeTypeConverter::AppendToArgumentsList - MemRefType memref = converted.dyn_cast(); - if (!memref) - return callOp.emitError("Cannot allocate for a non-Memref type"); - Value alloc = rewriter.create(loc, memref); - newOperands.push_back(alloc); - resultMapping.addMapping(alloc); - } + newResultTypes.push_back(converted); + // The result value is not yet available. Its index is kept and it is + // replaced with the actual value of the new `CallOp` later. + resultMapping.push_back(newResultTypes.size() - 1); } } @@ -278,12 +193,12 @@ // Build a replacing value for each result to replace its uses. If a result // has multiple mapping values, it needs to be packed to a single value. - OpBuilder nextBuilder(callOp.getOperation()->getNextNode()); SmallVector replacedValues; replacedValues.reserve(callOp.getNumResults()); for (unsigned i = 0, e = callOp.getNumResults(); i < e; ++i) { - SmallVector valuesToPack; - mappings[i].getMappingValues(newCallOp.getResults(), valuesToPack); + auto valuesToPack = llvm::to_vector<6>( + llvm::map_range(expandedResultIndices[i], + [&](int i) { return newCallOp.getResult(i); })); if (valuesToPack.empty()) { // No replacement is required. replacedValues.push_back(nullptr); @@ -293,7 +208,7 @@ // Values need to be packed using callback function. The same callback // that is used for materializeArgumentConversion is used for packing. Value packed = converter.materializeArgumentConversion( - nextBuilder, loc, callOp.getType(i), valuesToPack); + rewriter, loc, callOp.getType(i), valuesToPack); replacedValues.push_back(packed); } } diff --git a/mlir/test/Transforms/finalizing-bufferize-allowed-memref-results.mlir b/mlir/test/Transforms/finalizing-bufferize-allowed-memref-results.mlir deleted file mode 100644 --- a/mlir/test/Transforms/finalizing-bufferize-allowed-memref-results.mlir +++ /dev/null @@ -1,184 +0,0 @@ -// RUN: mlir-opt -test-finalizing-bufferize-with-allowed-memref-results -split-input-file %s | FileCheck %s - -// Since allowMemrefEscaping is active for Bufferization in this test pass, -// all tensor typed function results are converted to memref and remain as -// function results. All memref typed function results will escape from the -// deallocation phase of Bufferization. - -// CHECK-LABEL: func @void_function_signature_conversion -func @void_function_signature_conversion(%arg0: tensor<4x8xf32>) { - return -} -// CHECK: ({{.*}}: memref<4x8xf32>) - -// ----- - -// CHECK-LABEL: func @complex_signature_conversion -func @complex_signature_conversion( - %arg0: tensor<5xf32>, - %arg1: memref<10xf32>, - %arg2: i1, %arg3: f16) -> ( - i1, - tensor<5xf32>, - memref<10xf32>, - memref<15xf32>, - f16) { - %0 = alloc() : memref<15xf32> - %1 = test.tensor_based in(%arg0 : tensor<5xf32>) -> tensor<5xf32> - return %arg2, %1, %arg1, %0, %arg3 : - i1, tensor<5xf32>, memref<10xf32>, memref<15xf32>, f16 -} -// CHECK: (%[[ARG0:.*]]: memref<5xf32>, %[[ARG1:.*]]: memref<10xf32>, -// CHECK-SAME: %[[ARG2:.*]]: i1, %[[ARG3:.*]]: f16) -// CHECK-SAME: (i1, memref<5xf32>, memref<10xf32>, memref<15xf32>, f16) -// CHECK: %[[FIRST_ALLOC:.*]] = alloc() -// CHECK: %[[TENSOR_ALLOC:.*]] = alloc() -// CHECK: return %[[ARG2]], %[[TENSOR_ALLOC]], %[[ARG1]], %[[FIRST_ALLOC]], -// CHECK-SAME: %[[ARG3]] - -// ----- - -// CHECK-LABEL: func @no_signature_conversion_is_needed -func @no_signature_conversion_is_needed(%arg0: memref<4x8xf32>) { - return -} -// CHECK: ({{.*}}: memref<4x8xf32>) - -// ----- - -// CHECK-LABEL: func @no_signature_conversion_is_needed -func @no_signature_conversion_is_needed(%arg0: i1, %arg1: f16) -> (i1, f16){ - return %arg0, %arg1 : i1, f16 -} -// CHECK: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: f16) -> (i1, f16) -// CHECK: return %[[ARG0]], %[[ARG1]] - -// ----- - -// CHECK-LABEL: func @simple_signature_conversion -func @simple_signature_conversion(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> { - return %arg0 : tensor<4x8xf32> -} -// CHECK: (%[[ARG0:.*]]: [[TYPE:.*]]<[[RANK:.*]]>) -> [[TYPE]]<[[RANK]]> -// CHECK-NEXT: return %[[ARG0]] - -// ----- - -// CHECK-LABEL: func @func_with_unranked_arg_and_result -func @func_with_unranked_arg_and_result(%arg0: tensor<*xf32>) -> tensor<*xf32> { - return %arg0 : tensor<*xf32> -} -// CHECK-SAME: ([[ARG:%.*]]: memref<*xf32>) -> memref<*xf32> -// CHECK-NEXT: return [[ARG]] : memref<*xf32> - -// ----- - -// CHECK-LABEL: func @func_and_block_signature_conversion -func @func_and_block_signature_conversion(%arg0 : tensor<2xf32>, %cond : i1, %arg1: tensor<4x4xf32>) -> tensor<4x4xf32>{ - cond_br %cond, ^bb1, ^bb2 - ^bb1: - br ^exit(%arg0 : tensor<2xf32>) - ^bb2: - br ^exit(%arg0 : tensor<2xf32>) - ^exit(%arg2: tensor<2xf32>): - return %arg1 : tensor<4x4xf32> -} -// CHECK: (%[[ARG0:.*]]: [[ARG0_TYPE:.*]], %[[COND:.*]]: i1, %[[ARG1:.*]]: [[ARG1_TYPE:.*]]) -> [[RESULT_TYPE:.*]] -// CHECK: br ^[[EXIT_BLOCK:.*]](%[[ARG0]] : [[ARG0_TYPE]]) -// CHECK: br ^[[EXIT_BLOCK]](%[[ARG0]] : [[ARG0_TYPE]]) -// CHECK: ^[[EXIT_BLOCK]](%{{.*}}: [[ARG0_TYPE]]) -// CHECK-NEXT: return %[[ARG1]] - -// ----- - -// CHECK-LABEL: func @callee -func @callee(%arg1: tensor<5xf32>) -> (tensor<5xf32>, memref<2xf32>) { - %buff = alloc() : memref<2xf32> - return %arg1, %buff : tensor<5xf32>, memref<2xf32> -} -// CHECK: (%[[CALLEE_ARG:.*]]: memref<5xf32>) -> (memref<5xf32>, memref<2xf32>) -// CHECK: %[[ALLOC:.*]] = alloc() -// CHECK: return %[[CALLEE_ARG]], %[[ALLOC]] - -// CHECK-LABEL: func @caller -func @caller(%arg0: tensor<5xf32>) -> tensor<5xf32> { - %x:2 = call @callee(%arg0) : (tensor<5xf32>) -> (tensor<5xf32>, memref<2xf32>) - %y:2 = call @callee(%x#0) : (tensor<5xf32>) -> (tensor<5xf32>, memref<2xf32>) - return %y#0 : tensor<5xf32> -} -// CHECK: (%[[CALLER_ARG:.*]]: memref<5xf32>) -> memref<5xf32> -// CHECK: %[[X:.*]]:2 = call @callee(%[[CALLER_ARG]]) -// CHECK: %[[Y:.*]]:2 = call @callee(%[[X]]#0) -// CHECK: return %[[Y]]#0 - -// ----- - -// Test case: Testing BufferizeCallOpConverter to see if it matches with the -// signature of the new signature of the callee function when there are tuple -// typed args and results. BufferizeTypeConverter is set to flatten tuple typed -// arguments. The tuple typed values should be decomposed and composed using -// get_tuple_element and make_tuple operations of test dialect. Tensor types are -// converted to Memref. Memref typed function results remain as function -// results. - -// CHECK-LABEL: func @callee -func @callee(%arg0: tuple,i1, tensor<5xf32>>) -> (tuple,i1, tensor<5xf32>>){ - return %arg0 : tuple,i1, tensor<5xf32>> -} -// CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[ARG1:.*]]: i1, %[[ARG2:.*]]: memref<5xf32>) -// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>) -// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) -// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) {index = 0 : i32} -// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) {index = 1 : i32} -// CHECK-NEXT: %[[THIRD_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) {index = 2 : i32} -// CHECK-NEXT: return %[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]] - -// CHECK-LABEL: func @caller -func @caller(%arg0: tuple,i1, tensor<5xf32>>) -> tuple,i1, tensor<5xf32>>{ - %x0 = call @callee(%arg0) : (tuple,i1, tensor<5xf32>>) -> (tuple,i1, tensor<5xf32>>) - %y0 = call @callee(%x0) : (tuple,i1, tensor<5xf32>>) -> (tuple,i1, tensor<5xf32>>) - return %y0 : tuple,i1, tensor<5xf32>> -} -// CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[ARG1:.*]]: i1, %[[ARG2:.*]]: memref<5xf32>) -// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>) -// CHECK-NEXT: %[[ARG_TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) -// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 0 : i32} -// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 1 : i32} -// CHECK-NEXT: %[[THIRD_ELEM:.*]] = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 2 : i32} -// CHECK-NEXT: %[[CALLEE_RESULTS:.*]]:3 = call @callee(%[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]]) -// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>) -> (memref<2xf32>, i1, memref<5xf32>) -// CHECK-NEXT: %[[RESULT_TUPLE:.*]] = "test.make_tuple"(%[[CALLEE_RESULTS]]#0, %[[CALLEE_RESULTS]]#1, %[[CALLEE_RESULTS]]#2) -// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[RESULT_TUPLE]]) {index = 0 : i32} -// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[RESULT_TUPLE]]) {index = 1 : i32} -// CHECK-NEXT: %[[THIRD_ELEM:.*]] = "test.get_tuple_element"(%[[RESULT_TUPLE]]) {index = 2 : i32} -// CHECK-NEXT: %[[CALLEE_RESULTS:.*]]:3 = call @callee(%[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]]) -// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>) -> (memref<2xf32>, i1, memref<5xf32>) -// CHECK-NEXT: %[[RETURN_TUPLE:.*]] = "test.make_tuple"(%[[CALLEE_RESULTS]]#0, %[[CALLEE_RESULTS]]#1, %[[CALLEE_RESULTS]]#2) -// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[RETURN_TUPLE]]) {index = 0 : i32} -// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[RETURN_TUPLE]]) {index = 1 : i32} -// CHECK-NEXT: %[[THIRD_ELEM:.*]] = "test.get_tuple_element"(%[[RETURN_TUPLE]]) {index = 2 : i32} -// CHECK-NEXT: return %[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]] - -// ----- - -// Test case: Testing BufferizeFuncOpConverter and -// BufferizeReturnOpConverter to see if the return operation matches with the -// new function signature when there are tuple typed args and results. -// BufferizeTypeConverter is set to flatten tuple typed arguments. The tuple -// typed values should be decomposed and composed using get_tuple_element and -// make_tuple operations of test dialect. Tensor types are converted to Memref. -// Memref typed function results remain as function results. - -// CHECK-LABEL: func @decompose_tuple_typed_function_args_and_results -func @decompose_tuple_typed_function_args_and_results(%arg0: tuple, %arg1: tensor<10xf32>, %arg2: tuple>) -> (tuple>, tensor<10xf32>, tuple){ - return %arg2, %arg1, %arg0 : tuple>, tensor<10xf32>, tuple -} -// CHECK-SAME: %[[ARG0:.*]]: i1, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: memref<10xf32>, %[[ARG3:.*]]: i1, %[[ARG4:.*]]: memref<5xf32> -// CHECK-SAME: (i1, memref<5xf32>, memref<10xf32>, i1, f32) -// CHECK-NEXT: %[[FIRST_TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]]) -// CHECK-NEXT: %[[SECOND_TUPLE:.*]] = "test.make_tuple"(%[[ARG3]], %[[ARG4]]) -// CHECK-NEXT: %[[SECOND_TUPLE_FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[SECOND_TUPLE]]) {index = 0 : i32} -// CHECK-NEXT: %[[SECOND_TUPLE_SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[SECOND_TUPLE]]) {index = 1 : i32} -// CHECK-NEXT: %[[FIRST_TUPLE_FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[FIRST_TUPLE]]) {index = 0 : i32} -// CHECK-NEXT: %[[FIRST_TUPLE_SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[FIRST_TUPLE]]) {index = 1 : i32} -// CHECK-NEXT: return %[[SECOND_TUPLE_FIRST_ELEM]], %[[SECOND_TUPLE_SECOND_ELEM]], %[[ARG2]], %[[FIRST_TUPLE_FIRST_ELEM]], %[[FIRST_TUPLE_SECOND_ELEM]] diff --git a/mlir/test/Transforms/finalizing-bufferize.mlir b/mlir/test/Transforms/finalizing-bufferize.mlir --- a/mlir/test/Transforms/finalizing-bufferize.mlir +++ b/mlir/test/Transforms/finalizing-bufferize.mlir @@ -1,31 +1,36 @@ // RUN: mlir-opt -test-finalizing-bufferize -split-input-file %s | FileCheck %s -// CHECK-LABEL: func @func_signature_conversion -func @func_signature_conversion(%arg0: tensor<4x8xf32>) { +// CHECK-LABEL: func @void_function_signature_conversion +func @void_function_signature_conversion(%arg0: tensor<4x8xf32>) { return } -// CHECK: ({{.*}}: memref<4x8xf32>) { +// CHECK: ({{.*}}: memref<4x8xf32>) // ----- -// Only tensor typed function result should be converted to memref and move to -// the function arguments list. The other memref function results remain as -// function results. - -// CHECK-LABEL: func @memref_in_function_results -func @memref_in_function_results(%arg0: tensor<5xf32>, %arg1: memref<10xf32>) - -> (tensor<5xf32>, memref<10xf32>, memref<15xf32>) { +// CHECK-LABEL: func @complex_signature_conversion +func @complex_signature_conversion( + %arg0: tensor<5xf32>, + %arg1: memref<10xf32>, + %arg2: i1, + %arg3: f16) -> ( + i1, + tensor<5xf32>, + memref<10xf32>, + memref<15xf32>, + f16) { %0 = alloc() : memref<15xf32> %1 = test.tensor_based in(%arg0 : tensor<5xf32>) -> tensor<5xf32> - return %1, %arg1, %0 : tensor<5xf32>, memref<10xf32>, memref<15xf32> + return %arg2, %1, %arg1, %0, %arg3 : + i1, tensor<5xf32>, memref<10xf32>, memref<15xf32>, f16 } // CHECK: (%[[ARG0:.*]]: memref<5xf32>, %[[ARG1:.*]]: memref<10xf32>, -// CHECK-SAME: %[[RESULT:.*]]: memref<5xf32>) -// CHECK-SAME: (memref<10xf32>, memref<15xf32>) +// CHECK-SAME: %[[ARG2:.*]]: i1, %[[ARG3:.*]]: f16) +// CHECK-SAME: (i1, memref<5xf32>, memref<10xf32>, memref<15xf32>, f16) // CHECK: %[[FIRST_ALLOC:.*]] = alloc() // CHECK: %[[TENSOR_ALLOC:.*]] = alloc() -// CHECK: test.copy(%[[TENSOR_ALLOC]], %[[RESULT]]) -// CHECK: return %[[ARG1]], %[[FIRST_ALLOC]] +// CHECK: return %[[ARG2]], %[[TENSOR_ALLOC]], %[[ARG1]], %[[FIRST_ALLOC]], +// CHECK-SAME: %[[ARG3]] // ----- @@ -33,7 +38,7 @@ func @no_signature_conversion_is_needed(%arg0: memref<4x8xf32>) { return } -// CHECK: ({{.*}}: memref<4x8xf32>) { +// CHECK: ({{.*}}: memref<4x8xf32>) // ----- @@ -46,39 +51,26 @@ // ----- -// CHECK-LABEL: func @complex_signature_conversion -func @complex_signature_conversion(%arg0: tensor<4x8xf32>, %arg1: i1, - %arg2: tensor<5x5xf64>,%arg3: f16) -> - (i1, tensor<5x5xf64>, f16, tensor<4x8xf32>) { - return %arg1, %arg2, %arg3, %arg0 : i1, tensor<5x5xf64>, f16, - tensor<4x8xf32> +// CHECK-LABEL: func @simple_signature_conversion +func @simple_signature_conversion(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> { + return %arg0 : tensor<4x8xf32> } -// CHECK: (%[[ARG0:.*]]: memref<4x8xf32>, %[[ARG1:.*]]: i1 -// CHECK-SAME: %[[ARG2:.*]]: memref<5x5xf64>, %[[ARG3:.*]]: f16 -// CHECK-SAME: %[[RESULT1:.*]]: memref<5x5xf64> -// CHECK-SAME: %[[RESULT2:.*]]: memref<4x8xf32>) -> (i1, f16) { -// CHECK-NEXT: test.copy(%[[ARG2]], %[[RESULT1]]) -// CHECK-NEXT: test.copy(%[[ARG0]], %[[RESULT2]]) -// CHECK-NEXT: return %[[ARG1]], %[[ARG3]] +// CHECK: (%[[ARG0:.*]]: [[TYPE:.*]]<[[RANK:.*]]>) -> [[TYPE]]<[[RANK]]> +// CHECK-NEXT: return %[[ARG0]] // ----- -// CHECK-LABEL: func @non_void_to_void_return_op_converter -func @non_void_to_void_return_op_converter(%arg0: tensor<4x8xf32>) - -> tensor<4x8xf32> { - return %arg0 : tensor<4x8xf32> +// CHECK-LABEL: func @func_with_unranked_arg_and_result +func @func_with_unranked_arg_and_result(%arg0: tensor<*xf32>) -> tensor<*xf32> { + return %arg0 : tensor<*xf32> } -// CHECK: (%[[ARG0:.*]]: [[TYPE:.*]]<[[RANK:.*]]>, -// CHECK-SAME: %[[RESULT:.*]]: [[TYPE]]<[[RANK]]>) { -// CHECK-NEXT: test.copy(%[[ARG0]], %[[RESULT]]) -// CHECK-NEXT: return +// CHECK-SAME: ([[ARG:%.*]]: memref<*xf32>) -> memref<*xf32> +// CHECK-NEXT: return [[ARG]] : memref<*xf32> // ----- // CHECK-LABEL: func @func_and_block_signature_conversion -func @func_and_block_signature_conversion(%arg0 : tensor<2xf32>, %cond : i1, - %arg1: tensor<4x4xf32>) - -> tensor<4x4xf32>{ +func @func_and_block_signature_conversion(%arg0 : tensor<2xf32>, %cond : i1, %arg1: tensor<4x4xf32>) -> tensor<4x4xf32>{ cond_br %cond, ^bb1, ^bb2 ^bb1: br ^exit(%arg0 : tensor<2xf32>) @@ -87,293 +79,102 @@ ^exit(%arg2: tensor<2xf32>): return %arg1 : tensor<4x4xf32> } -// CHECK: (%[[ARG0:.*]]: [[ARG0_TYPE:.*]], %[[COND:.*]]: i1, -// CHECK-SAME: %[[ARG1:.*]]: [[ARG1_TYPE:.*]], -// CHECK-SAME: %[[RESULT:.*]]: [[RESULT_TYPE:.*]]) { +// CHECK: (%[[ARG0:.*]]: [[ARG0_TYPE:.*]], %[[COND:.*]]: i1, %[[ARG1:.*]]: [[ARG1_TYPE:.*]]) -> [[RESULT_TYPE:.*]] { // CHECK: br ^[[EXIT_BLOCK:.*]](%[[ARG0]] : [[ARG0_TYPE]]) // CHECK: br ^[[EXIT_BLOCK]](%[[ARG0]] : [[ARG0_TYPE]]) // CHECK: ^[[EXIT_BLOCK]](%{{.*}}: [[ARG0_TYPE]]) -// CHECK-NEXT: test.copy(%[[ARG1]], %[[RESULT]]) -// CHECK-NEXT: return - -// ----- - -// Test Case: Simple case for checking if BufferizePlacer creates AllocOps -// right before TensorBasedOp. - -// CHECK-LABEL: func @compute_allocs_position_simple -func @compute_allocs_position_simple(%cond: i1, %arg0: tensor<2xf32>) - -> tensor<2xf32>{ - %0 = test.tensor_based in(%arg0 : tensor<2xf32>) -> tensor<2xf32> - %1 = test.tensor_based in(%0 : tensor<2xf32>) -> tensor<2xf32> - return %1 : tensor<2xf32> -} -// CHECK: (%{{.*}}: {{.*}}, %[[ARG0:.*]]: memref<2xf32>, -// CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc() -// CHECK-NEXT: test.buffer_based in(%[[ARG0]]{{.*}} out(%[[FIRST_ALLOC]] -// CHECK: %[[SECOND_ALLOC:.*]] = alloc() -// CHECK-NEXT: test.buffer_based in(%[[FIRST_ALLOC]]{{.*}} out(%[[SECOND_ALLOC]] +// CHECK-NEXT: return %[[ARG1]] : [[RESULT_TYPE]] // ----- -// Test Case: if-else case for checking if BufferizePlacer creates AllocOps -// right before TensorBasedOp. - -// CHECK-LABEL: func @compute_allocs_position -func @compute_allocs_position(%cond: i1, %arg0: tensor<2xf32>) -> tensor<2xf32>{ - %0 = test.tensor_based in(%arg0 : tensor<2xf32>) -> tensor<2xf32> - %1 = test.tensor_based in(%0 : tensor<2xf32>) -> tensor<2xf32> - cond_br %cond, ^bb1(%arg0, %0: tensor<2xf32>, tensor<2xf32>), - ^bb2(%0, %arg0: tensor<2xf32>, tensor<2xf32>) - ^bb1(%arg1 : tensor<2xf32>, %arg2 : tensor<2xf32>): - %2 = test.tensor_based in(%arg0 : tensor<2xf32>) -> tensor<2xf32> - %3 = test.tensor_based in(%2 : tensor<2xf32>) -> tensor<2xf32> - br ^exit(%arg1, %arg2 : tensor<2xf32>, tensor<2xf32>) - ^bb2(%arg3 : tensor<2xf32>, %arg4 : tensor<2xf32>): - %4 = test.tensor_based in(%arg0 : tensor<2xf32>) -> tensor<2xf32> - %5 = test.tensor_based in(%4 : tensor<2xf32>) -> tensor<2xf32> - br ^exit(%arg3, %arg4 : tensor<2xf32>, tensor<2xf32>) - ^exit(%arg5 : tensor<2xf32>, %arg6 : tensor<2xf32>): - %6 = test.tensor_based in(%arg0 : tensor<2xf32>) -> tensor<2xf32> - %7 = test.tensor_based in(%6 : tensor<2xf32>) -> tensor<2xf32> - return %7 : tensor<2xf32> -} -// CHECK: (%{{.*}}: {{.*}}, %[[ARG0:.*]]: memref<2xf32>, -// CHECK-NEXT: %[[ALLOC0:.*]] = alloc() -// CHECK-NEXT: test.buffer_based in(%[[ARG0]]{{.*}} out(%[[ALLOC0]] -// CHECK: %[[ALLOC1:.*]] = alloc() -// CHECK-NEXT: test.buffer_based in(%[[ALLOC0]]{{.*}} out(%[[ALLOC1]] -// CHECK: cond_br %{{.*}}, ^[[BB0:.*]]({{.*}}), ^[[BB1:.*]]( -// CHECK-NEXT: ^[[BB0]] -// CHECK-NEXT: %[[ALLOC2:.*]] = alloc() -// CHECK-NEXT: test.buffer_based in(%[[ARG0]]{{.*}} out(%[[ALLOC2]] -// CHECK: %[[ALLOC3:.*]] = alloc() -// CHECK-NEXT: test.buffer_based in(%[[ALLOC2]]{{.*}} out(%[[ALLOC3]] -// CHECK: br ^[[EXIT:.*]]({{.*}}) -// CHECK-NEXT: ^[[BB1]] -// CHECK-NEXT: %[[ALLOC4:.*]] = alloc() -// CHECK-NEXT: test.buffer_based in(%[[ARG0]]{{.*}} out(%[[ALLOC4]] -// CHECK: %[[ALLOC5:.*]] = alloc() -// CHECK-NEXT: test.buffer_based in(%[[ALLOC4]]{{.*}} out(%[[ALLOC5]] -// CHECK: br ^[[EXIT]] -// CHECK-NEXT: ^[[EXIT]] -// CHECK-NEXT: %[[ALLOC6:.*]] = alloc() -// CHECK-NEXT: test.buffer_based in(%[[ARG0]]{{.*}} out(%[[ALLOC6]] -// CHECK: %[[ALLOC7:.*]] = alloc() -// CHECK-NEXT: test.buffer_based in(%[[ALLOC6]]{{.*}} out(%[[ALLOC7]] - -// ----- - -// Test case: Checking BufferizeCallOpConverter and -// BufferizeFuncOpConverter and BufferizeReturnOpConverter all -// together. The signature of `callee` after signature conversion would be: - -// func @callee(%arg0: memref<5xf32>,%arg1: memref<5xf32>) -> () - -// The operands and results of caller and return operations must be matched -// respectively. - -// CHECK-LABEL: func @callee -func @callee(%arg1: tensor<5xf32>) -> tensor<5xf32> { - %0 = test.tensor_based in(%arg1 : tensor<5xf32>) -> tensor<5xf32> - return %0 : tensor<5xf32> -} -// CHECK: (%[[CALLEE_ARG:.*]]: memref<5xf32>, -// CHECK-SAME: %[[CALLEE_RESULT:.*]]: memref<5xf32>) -// CHECK: %[[ALLOC:.*]] = alloc() -// CHECK: test.buffer_based -// CHECK: test.copy(%[[ALLOC]], %[[CALLEE_RESULT]]) -// CHECK: return - -// CHECK-LABEL: func @caller -func @caller(%arg0: tensor<5xf32>) -> tensor<5xf32> { - %x = call @callee(%arg0) : (tensor<5xf32>) -> tensor<5xf32> - %y = call @callee(%x) : (tensor<5xf32>) -> tensor<5xf32> - return %y : tensor<5xf32> -} -// CHECK: (%[[CALLER_ARG:.*]]: memref<5xf32>, -// CHECK-SAME: %[[CALLER_RESULT:.*]]: memref<5xf32>) -// CHECK: %[[FIRST_ALLOC:.*]] = alloc() -// CHECK: call @callee(%[[CALLER_ARG]], %[[FIRST_ALLOC]]) -// CHECK: %[[SECOND_ALLOC:.*]] = alloc() -// CHECK: call @callee(%[[FIRST_ALLOC]], %[[SECOND_ALLOC]]) -// CHECK: test.copy(%[[SECOND_ALLOC]], %[[CALLER_RESULT]]) -// CHECK: return - -// ----- - -// Test case: Checking BufferizeCallOpConverter and -// BufferizeFuncOpConverter and BufferizeReturnOpConverter all -// together on functions that also have memref typed results. The signature of -// `callee` after signature conversion would be: - -// func @callee(%arg0: memref<5xf32>,%arg1: memref<5xf32>)-> memref<2xf32> - -// where %arg0 is the input and %arg1 is the output buffer and the original -// memref type result remain as the function result. Then, the rewriter should -// match the caller's signature with the callee. Thus, two buffers will be -// allocated instead of %x0 and %y0 and they are passed to the callers' operands -// list as the output buffers. %x1 and %y1 remain as callers' results. - // CHECK-LABEL: func @callee func @callee(%arg1: tensor<5xf32>) -> (tensor<5xf32>, memref<2xf32>) { %buff = alloc() : memref<2xf32> return %arg1, %buff : tensor<5xf32>, memref<2xf32> } -// CHECK: (%[[CALLEE_ARG:.*]]: memref<5xf32>, -// CHECK-SAME: %[[CALLEE_RESULT:.*]]: memref<5xf32>) -// CHECK-SAME: memref<2xf32> -// CHECK: %[[ALLOC:.*]] = alloc() -// CHECK: test.copy(%[[CALLEE_ARG]], %[[CALLEE_RESULT]]) -// CHECK: return %[[ALLOC]] +// CHECK: (%[[CALLEE_ARG:.*]]: memref<5xf32>) -> (memref<5xf32>, memref<2xf32>) +// CHECK: %[[ALLOC:.*]] = alloc() +// CHECK: return %[[CALLEE_ARG]], %[[ALLOC]] // CHECK-LABEL: func @caller func @caller(%arg0: tensor<5xf32>) -> tensor<5xf32> { - %x0, %x1 = call @callee(%arg0) : (tensor<5xf32>) - -> (tensor<5xf32>, memref<2xf32>) - %y0, %y1 = call @callee(%x0) : (tensor<5xf32>) - -> (tensor<5xf32>, memref<2xf32>) - return %y0 : tensor<5xf32> -} -// CHECK: (%[[CALLER_ARG:.*]]: memref<5xf32>, -// CHECK-SAME: %[[CALLER_RESULT:.*]]: memref<5xf32>) -// CHECK: %[[X0:.*]] = alloc() -// CHECK: %[[X1:.*]] = call @callee(%[[CALLER_ARG]], %[[X0]]) -// CHECK: %[[Y0:.*]] = alloc() -// CHECK: %[[Y1:.*]] = call @callee(%[[X0]], %[[Y0]]) -// CHECK: test.copy(%[[Y0]], %[[CALLER_RESULT]]) -// CHECK: return - -// ----- - -// CHECK-LABEL: func @func_with_unranked_arg -func @func_with_unranked_arg(%arg0: tensor<*xf32>) { - return + %x:2 = call @callee(%arg0) : (tensor<5xf32>) -> (tensor<5xf32>, memref<2xf32>) + %y:2 = call @callee(%x#0) : (tensor<5xf32>) -> (tensor<5xf32>, memref<2xf32>) + return %y#0 : tensor<5xf32> } -// CHECK-SAME: ([[ARG:%.*]]: memref<*xf32>) +// CHECK: (%[[CALLER_ARG:.*]]: memref<5xf32>) -> memref<5xf32> +// CHECK: %[[X:.*]]:2 = call @callee(%[[CALLER_ARG]]) +// CHECK: %[[Y:.*]]:2 = call @callee(%[[X]]#0) +// CHECK: return %[[Y]]#0 // ----- // Test case: Testing BufferizeCallOpConverter to see if it matches with the // signature of the new signature of the callee function when there are tuple -// typed args and results. BufferizeTypeConverter is set to flatten tuple -// typed arguments. The tuple typed values should be decomposed and composed -// using get_tuple_element and make_tuple operations of test dialect. Tensor -// types are converted to Memref. Memref typed function results are appended to -// the function arguments list. +// typed args and results. BufferizeTypeConverter is set to flatten tuple typed +// arguments. The tuple typed values should be decomposed and composed using +// get_tuple_element and make_tuple operations of test dialect. Tensor types are +// converted to Memref. Memref typed function results remain as function +// results. // CHECK-LABEL: func @callee -func @callee(%arg0: tuple,i1, tensor<5xf32>>) - -> (tuple,i1, tensor<5xf32>>){ +func @callee(%arg0: tuple,i1, tensor<5xf32>>) -> (tuple,i1, tensor<5xf32>>){ return %arg0 : tuple,i1, tensor<5xf32>> } -// CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[ARG1:.*]]: i1, -// CHECK-SAME: %[[ARG2:.*]]: memref<5xf32>, %[[RESULT0:.*]]: memref<2xf32>, -// CHECK-SAME: %[[RESULT1:.*]]: memref<5xf32>) -> i1 -// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]], -// CHECK-SAME: %[[ARG2]]) -// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) -// CHECK-SAME: {index = 0 : i32} -// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) -// CHECK-SAME: {index = 1 : i32} -// CHECK-NEXT: %[[THIRD_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) -// CHECK-SAME: {index = 2 : i32} -// CHECK-NEXT: test.copy(%[[FIRST_ELEM]], %[[RESULT0]]) -// CHECK-NEXT: test.copy(%[[THIRD_ELEM]], %[[RESULT1]]) -// CHECK-NEXT: return %[[SECOND_ELEM]] +// CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[ARG1:.*]]: i1, %[[ARG2:.*]]: memref<5xf32>) +// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>) +// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) +// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) {index = 0 : i32} +// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) {index = 1 : i32} +// CHECK-NEXT: %[[THIRD_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) {index = 2 : i32} +// CHECK-NEXT: return %[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]] // CHECK-LABEL: func @caller -func @caller(%arg0: tuple,i1, tensor<5xf32>>) - -> tuple,i1, tensor<5xf32>>{ - %x0 = call @callee(%arg0) : (tuple,i1, tensor<5xf32>>) - -> (tuple,i1, tensor<5xf32>>) - %y0 = call @callee(%x0) : (tuple,i1, tensor<5xf32>>) - -> (tuple,i1, tensor<5xf32>>) +func @caller(%arg0: tuple,i1, tensor<5xf32>>) -> tuple,i1, tensor<5xf32>>{ + %x0 = call @callee(%arg0) : (tuple,i1, tensor<5xf32>>) -> (tuple,i1, tensor<5xf32>>) + %y0 = call @callee(%x0) : (tuple,i1, tensor<5xf32>>) -> (tuple,i1, tensor<5xf32>>) return %y0 : tuple,i1, tensor<5xf32>> } -// CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[ARG1:.*]]: i1, -// CHECK-SAME: %[[ARG2:.*]]: memref<5xf32>, %[[RESULT0:.*]]: memref<2xf32>, -// CHECK-SAME: %[[RESULT1:.*]]: memref<5xf32>) -> i1 -// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]], -// CHECK-SAME: %[[ARG2]]) -// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) -// CHECK-SAME: {index = 0 : i32} -// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) -// CHECK-SAME: {index = 1 : i32} -// CHECK-NEXT: %[[THIRD_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) -// CHECK-SAME: {index = 2 : i32} -// CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc() -// CHECK-NEXT: %[[SECOND_ALLOC:.*]] = alloc() -// CHECK-NEXT: %[[CALLEE_RESULT:.*]] = call @callee(%[[FIRST_ELEM]], -// CHECK-SAME: %[[SECOND_ELEM]], %[[THIRD_ELEM]], %[[FIRST_ALLOC]], -// CHECK-SAME: %[[SECOND_ALLOC]]) -// CHECK-SAME: (memref<2xf32>, i1, -// CHECK-SAME: memref<5xf32>, memref<2xf32>, memref<5xf32>) -> i1 -// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[FIRST_ALLOC]], -// CHECK-SAME: %[[CALLEE_RESULT]], %[[SECOND_ALLOC]]) -// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) -// CHECK-SAME: {index = 0 : i32} -// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) -// CHECK-SAME: {index = 1 : i32} -// CHECK-NEXT: %[[THIRD_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) -// CHECK-SAME: {index = 2 : i32} -// CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc() -// CHECK-NEXT: %[[SECOND_ALLOC:.*]] = alloc() -// CHECK-NEXT: %[[CALLEE_RESULT:.*]] = call @callee(%[[FIRST_ELEM]], -// CHECK-SAME: %[[SECOND_ELEM]], %[[THIRD_ELEM]], %[[FIRST_ALLOC]], -// CHECK-SAME: %[[SECOND_ALLOC]]) -// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>, memref<2xf32>, memref<5xf32>) -// CHECK-SAME: i1 -// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[FIRST_ALLOC]], -// CHECK-SAME: %[[CALLEE_RESULT]], %[[SECOND_ALLOC]]) -// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) -// CHECK-SAME: {index = 0 : i32} -// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) -// CHECK-SAME: {index = 1 : i32} -// CHECK-NEXT: %[[THIRD_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) -// CHECK-SAME: {index = 2 : i32} -// CHECK-NEXT: test.copy(%[[FIRST_ELEM]], %[[RESULT0]]) -// CHECK-NEXT: test.copy(%[[THIRD_ELEM]], %[[RESULT1]]) -// CHECK-NEXT: return %[[SECOND_ELEM]] +// CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[ARG1:.*]]: i1, %[[ARG2:.*]]: memref<5xf32>) +// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>) +// CHECK-NEXT: %[[ARG_TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) +// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 0 : i32} +// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 1 : i32} +// CHECK-NEXT: %[[THIRD_ELEM:.*]] = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 2 : i32} +// CHECK-NEXT: %[[CALLEE_RESULTS:.*]]:3 = call @callee(%[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]]) +// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>) -> (memref<2xf32>, i1, memref<5xf32>) +// CHECK-NEXT: %[[RESULT_TUPLE:.*]] = "test.make_tuple"(%[[CALLEE_RESULTS]]#0, %[[CALLEE_RESULTS]]#1, %[[CALLEE_RESULTS]]#2) +// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[RESULT_TUPLE]]) {index = 0 : i32} +// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[RESULT_TUPLE]]) {index = 1 : i32} +// CHECK-NEXT: %[[THIRD_ELEM:.*]] = "test.get_tuple_element"(%[[RESULT_TUPLE]]) {index = 2 : i32} +// CHECK-NEXT: %[[CALLEE_RESULTS:.*]]:3 = call @callee(%[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]]) +// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>) -> (memref<2xf32>, i1, memref<5xf32>) +// CHECK-NEXT: %[[RETURN_TUPLE:.*]] = "test.make_tuple"(%[[CALLEE_RESULTS]]#0, %[[CALLEE_RESULTS]]#1, %[[CALLEE_RESULTS]]#2) +// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[RETURN_TUPLE]]) {index = 0 : i32} +// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[RETURN_TUPLE]]) {index = 1 : i32} +// CHECK-NEXT: %[[THIRD_ELEM:.*]] = "test.get_tuple_element"(%[[RETURN_TUPLE]]) {index = 2 : i32} +// CHECK-NEXT: return %[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]] // ----- -// Test case: Testing BufferizeFuncOpConverter and BufferizeReturnOpConverter -// to see if the return operation matches with the new function signature when -// there are tuple typed args and results. BufferizeTypeConverter is set to -// flatten tuple typed arguments. The tuple typed values should be decomposed -// and composed using get_tuple_element and make_tuple operations of test -// dialect. Tensor types are converted to Memref. Memref typed function results -// are appended to the function arguments list. +// Test case: Testing BufferizeFuncOpConverter and +// BufferizeReturnOpConverter to see if the return operation matches with the +// new function signature when there are tuple typed args and results. +// BufferizeTypeConverter is set to flatten tuple typed arguments. The tuple +// typed values should be decomposed and composed using get_tuple_element and +// make_tuple operations of test dialect. Tensor types are converted to Memref. +// Memref typed function results remain as function results. // CHECK-LABEL: func @decompose_tuple_typed_function_args_and_results -func @decompose_tuple_typed_function_args_and_results(%arg0: tuple, - %arg1: tensor<10xf32>, - %arg2: tuple>) - -> (tuple>, - tensor<10xf32>, - tuple){ - return %arg2, %arg1, %arg0 : tuple>, tensor<10xf32>, - tuple +func @decompose_tuple_typed_function_args_and_results(%arg0: tuple, %arg1: tensor<10xf32>, %arg2: tuple>) -> (tuple>, tensor<10xf32>, tuple){ + return %arg2, %arg1, %arg0 : tuple>, tensor<10xf32>, tuple } -// CHECK-SAME: %[[ARG0:.*]]: i1, %[[ARG1:.*]]: f32, -// CHECK-SAME: %[[ARG2:.*]]: memref<10xf32>, %[[ARG3:.*]]: i1, -// CHECK-SAME: %[[ARG4:.*]]: memref<5xf32>, %[[RESULT0:.*]]: memref<5xf32>, -// CHECK-SAME: %[[RESULT1:.*]]: memref<10xf32> -// CHECK-SAME: (i1, i1, f32) +// CHECK-SAME: %[[ARG0:.*]]: i1, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: memref<10xf32>, %[[ARG3:.*]]: i1, %[[ARG4:.*]]: memref<5xf32> +// CHECK-SAME: (i1, memref<5xf32>, memref<10xf32>, i1, f32) // CHECK-NEXT: %[[FIRST_TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]]) // CHECK-NEXT: %[[SECOND_TUPLE:.*]] = "test.make_tuple"(%[[ARG3]], %[[ARG4]]) -// CHECK-NEXT: %[[SECOND_TUPLE_FIRST_ELEM:.*]] = "test.get_tuple_element" -// CHECK-SAME: (%[[SECOND_TUPLE]]) {index = 0 : i32} -// CHECK-NEXT: %[[SECOND_TUPLE_SECOND_ELEM:.*]] = "test.get_tuple_element" -// CHECK-SAME: (%[[SECOND_TUPLE]]) {index = 1 : i32} -// CHECK-NEXT: %[[FIRST_TUPLE_FIRST_ELEM:.*]] = "test.get_tuple_element" -// CHECK-SAME: (%[[FIRST_TUPLE]]) {index = 0 : i32} -// CHECK-NEXT: %[[FIRST_TUPLE_SECOND_ELEM:.*]] = "test.get_tuple_element" -// CHECK-SAME: (%[[FIRST_TUPLE]]) {index = 1 : i32} -// CHECK-NEXT: test.copy(%[[SECOND_TUPLE_SECOND_ELEM]], %[[RESULT0]]) -// CHECK-NEXT: test.copy(%[[ARG2]], %[[RESULT1]]) -// CHECK-NEXT: return %[[SECOND_TUPLE_FIRST_ELEM]], %[[FIRST_TUPLE_FIRST_ELEM]], -// CHECK-SAME: %[[FIRST_TUPLE_SECOND_ELEM]] +// CHECK-NEXT: %[[SECOND_TUPLE_FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[SECOND_TUPLE]]) {index = 0 : i32} +// CHECK-NEXT: %[[SECOND_TUPLE_SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[SECOND_TUPLE]]) {index = 1 : i32} +// CHECK-NEXT: %[[FIRST_TUPLE_FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[FIRST_TUPLE]]) {index = 0 : i32} +// CHECK-NEXT: %[[FIRST_TUPLE_SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[FIRST_TUPLE]]) {index = 1 : i32} +// CHECK-NEXT: return %[[SECOND_TUPLE_FIRST_ELEM]], %[[SECOND_TUPLE_SECOND_ELEM]], %[[ARG2]], %[[FIRST_TUPLE_FIRST_ELEM]], %[[FIRST_TUPLE_SECOND_ELEM]] diff --git a/mlir/test/lib/Transforms/TestFinalizingBufferize.cpp b/mlir/test/lib/Transforms/TestFinalizingBufferize.cpp --- a/mlir/test/lib/Transforms/TestFinalizingBufferize.cpp +++ b/mlir/test/lib/Transforms/TestFinalizingBufferize.cpp @@ -35,17 +35,9 @@ /// otherwise the IR will end up invalid. Thus, finalizing bufferization passes /// require an atomic change to the entire program (e.g. the whole module). /// -/// `allowMemrefFunctionResults` informs the buffer finalization policy to allow -/// functions that have memref typed results. Patterns involved with converting -/// func/call/return respect the finalization policy to ensure a consistent -/// atomic conversion of the entire module. `allowMemrefFunctionResults` also -/// allows memref typed results to escape from the deallocation. -/// /// TODO: Split out BufferizeFinalizationPolicy from BufferizeTypeConverter. -template struct TestFinalizingBufferizePass - : mlir::PassWrapper, - OperationPass> { + : mlir::PassWrapper> { /// Converts tensor based test operations to buffer based ones using /// bufferize. @@ -123,13 +115,6 @@ converter.isLegal(&funcOp.getBody()); }); - auto kind = allowMemrefFunctionResults - ? BufferizeTypeConverter::KeepAsFunctionResult - : BufferizeTypeConverter::AppendToArgumentsList; - converter.setResultConversionKind(kind); - converter.setResultConversionKind( - kind); - converter.addDecomposeTypeConversion( [](TupleType tupleType, SmallVectorImpl &types) { tupleType.getFlattenedTypes(types); @@ -175,17 +160,8 @@ namespace mlir { namespace test { void registerTestFinalizingBufferizePass() { - PassRegistration< - TestFinalizingBufferizePass>( + PassRegistration( "test-finalizing-bufferize", "Tests finalizing bufferize conversions"); } - -void registerTestPreparationPassWithAllowedMemrefResults() { - PassRegistration< - TestFinalizingBufferizePass>( - "test-finalizing-bufferize-with-allowed-memref-results", - "Tests finalizing buffierize conversions, allowing functions to have " - "memref typed results."); -} } // namespace test } // namespace mlir diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp --- a/mlir/tools/mlir-opt/mlir-opt.cpp +++ b/mlir/tools/mlir-opt/mlir-opt.cpp @@ -146,7 +146,6 @@ test::registerTestMemRefDependenceCheck(); test::registerTestMemRefStrideCalculation(); test::registerTestOpaqueLoc(); - test::registerTestPreparationPassWithAllowedMemrefResults(); test::registerTestRecursiveTypesPass(); test::registerTestSCFUtilsPass(); test::registerTestVectorConversions();