[flang-commits] [flang] Change induction variable from i32 to index for doconcurrent inside cuf kernel directive (PR #129924)

[Zhen Wang via flang-commits]

https://github.com/wangzpgi created https://github.com/llvm/llvm-project/pull/129924

Use `index` instead of `i32` for induction variables for doconcurrent inside cuf kernel directive. Regular do loop inside cuf kernel directive also uses `index`:
```
cuf.kernel<<<*, *>>> (%arg0 : index) = ...
```

>From 3bcfb0f8bbba5b2fc156454325534db4dc26d5ad Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Wed, 5 Mar 2025 11:54:03 -0800
Subject: [PATCH] Changedoconcurrent inside cuf kernel directive induction
 variable type from i32 to index to be the same as regular do loops inside cuf
 kernel.

---
 flang/lib/Lower/Bridge.cpp            | 15 +++++----------
 flang/test/Lower/CUDA/cuda-doconc.cuf | 14 +++++++-------
 2 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index 4c6e47d250329..95f431983d442 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -3150,10 +3150,9 @@ class FirConverter : public Fortran::lower::AbstractConverter {
               loc, 1); // Use index type directly
 
         // Ensure lb, ub, and step are of index type using fir.convert
-        mlir::Type indexType = builder->getIndexType();
-        lb = builder->create<fir::ConvertOp>(loc, indexType, lb);
-        ub = builder->create<fir::ConvertOp>(loc, indexType, ub);
-        step = builder->create<fir::ConvertOp>(loc, indexType, step);
+        lb = builder->create<fir::ConvertOp>(loc, idxTy, lb);
+        ub = builder->create<fir::ConvertOp>(loc, idxTy, ub);
+        step = builder->create<fir::ConvertOp>(loc, idxTy, step);
 
         lbs.push_back(lb);
         ubs.push_back(ub);
@@ -3163,10 +3162,6 @@ class FirConverter : public Fortran::lower::AbstractConverter {
 
         // Handle induction variable
         mlir::Value ivValue = getSymbolAddress(*name.symbol);
-        std::size_t ivTypeSize = name.symbol->size();
-        if (ivTypeSize == 0)
-          llvm::report_fatal_error("unexpected induction variable size");
-        mlir::Type ivTy = builder->getIntegerType(ivTypeSize * 8);
 
         if (!ivValue) {
           // DO CONCURRENT induction variables are not mapped yet since they are
@@ -3174,7 +3169,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
           mlir::OpBuilder::InsertPoint insPt = builder->saveInsertionPoint();
           builder->setInsertionPointToStart(builder->getAllocaBlock());
           ivValue = builder->createTemporaryAlloc(
-              loc, ivTy, toStringRef(name.symbol->name()));
+              loc, idxTy, toStringRef(name.symbol->name()));
           builder->restoreInsertionPoint(insPt);
         }
 
@@ -3186,7 +3181,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
         // Bind the symbol to the declared variable
         bindSymbol(*name.symbol, ivValue);
         ivValues.push_back(ivValue);
-        ivTypes.push_back(ivTy);
+        ivTypes.push_back(idxTy);
         ivLocs.push_back(loc);
       }
     } else {
diff --git a/flang/test/Lower/CUDA/cuda-doconc.cuf b/flang/test/Lower/CUDA/cuda-doconc.cuf
index 32cd1676b22f4..e240b1adc206a 100644
--- a/flang/test/Lower/CUDA/cuda-doconc.cuf
+++ b/flang/test/Lower/CUDA/cuda-doconc.cuf
@@ -15,9 +15,9 @@ subroutine doconc1
 end
 
 ! CHECK: func.func @_QPdoconc1() {
-! CHECK: %[[DECL:.*]]:2 = hlfir.declare %{{.*}}#0 {uniq_name = "_QFdoconc1Ei"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+! CHECK: %[[DECL:.*]]:2 = hlfir.declare %{{.*}}#0 {uniq_name = "_QFdoconc1Ei"} : (!fir.ref<index>) -> (!fir.ref<index>, !fir.ref<index>)
 ! CHECK: cuf.kernel<<<*, *>>>
-! CHECK: %{{.*}} = fir.load %[[DECL]]#0 : !fir.ref<i32>
+! CHECK: %{{.*}} = fir.load %[[DECL]]#0 : !fir.ref<index>
 
 subroutine doconc2
   integer :: i, j, m, n
@@ -32,8 +32,8 @@ subroutine doconc2
 end
 
 ! CHECK: func.func @_QPdoconc2() {
-! CHECK: %[[DECLI:.*]]:2 = hlfir.declare %{{.*}}#0 {uniq_name = "_QFdoconc2Ei"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
-! CHECK: %[[DECLJ:.*]]:2 = hlfir.declare %{{.*}}#0 {uniq_name = "_QFdoconc2Ej"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
-! CHECK: cuf.kernel<<<*, *>>> (%arg0 : i32, %arg1 : i32) = (%{{.*}}, %{{.*}} : index, index) to (%{{.*}}, %{{.*}} : index, index)  step (%{{.*}}, %{{.*}} : index, index) {
-! CHECK: %{{.*}} = fir.load %[[DECLI]]#0 : !fir.ref<i32>
-! CHECK: %{{.*}} = fir.load %[[DECLJ]]#0 : !fir.ref<i32> 
+! CHECK: %[[DECLI:.*]]:2 = hlfir.declare %{{.*}}#0 {uniq_name = "_QFdoconc2Ei"} : (!fir.ref<index>) -> (!fir.ref<index>, !fir.ref<index>)
+! CHECK: %[[DECLJ:.*]]:2 = hlfir.declare %{{.*}}#0 {uniq_name = "_QFdoconc2Ej"} : (!fir.ref<index>) -> (!fir.ref<index>, !fir.ref<index>)
+! CHECK: cuf.kernel<<<*, *>>> (%arg0 : index, %arg1 : index) = (%{{.*}}, %{{.*}} : index, index) to (%{{.*}}, %{{.*}} : index, index)  step (%{{.*}}, %{{.*}} : index, index) {
+! CHECK: %{{.*}} = fir.load %[[DECLI]]#0 : !fir.ref<index>
+! CHECK: %{{.*}} = fir.load %[[DECLJ]]#0 : !fir.ref<index>