diff --git a/mlir/lib/Dialect/SCF/SCF.cpp b/mlir/lib/Dialect/SCF/SCF.cpp --- a/mlir/lib/Dialect/SCF/SCF.cpp +++ b/mlir/lib/Dialect/SCF/SCF.cpp @@ -660,6 +660,7 @@ scf::ForOp newForOp = rewriter.create( forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), forOp.getStep(), newIterArgs); + newForOp->setAttrs(forOp->getAttrs()); Block &newBlock = newForOp.getRegion().front(); // Replace the null placeholders with newly constructed values. @@ -802,6 +803,7 @@ scf::ForOp newForOp = rewriter.create( forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), forOp.getStep(), newIterOperands); + newForOp->setAttrs(forOp->getAttrs()); Block &newBlock = newForOp.getRegion().front(); SmallVector newBlockTransferArgs(newBlock.getArguments().begin(), newBlock.getArguments().end()); diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp --- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp +++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp @@ -491,6 +491,7 @@ auto newForOp = rewriter.create( forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), forOp.getStep(), initArgs); + newForOp->setAttrs(forOp->getAttrs()); ValueRange initArgsRange(initArgs); TypeRange initArgsTypes(initArgsRange); Block *loopBody = &newForOp.getLoopBody().front(); diff --git a/mlir/test/Dialect/SCF/bufferize.mlir b/mlir/test/Dialect/SCF/bufferize.mlir --- a/mlir/test/Dialect/SCF/bufferize.mlir +++ b/mlir/test/Dialect/SCF/bufferize.mlir @@ -30,14 +30,14 @@ // CHECK: %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref // CHECK: %[[RESULT_MEMREF:.*]] = scf.for %[[VAL_6:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] iter_args(%[[ITER:.*]] = %[[MEMREF]]) -> (memref) { // CHECK: scf.yield %[[ITER]] : memref -// CHECK: } +// CHECK: } {some_attr} // CHECK: %[[VAL_8:.*]] = bufferization.to_tensor %[[VAL_9:.*]] : memref // CHECK: return %[[VAL_8]] : tensor // CHECK: } func.func @for(%arg0: tensor, %lb: index, %ub: index, %step: index) -> tensor { %ret = scf.for %iv = %lb to %ub step %step iter_args(%iter = %arg0) -> tensor { scf.yield %iter : tensor - } + } {some_attr} return %ret : tensor } diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir --- a/mlir/test/Dialect/SCF/canonicalize.mlir +++ b/mlir/test/Dialect/SCF/canonicalize.mlir @@ -372,7 +372,7 @@ %r:3 = scf.for %i = %lb to %ub step %step iter_args(%0 = %a, %1 = %a, %2 = %b) -> (i32, i32, i32) { %c = func.call @make_i32() : () -> (i32) scf.yield %0, %c, %2 : i32, i32, i32 - } + } {some_attr} return %r#0, %r#1, %r#2 : i32, i32, i32 } @@ -382,7 +382,7 @@ // CHECK-NEXT: %[[r1:.*]] = scf.for {{.*}} iter_args(%arg4 = %[[a]]) -> (i32) { // CHECK-NEXT: %[[c:.*]] = func.call @make_i32() : () -> i32 // CHECK-NEXT: scf.yield %[[c]] : i32 -// CHECK-NEXT: } +// CHECK-NEXT: } {some_attr} // CHECK-NEXT: return %[[a]], %[[r1]], %[[b]] : i32, i32, i32 // ----- @@ -846,11 +846,12 @@ // CHECK: %[[DONE:.*]] = func.call @do(%[[CAST]]) : (tensor) -> tensor // CHECK: %[[UNCAST:.*]] = tensor.cast %[[DONE]] : tensor to tensor<32x1024xf32> // CHECK: scf.yield %[[UNCAST]] : tensor<32x1024xf32> +// CHECK: } {some_attr} %0 = tensor.cast %t0 : tensor<32x1024xf32> to tensor %1 = scf.for %i = %c0 to %c1024 step %c32 iter_args(%iter_t0 = %0) -> (tensor) { %2 = func.call @do(%iter_t0) : (tensor) -> tensor scf.yield %2 : tensor - } + } {some_attr} // CHECK-NOT: tensor.cast // CHECK: %[[RES:.*]] = tensor.insert_slice %[[FOR_RES]] into %[[T1]][0, 0] [32, 1024] [1, 1] : tensor<32x1024xf32> into tensor<1024x1024xf32> // CHECK: return %[[RES]] : tensor<1024x1024xf32>