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

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.

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D155825

Author: grypp

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

@@ -744,14 +744,13 @@ struct NVGPUMBarrierCreateLowering
   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;
-    if (auto moduleOp = dyn_cast<gpu::GPUModuleOp>(mOp))
+    if (auto moduleOp = funcOp->getParentOfType<gpu::GPUModuleOp>())
       global = generateGlobalBarrier(rewriter, funcOp, moduleOp, barrierType);
-    else if (auto moduleOp = dyn_cast<ModuleOp>(mOp))
+    else if (auto moduleOp = funcOp->getParentOfType<ModuleOp>())
       global = generateGlobalBarrier(rewriter, funcOp, moduleOp, barrierType);
 
     rewriter.setInsertionPoint(op);

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

@@ -635,3 +635,18 @@ func.func @async_tma_load(%tensorMap1d: !tensorMap1d, %tensorMap2d: !tensorMap2d
   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 
+  }
+}