diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp --- a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp +++ b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp @@ -537,7 +537,7 @@ return; // Caller is expected to fully compose map/operands if necessary. - assert((isTopLevelValue(val) || isAffineForInductionVar(val)) && + assert((isTopLevelValue(val) || isAffineInductionVar(val)) && "non-terminal symbol / loop IV expected"); // Outer loop IVs could be used in forOp's bounds. if (auto loop = getForInductionVarOwner(val)) { @@ -547,6 +547,14 @@ loop.emitWarning("failed to add domain info to constraint system")); return; } + if (auto parallel = getAffineParallelInductionVarOwner(val)) { + appendDimVar(parallel.getIVs()); + if (failed(this->addAffineParallelOpDomain(parallel))) + LLVM_DEBUG(parallel.emitWarning( + "failed to add domain info to constraint system")); + return; + } + // Add top level symbol. appendSymbolVar(val); // Check if the symbol is a constant. diff --git a/mlir/test/Dialect/Affine/affine-data-copy.mlir b/mlir/test/Dialect/Affine/affine-data-copy.mlir --- a/mlir/test/Dialect/Affine/affine-data-copy.mlir +++ b/mlir/test/Dialect/Affine/affine-data-copy.mlir @@ -286,3 +286,26 @@ // CHECK-NOT: memref.alloc // CHECK: return } + +#map16 = affine_map<(d0, d1, d2) -> (d0 * 40 + d1 * 8 + d2 * 2)> +#map17 = affine_map<(d0, d1, d2) -> (d0 * 40 + d1 * 8 + d2 * 2 + 2)> +func.func @main(%85:memref<2x5x4x2xi64>) { + affine.for %arg0 = 0 to 2 { + affine.parallel (%arg1) = (0) to (5) { + affine.parallel (%arg2) = (0) to (4) { + affine.for %arg3 = #map16(%arg0, %arg1, %arg2) to #map17(%arg0, %arg1, %arg2) { + %105 = affine.load %85[((%arg3 floordiv 2) floordiv 4) floordiv 5, ((%arg3 floordiv 2) floordiv 4) mod 5, (%arg3 floordiv 2) mod 4, %arg3 mod 2] : memref<2x5x4x2xi64> + } + } + } + } + // CHECK: affine.for + // CHECK-NEXT: affine.for %{{.*}} = 0 to 5 + // CHECK-NEXT: affine.for %{{.*}} = 0 to 4 + // CHECK-NEXT: affine.for %{{.*}} = 0 to 2 + + // CHECK: affine.for + // CHECK-NEXT: affine.parallel + // CHECK-NEXT: affine.parallel + return +}