[Mlir-commits] [mlir] 6ca0e4a - [mlir][OpenMP] Fix update of linear iteration variables (#183800)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jul 27 07:31:36 PDT 2026
Author: Leandro Lupori
Date: 2026-07-27T11:31:30-03:00
New Revision: 6ca0e4a67193836fe696eb0273d523da1950cfda
URL: https://github.com/llvm/llvm-project/commit/6ca0e4a67193836fe696eb0273d523da1950cfda
DIFF: https://github.com/llvm/llvm-project/commit/6ca0e4a67193836fe696eb0273d523da1950cfda.diff
LOG: [mlir][OpenMP] Fix update of linear iteration variables (#183800)
The final value of a linear iteration variable must be the loop
limit_value + step. Before this patch it was limit_value.
Fixes #170784.
Added:
Modified:
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
mlir/test/Target/LLVMIR/openmp-llvm-invalid.mlir
mlir/test/Target/LLVMIR/openmp-llvm.mlir
mlir/test/Target/LLVMIR/openmp-simd-guided.mlir
mlir/test/Target/LLVMIR/openmp-simd-ordered.mlir
mlir/test/Target/LLVMIR/openmp-wsloop-simd-ordered.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 95ccb53bce3d4..8605b39f17c16 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -3480,6 +3480,16 @@ LogicalResult SimdOp::verify() {
if (getLinearVars().size() &&
getLinearVarTypes().value().size() != getLinearVars().size())
return emitError() << "Ill-formed type attributes for linear variables";
+
+ llvm::DenseSet<Value> privateVars(llvm::from_range, getPrivateVars());
+ llvm::DenseSet<Value> reductionVars(llvm::from_range, getReductionVars());
+ // TODO Check lastprivate vars when their support is added to SimdOp.
+ for (Value var : getLinearVars()) {
+ if (privateVars.contains(var) || reductionVars.contains(var))
+ return emitOpError()
+ << "linear variables cannot appear in other data-sharing clauses";
+ }
+
return success();
}
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 97993a80926eb..1c50ff192c3d5 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -150,6 +150,7 @@ class LinearClauseProcessor {
llvm::BasicBlock *linearFinalizationBB;
llvm::BasicBlock *linearExitBB;
llvm::BasicBlock *linearLastIterExitBB;
+ Value linearLoopIV;
public:
// Register type for the linear variables
@@ -189,6 +190,35 @@ class LinearClauseProcessor {
}
}
+ // Find linear iteration variable and save it for later updates
+ LogicalResult initLinearIV(omp::SimdOp simdOp) {
+ auto loopOp = cast<omp::LoopNestOp>(simdOp.getWrappedLoop());
+ // NOTE iteration variables can only be linear in non-nested loops.
+ if (loopOp.getIVs().size() != 1)
+ return success();
+ // Currently, frontends using `omp.simd` always generate a store from the
+ // `omp.loop_nest`'s IV to the corresponding iteration variable.
+ // We leverage this to find the linear iteration variable.
+ //
+ // TODO Add an attribute to `omp.loop_nest` that explicitly lists the
+ // variables that correspond to the loop induction variables.
+ BlockArgument arg = loopOp.getIVs().front();
+ for (const Operation *user : arg.getUsers()) {
+ if (auto storeOp = dyn_cast<LLVM::StoreOp>(user)) {
+ for (Value linearVar : simdOp.getLinearVars()) {
+ if (linearVar == storeOp.getAddr()) {
+ if (linearLoopIV && linearLoopIV != linearVar)
+ return simdOp.emitError(
+ "Could not determine the linear variable associated with the "
+ "loop nest induction variable");
+ linearLoopIV = linearVar;
+ }
+ }
+ }
+ }
+ return success();
+ }
+
// Emit IR for updating Linear variables
void updateLinearVar(llvm::IRBuilderBase &builder, llvm::BasicBlock *loopBody,
llvm::Value *loopInductionVar) {
@@ -229,6 +259,34 @@ class LinearClauseProcessor {
}
}
+ // Emit IR for updating linear iteration variables on loop exit
+ void updateLinearIV(llvm::IRBuilderBase &builder,
+ LLVM::ModuleTranslation &moduleTranslation) {
+ if (!linearLoopIV)
+ return;
+ llvm::Value *linearIV = moduleTranslation.lookupValue(linearLoopIV);
+
+ // Find linearIV's index
+ size_t index;
+ for (index = 0; index < linearOrigVal.size(); index++)
+ if (linearIV == linearOrigVal[index])
+ break;
+ if (index == linearOrigVal.size())
+ return;
+
+ // Add one more step to the linear iteration variable
+ llvm::Type *varType = linearVarTypes[index];
+ llvm::Value *var = linearLoopBodyTemps[index];
+ llvm::Value *step = linearSteps[index];
+ if (!varType->isIntegerTy())
+ llvm_unreachable("Linear iteration variable must be of integer type");
+
+ step = builder.CreateSExtOrTrunc(step, varType);
+ llvm::Value *val = builder.CreateLoad(varType, var);
+ llvm::Value *addInst = builder.CreateAdd(val, step);
+ builder.CreateStore(addInst, var);
+ }
+
// Linear variable finalization is conditional on the last logical iteration.
// Create BB splits to manage the same.
void splitLinearFiniBB(llvm::IRBuilderBase &builder,
@@ -4968,6 +5026,8 @@ convertOmpSimd(Operation &opInst, llvm::IRBuilderBase &builder,
// Initialize linear variables and linear step
LinearClauseProcessor linearClauseProcessor;
+ if (linearClauseProcessor.initLinearIV(simdOp).failed())
+ return failure();
if (!simdOp.getLinearVars().empty()) {
auto linearVarTypes = simdOp.getLinearVarTypes().value();
@@ -5081,6 +5141,7 @@ convertOmpSimd(Operation &opInst, llvm::IRBuilderBase &builder,
: nullptr,
order, simdlen, safelen);
+ linearClauseProcessor.updateLinearIV(builder, moduleTranslation);
linearClauseProcessor.emitStoresForLinearVar(builder);
// We now need to reduce the per-simd-lane reduction variable into the
diff --git a/mlir/test/Target/LLVMIR/openmp-llvm-invalid.mlir b/mlir/test/Target/LLVMIR/openmp-llvm-invalid.mlir
index dd427097bd600..234f404ff5690 100644
--- a/mlir/test/Target/LLVMIR/openmp-llvm-invalid.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-llvm-invalid.mlir
@@ -116,6 +116,35 @@ llvm.func @simd_linear(%lb : i32, %ub : i32, %step : i32, %x : !llvm.ptr) {
// -----
+omp.private {type = private} @i_private_i32 : i32
+llvm.func @simd_linear_private(%lb : i32, %ub : i32, %step : i32, %i : !llvm.ptr) {
+ // expected-error @below {{linear variables cannot appear in other data-sharing clauses}}
+ omp.simd linear(%i : !llvm.ptr = %step : i32)
+ private(@i_private_i32 %i -> %priv_i : !llvm.ptr) {
+ omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
+ omp.yield
+ }
+ } {linear_var_types = [i32]}
+ llvm.return
+}
+
+// -----
+
+llvm.func @simd_linear_ambiguous_iv(%lb : i32, %ub : i32, %step : i32, %x : !llvm.ptr, %y : !llvm.ptr) {
+ // expected-error @below {{Could not determine the linear variable associated with the loop nest induction variable}}
+ // expected-error @below {{LLVM Translation failed for operation: omp.simd}}
+ omp.simd linear(%x : !llvm.ptr = %step : i32, %y : !llvm.ptr = %step : i32) {
+ omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
+ llvm.store %iv, %x : i32, !llvm.ptr
+ llvm.store %iv, %y : i32, !llvm.ptr
+ omp.yield
+ }
+ } {linear_var_types = [i32, i32]}
+ llvm.return
+}
+
+// -----
+
module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_target_device = true} {
llvm.func @host_op_in_device(%arg0 : !llvm.ptr) {
// expected-error @below {{unsupported host op found in device}}
diff --git a/mlir/test/Target/LLVMIR/openmp-llvm.mlir b/mlir/test/Target/LLVMIR/openmp-llvm.mlir
index d26557d57d0dd..e867dd8afcb9b 100644
--- a/mlir/test/Target/LLVMIR/openmp-llvm.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-llvm.mlir
@@ -767,6 +767,7 @@ llvm.func @simd_simple(%lb : i64, %ub : i64, %step : i64, %arg0: !llvm.ptr) {
llvm.func @simd_linear(%lb : i32, %ub : i32, %step : i32, %x : !llvm.ptr) {
// CHECK-LABEL: @simd_linear
+// CHECK-SAME: (i32 %{{.*}}, i32 %{{.*}}, i32 %[[STEP:.*]], ptr %[[X:.*]])
// CHECK: %[[LINEAR_VAR:.*]] = alloca i32, align 4
// CHECK: %[[LINEAR_RESULT:.*]] = alloca i32, align 4
@@ -782,8 +783,17 @@ llvm.func @simd_linear(%lb : i32, %ub : i32, %step : i32, %x : !llvm.ptr) {
// CHECK: %[[MUL:.*]] = mul i32 %omp_loop.iv, {{.*}}
// CHECK: %[[ADD:.*]] = add i32 %[[LOAD]], %[[MUL]]
// CHECK: store i32 %[[ADD]], ptr %[[LINEAR_RESULT]], align 4, !llvm.access.group !1
+
+// CHECK: omp.region.cont:
+// CHECK: %[[VAL:.*]] = load i32, ptr %[[LINEAR_RESULT]]
+// CHECK-NEXT: %[[ADD:.*]] = add i32 %[[VAL]], %[[STEP]]
+// CHECK-NEXT: store i32 %[[ADD]], ptr %[[LINEAR_RESULT]]
+// CHECK-NEXT: %[[LOAD:.*]] = load i32, ptr %[[LINEAR_RESULT]]
+// CHECK-NEXT: store i32 %[[LOAD]], ptr %[[X]], align 4
omp.simd linear(%x : !llvm.ptr = %step : i32) {
omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
+ llvm.store %iv, %x : i32, !llvm.ptr
+ llvm.store %iv, %x : i32, !llvm.ptr
omp.yield
}
} {linear_var_types = [i32]}
@@ -856,6 +866,36 @@ llvm.func @simd_linear_f64_var_i32_step(%lb : i32, %ub : i32, %x : !llvm.ptr) {
// -----
+// Test the update of omp.simd linear iteration variables, when nested inside
+// omp.wsloop.
+llvm.func @wsloop_simd_linear(%x : !llvm.ptr) {
+
+// CHECK-LABEL: @wsloop_simd_linear
+
+// CHECK: omp.wsloop.region:
+// CHECK: %[[LINEAR_VAR:.*]] = alloca i32
+// CHECK: %[[LINEAR_RESULT:.*]] = alloca i32
+
+// CHECK: omp.region.cont2:
+// CHECK: %[[VAL:.*]] = load i32, ptr %[[LINEAR_RESULT]]
+// CHECK-NEXT: %[[ADD:.*]] = add i32 %[[VAL]], 25
+// CHECK-NEXT: store i32 %[[ADD]], ptr %[[LINEAR_RESULT]]
+ %lb = llvm.mlir.constant(1 : i32) : i32
+ %ub = llvm.mlir.constant(100 : i32) : i32
+ %step = llvm.mlir.constant(25 : i32) : i32
+ omp.wsloop {
+ omp.simd linear(%x : !llvm.ptr = %step : i32) {
+ omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
+ llvm.store %iv, %x : i32, !llvm.ptr
+ omp.yield
+ }
+ } {linear_var_types = [i32], omp.composite}
+ } {omp.composite}
+ llvm.return
+}
+
+// -----
+
// CHECK-LABEL: @simd_simple_multiple
llvm.func @simd_simple_multiple(%lb1 : i64, %ub1 : i64, %step1 : i64, %lb2 : i64, %ub2 : i64, %step2 : i64, %arg0: !llvm.ptr, %arg1: !llvm.ptr) {
omp.simd {
diff --git a/mlir/test/Target/LLVMIR/openmp-simd-guided.mlir b/mlir/test/Target/LLVMIR/openmp-simd-guided.mlir
index b1f18b9f85417..496669c5b90d5 100644
--- a/mlir/test/Target/LLVMIR/openmp-simd-guided.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-simd-guided.mlir
@@ -1,7 +1,6 @@
// Ensure that schedule can be used with the guided kind-type and simd construct
// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s
-omp.private {type = private} @_QFEi_private_i32 : i32
llvm.func @test_simd_guided() {
%0 = llvm.mlir.constant (1 : i64) : i64
%c1_i32 = llvm.mlir.constant (1 : i32) : i32
@@ -10,7 +9,7 @@ llvm.func @test_simd_guided() {
%c4_i32 = llvm.mlir.constant (4 : i32) : i32
%1 = llvm.alloca %0 x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr
omp.wsloop schedule(guided = %c4_i32 : i32, simd) {
- omp.simd linear(%1 : !llvm.ptr = %c1_i32 : i32) private(@_QFEi_private_i32 %1 -> %arg0 : !llvm.ptr) {
+ omp.simd linear(%1 : !llvm.ptr = %c1_i32 : i32) {
omp.loop_nest (%arg1) : i32 = (%c0_i32) to (%c64_i32) inclusive step (%c1_i32) {
omp.yield
}
diff --git a/mlir/test/Target/LLVMIR/openmp-simd-ordered.mlir b/mlir/test/Target/LLVMIR/openmp-simd-ordered.mlir
index 8da1fab1f48af..a84c51f08daa7 100644
--- a/mlir/test/Target/LLVMIR/openmp-simd-ordered.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-simd-ordered.mlir
@@ -17,8 +17,6 @@
// !$omp end simd
module {
- omp.private {type = private} @i_private_i32 : i32
-
// CHECK-LABEL: define void @simd_ordered_linear
llvm.func @simd_ordered_linear() {
%c0_i64 = llvm.mlir.constant(0 : i64) : i64
@@ -38,14 +36,14 @@ module {
// CHECK: %.linear_var = alloca i32
// CHECK: %.linear_result = alloca i32
- omp.simd linear(%i : !llvm.ptr = %c1_i32 : i32) private(@i_private_i32 %i -> %arg0 : !llvm.ptr) {
+ omp.simd linear(%i : !llvm.ptr = %c1_i32 : i32) {
omp.loop_nest (%iv) : i32 = (%c1_i32) to (%c10_i32) inclusive step (%c1_i32) {
// CHECK: omp.loop_nest.region:
// CHECK: load i32, ptr %.linear_result
- llvm.store %iv, %arg0 : i32, !llvm.ptr
+ llvm.store %iv, %i : i32, !llvm.ptr
// Compute a[i] = b[i] * 10
- %i_val = llvm.load %arg0 : !llvm.ptr -> i32
+ %i_val = llvm.load %i : !llvm.ptr -> i32
%i_idx = llvm.sext %i_val : i32 to i64
%i_off = llvm.sub %i_idx, %c1_i64 : i64
%b_ptr = llvm.getelementptr %b[%i_off] : (!llvm.ptr, i64) -> !llvm.ptr, i32
@@ -58,7 +56,7 @@ module {
omp.ordered.region par_level_simd {
// CHECK: omp.ordered.region:
// CHECK: load i32, ptr %.linear_result
- %i_ord = llvm.load %arg0 : !llvm.ptr -> i32
+ %i_ord = llvm.load %i : !llvm.ptr -> i32
%i_ord_idx = llvm.sext %i_ord : i32 to i64
%i_ord_off = llvm.sub %i_ord_idx, %c1_i64 : i64
%a_ord_ptr = llvm.getelementptr %a[%i_ord_off] : (!llvm.ptr, i64) -> !llvm.ptr, i32
@@ -69,7 +67,7 @@ module {
// Compute c[i] = a[i] * 2 (code after ordered region)
// CHECK: omp_region.finalize:
// CHECK: load i32, ptr %.linear_result
- %i_post = llvm.load %arg0 : !llvm.ptr -> i32
+ %i_post = llvm.load %i : !llvm.ptr -> i32
%i_post_idx = llvm.sext %i_post : i32 to i64
%i_post_off = llvm.sub %i_post_idx, %c1_i64 : i64
%a_post_ptr = llvm.getelementptr %a[%i_post_off] : (!llvm.ptr, i64) -> !llvm.ptr, i32
diff --git a/mlir/test/Target/LLVMIR/openmp-wsloop-simd-ordered.mlir b/mlir/test/Target/LLVMIR/openmp-wsloop-simd-ordered.mlir
index beb300c23651f..8d5b89ecbd00a 100644
--- a/mlir/test/Target/LLVMIR/openmp-wsloop-simd-ordered.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-wsloop-simd-ordered.mlir
@@ -17,8 +17,6 @@
// !$omp end do simd
module {
- omp.private {type = private} @i_private_i32 : i32
-
// CHECK-LABEL: define void @wsloop_simd_ordered_linear
llvm.func @wsloop_simd_ordered_linear() {
%c0_i64 = llvm.mlir.constant(0 : i64) : i64
@@ -39,14 +37,14 @@ module {
// CHECK: %.linear_result = alloca i32
omp.wsloop ordered(0) {
- omp.simd linear(%i : !llvm.ptr = %c1_i32 : i32) private(@i_private_i32 %i -> %arg0 : !llvm.ptr) {
+ omp.simd linear(%i : !llvm.ptr = %c1_i32 : i32) {
omp.loop_nest (%iv) : i32 = (%c1_i32) to (%c100_i32) inclusive step (%c1_i32) {
// CHECK: omp.loop_nest.region:
// CHECK: load i32, ptr %.linear_result
- llvm.store %iv, %arg0 : i32, !llvm.ptr
+ llvm.store %iv, %i : i32, !llvm.ptr
// Compute a[i] = b[i] * 10
- %i_val = llvm.load %arg0 : !llvm.ptr -> i32
+ %i_val = llvm.load %i : !llvm.ptr -> i32
%i_idx = llvm.sext %i_val : i32 to i64
%i_off = llvm.sub %i_idx, %c1_i64 : i64
%b_ptr = llvm.getelementptr %b[%i_off] : (!llvm.ptr, i64) -> !llvm.ptr, i32
@@ -59,7 +57,7 @@ module {
omp.ordered.region par_level_simd {
// CHECK: omp.ordered.region:
// CHECK: load i32, ptr %.linear_result
- %i_ord = llvm.load %arg0 : !llvm.ptr -> i32
+ %i_ord = llvm.load %i : !llvm.ptr -> i32
%i_ord_idx = llvm.sext %i_ord : i32 to i64
%i_ord_off = llvm.sub %i_ord_idx, %c1_i64 : i64
%a_ord_ptr = llvm.getelementptr %a[%i_ord_off] : (!llvm.ptr, i64) -> !llvm.ptr, i32
@@ -71,7 +69,7 @@ module {
// Compute c[i] = a[i] * 2 (code after ordered region)
// CHECK: omp_region.finalize:
// CHECK: load i32, ptr %.linear_result
- %i_post = llvm.load %arg0 : !llvm.ptr -> i32
+ %i_post = llvm.load %i : !llvm.ptr -> i32
%i_post_idx = llvm.sext %i_post : i32 to i64
%i_post_off = llvm.sub %i_post_idx, %c1_i64 : i64
%a_post_ptr = llvm.getelementptr %a[%i_post_off] : (!llvm.ptr, i64) -> !llvm.ptr, i32
More information about the Mlir-commits
mailing list