diff --git a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp @@ -71,18 +71,17 @@ return success(); } -/// Specialization for `linalg::GenericOp` and `linalg::IndexedGenericOp`. +/// Specialization for `linalg::GenericOp`. /// A pattern to convert Generic Linalg operations which work on tensors to /// use buffers. BufferPlacement pass should be later used to move /// Alloc operations to the correct positions and insert the missing Dealloc /// operations in the correct places. -template static void finalizeBufferAllocationForGenericOp(ConversionPatternRewriter &rewriter, - GenericOpTy genericOp, ValueRange inputs, + GenericOp genericOp, ValueRange inputs, ValueRange outputs) { // Generate a new linalg operation that works on buffers. - auto newGenericOp = rewriter.create( + auto newGenericOp = rewriter.create( genericOp.getLoc(), /*resultTensorTypes=*/llvm::None, /*inputs=*/inputs, @@ -116,7 +115,6 @@ linalg::LinalgOp linalgOp, ValueRange inputs, ValueRange outputs) { assert(!isa(linalgOp.getOperation())); - assert(!isa(linalgOp.getOperation())); SmallVector newOperands = inputs; newOperands.append(outputs.begin(), outputs.end()); auto otherOperands = linalgOp.getAssumedNonShapedOperands(); @@ -195,6 +193,10 @@ LogicalResult matchAndRewrite(LinalgOp op, ArrayRef operands, ConversionPatternRewriter &rewriter) const final { + // Canonicalize indexed generic operations before bufferization. + if (isa(op)) + return failure(); + // GenericOpAdaptor below expects an `operand_segment_sizes` attribute. if (!op->hasAttr("operand_segment_sizes")) return failure(); @@ -215,15 +217,8 @@ // Delegate to the linalg generic pattern. if (auto genericOp = dyn_cast(*op)) { - finalizeBufferAllocationForGenericOp( - rewriter, genericOp, adaptor.inputs(), newOutputBuffers); - return success(); - } - - // Delegate to the linalg indexed generic pattern. - if (auto genericOp = dyn_cast(*op)) { - finalizeBufferAllocationForGenericOp( - rewriter, genericOp, adaptor.inputs(), newOutputBuffers); + finalizeBufferAllocationForGenericOp(rewriter, genericOp, + adaptor.inputs(), newOutputBuffers); return success(); } diff --git a/mlir/test/Dialect/Linalg/bufferize.mlir b/mlir/test/Dialect/Linalg/bufferize.mlir --- a/mlir/test/Dialect/Linalg/bufferize.mlir +++ b/mlir/test/Dialect/Linalg/bufferize.mlir @@ -91,32 +91,6 @@ // ----- -#map0 = affine_map<(d0) -> (d0)> - -// CHECK-LABEL: func @multiple_results_indexed -// CHECK: %[[RESULT0:.*]] = memref.alloc() : memref<4xi32> -// CHECK: %[[RESULT1:.*]] = memref.alloc() : memref<4xi32> -// CHECK: linalg.generic -// CHECK-SAME: ins(%{{.*}} : memref<4xi32>) -// CHECK-SAME: outs(%[[RESULT0]], %[[RESULT1]] : memref<4xi32>, memref<4xi32>) -// CHECK-NEXT: ^bb0(%{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32): -func @multiple_results_indexed(%arg0: tensor<4xi32>) - -> (tensor<4xi32>, tensor<4xi32>) { - %0, %1 = linalg.indexed_generic { - indexing_maps = [#map0, #map0, #map0], - iterator_types = ["parallel"] - } ins(%arg0 : tensor<4xi32>) - outs (%arg0, %arg0 : tensor<4xi32>, tensor<4xi32>) { - ^bb0(%i: index, %gen_arg1: i32, %out1: i32, %out2: i32): - %i_i32 = index_cast %i : index to i32 - %tmp1 = addi %gen_arg1, %i_i32 : i32 - linalg.yield %tmp1, %tmp1 : i32, i32 - } -> tensor<4xi32>, tensor<4xi32> - return %0, %1 : tensor<4xi32>, tensor<4xi32> -} - -// ----- - #map_2d = affine_map<(d0, d1) -> (d0, d1)> // Check that the allocs properly consider the different shapes of the output