diff --git a/mlir/include/mlir/Dialect/LoopOps/LoopOps.td b/mlir/include/mlir/Dialect/LoopOps/LoopOps.td index 1e446efa9758..436a376a6f4f 100644 --- a/mlir/include/mlir/Dialect/LoopOps/LoopOps.td +++ b/mlir/include/mlir/Dialect/LoopOps/LoopOps.td @@ -312,10 +312,7 @@ def ParallelOp : Loop_Op<"parallel", let extraClassDeclaration = [{ Block *getBody() { return ®ion().front(); } - unsigned getNumInductionVars() { - return getBody()->getNumArguments(); - } - Block::BlockArgListType getInductionVars() { + ValueRange getInductionVars() { return getBody()->getArguments(); } unsigned getNumLoops() { return step().size(); } diff --git a/mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp b/mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp index 9697688ac850..177fdf0d668f 100644 --- a/mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp +++ b/mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp @@ -130,8 +130,8 @@ static void mapParallelOp(ParallelOp parallelOp, MLIRContext *ctx = parallelOp.getContext(); Builder b(ctx); SmallVector attrs; - attrs.reserve(parallelOp.getNumInductionVars()); - for (int i = 0, e = parallelOp.getNumInductionVars(); i < e; ++i) { + attrs.reserve(parallelOp.getNumLoops()); + for (int i = 0, e = parallelOp.getNumLoops(); i < e; ++i) { attrs.push_back(getParallelLoopDimMappingAttr( getHardwareIdForMapping(mappingLevel, i), b.getDimIdentityMap(), b.getDimIdentityMap()));