[mlir][nvgpu] Improve finding module Op to for `mbarrier.create`

[guraypp]

Summary

Current transformation expects module op to be two level higher, however, it is not always the case. This work searches module op in a while loop.

[guraypp]

add newline

[nicolasvasilache]

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

749

you could

if (auto moduleOp = mOp->getParentOfType<gpu::GPUModuleOp>())
  global = ...
else if (auto moduleOp = mOp->getParentOfType<ModuleOp>())
  global = ...
else 
  return op->emitError()...

yes there would be 2 potential traversals in the worst case but I'd err towards legibility and avoid handrolling our own while

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

memref::GlobalOp generateGlobalBarrier(ConversionPatternRewriter &rewriter,
     symbolTable.insert(global);
     return global;
   }
 
   LogicalResult
   matchAndRewrite(nvgpu::MBarrierCreateOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Operation *funcOp = op->getParentOp();
-    Operation *mOp = funcOp->getParentOp();
     MemRefType barrierType =
         createMBarrierMemrefType(rewriter, op.getBarrier().getType());
 
-    memref::GlobalOp global;
+    std::optional<memref::GlobalOp> global = std::nullopt;
+    Operation *mOp = funcOp->getParentOp();
+    while (mOp != nullptr && !global.has_value()) {
       if (auto moduleOp = dyn_cast<gpu::GPUModuleOp>(mOp))
         global = generateGlobalBarrier(rewriter, funcOp, moduleOp, barrierType);
       else if (auto moduleOp = dyn_cast<ModuleOp>(mOp))
         global = generateGlobalBarrier(rewriter, funcOp, moduleOp, barrierType);
-
+      mOp = mOp->getParentOp();
+    };
+    if (!global.has_value()) {
+      return op->emitError() << "Expected parent op to be either "
+                                "gpu::GPUModuleOp or ModuleOp "
+                                "to genere mbarrier object.";
+    }
     rewriter.setInsertionPoint(op);
     rewriter.replaceOpWithNewOp<memref::GetGlobalOp>(op, barrierType,
-                                                     global.getName());
+                                                     global->getName());
     return success();
   }
 };
 
 /// Lowers `nvgpu.mbarrier.init` to `nvvm.mbarrier.init`
 struct NVGPUMBarrierInitLowering
     : public ConvertOpToLLVMPattern<nvgpu::MBarrierInitOp> {
   using ConvertOpToLLVMPattern<nvgpu::MBarrierInitOp>::ConvertOpToLLVMPattern;

mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir

func.func @mbarrier_txcount() {
 
     %phase = arith.constant 0 : index
     %ticks = arith.constant 10000000 : index
     // CHECK: %[[barPtr3:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)> 
     // CHECK: nvvm.mbarrier.try_wait.parity.shared %[[barPtr3]]
     nvgpu.mbarrier.try_wait.parity %barrier, %phase, %ticks : !barrierType
 
     func.return 
+}
+// -----
+
+!barrierType = !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>>
+module @find_parent{
+  func.func @main()  {
+  %c1 = arith.constant 1 : index   
+    gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
+          threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1, %block_z = %c1) {
+          // CHECK: memref.get_global @__mbarrier : memref<1xi64, 3>
+          %barrier = nvgpu.mbarrier.create -> !barrierType
+          gpu.terminator
+    }
+    func.return 
+  }
 }
 No newline at end of file