diff --git a/mlir/extras/dialects/ext/gpu.py b/mlir/extras/dialects/ext/gpu.py index 8d7af826..b1ee6a21 100644 --- a/mlir/extras/dialects/ext/gpu.py +++ b/mlir/extras/dialects/ext/gpu.py @@ -184,7 +184,7 @@ def body(self): return self.regions[0].blocks[0] -module = region_op(GPUModuleOp, terminator=lambda *_args: module_end()) +module = region_op(GPUModuleOp) class GPUModuleMeta(ModuleMeta): @@ -202,12 +202,7 @@ def __prepare__(cls, name, bases, **kwargs): ) ip = InsertionPoint(gpu_module_op.body) ip.__enter__() - return { - "ip": ip, - "loc": loc, - "gpu_module_op": gpu_module_op, - "module_terminator": module_end, - } + return {"ip": ip, "gpu_module_op": gpu_module_op} class GPUFuncOp(GPUFuncOp): diff --git a/mlir/extras/util.py b/mlir/extras/util.py index 68e7d82d..cc6ca846 100644 --- a/mlir/extras/util.py +++ b/mlir/extras/util.py @@ -352,11 +352,7 @@ def region_adder_decorator(op, *args, **kwargs): class ModuleMeta(type): def __new__(cls, name, bases, classdict, **kwargs): ip = classdict.pop("ip") - loc = classdict.pop("loc") - module_terminator = classdict.pop("module_terminator", None) new = super().__new__(cls, name, bases, classdict) - if module_terminator is not None: - module_terminator(loc=loc, ip=ip) for k, v in classdict.items(): if callable(v): v.qualname = name