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 @@ -2180,6 +2180,16 @@ if (!innerOp) return failure(); + // Merging loops with attributes is not supported yet. + // We need to figure out how to properly merge those attributes without + // losing information. + auto isUnexpectedAttr = [](mlir::NamedAttribute attr) { + return attr.getName() != ParallelOp::getOperandSegmentSizeAttr(); + }; + if (llvm::any_of(op->getAttrs(), isUnexpectedAttr) || + llvm::any_of(innerOp->getAttrs(), isUnexpectedAttr)) + return failure(); + auto hasVal = [](const auto &range, Value val) { return llvm::find(range, val) != range.end(); }; 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 @@ -134,6 +134,36 @@ // ----- +func.func @nested_parallel_attr(%0: memref) -> memref { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %1 = memref.dim %0, %c0 : memref + %2 = memref.dim %0, %c1 : memref + %3 = memref.alloc(%1, %2) : memref + scf.parallel (%arg1) = (%c0) to (%1) step (%c1) { + scf.parallel (%arg2) = (%c0) to (%2) step (%c1) { + %4 = memref.load %0[%arg1, %arg2] : memref + memref.store %4, %3[%arg1, %arg2] : memref + scf.yield + } {mapping = [{bound = affine_map<(d0) -> (d0)>, map = affine_map<(d0) -> (d0)>, processor = 3 : i64}]} + scf.yield + } {mapping = [{bound = affine_map<(d0) -> (d0)>, map = affine_map<(d0) -> (d0)>, processor = 0 : i64}]} + return %3 : memref +} + +// CHECK-LABEL: func @nested_parallel_attr( +// CHECK-DAG: [[C0:%.*]] = arith.constant 0 : index +// CHECK-DAG: [[C1:%.*]] = arith.constant 1 : index +// CHECK-DAG: [[B0:%.*]] = memref.dim {{.*}}, [[C0]] +// CHECK-DAG: [[B1:%.*]] = memref.dim {{.*}}, [[C1]] +// CHECK: scf.parallel ([[V0:%.*]]) = ([[C0]]) to ([[B0]]) step ([[C1]]) +// CHECK: scf.parallel ([[V1:%.*]]) = ([[C0]]) to ([[B1]]) step ([[C1]]) +// CHECK: memref.load {{.*}}{{\[}}[[V0]], [[V1]]] +// CHECK: memref.store {{.*}}{{\[}}[[V0]], [[V1]]] +// CHECK: {mapping = {{.*}} processor = 3 +// CHECK: {mapping = {{.*}} processor = 0 +// ----- + func.func private @side_effect() func.func @one_unused(%cond: i1) -> (index) { %0, %1 = scf.if %cond -> (index, index) {