feat: early fail if not correct region

Author: avik-pal

deps/ReactantExtra/API.cpp

@@ -2404,3 +2404,7 @@ extern "C" void dump_operation(Operation *op, const char *filename) {
 extern "C" bool pjrt_device_is_addressable(PjRtDevice *device) {
   return device->IsAddressable();
 }
+
+extern "C" mlir::Operation *mlirGetParentOfTypeFunctionOp(mlir::Operation *op) {
+  return op->getParentOfType<mlir::FunctionOpInterface>();
+}

src/Ops.jl

@@ -126,6 +126,11 @@ end
             result_inference=false,
         )
 
+        parent_func_op = MLIR.IR.get_parent_of_type_function_op(cstop)
+        if parent_func_op == C_NULL
+            error("Constant must be created inside a Function Op.")
+        end
+
         res = MLIR.IR.result(cstop)
         tres = TracedRArray{T,N}((), res, size(x))
         constants[value] = tres

src/mlir/IR/Operation.jl

@@ -331,16 +331,38 @@ function create_operation_common(
     end
 end
 
+function create_operation_common_with_checks(args...; operands=nothing, kwargs...)
+    op = create_operation_common(args...; operands, kwargs...)
+    # if !isnothing(operands)
+    #     parent_function_op = get_parent_of_type_function_op(op)
+    #     if parent_function_op != C_NULL
+    #         function_op_region = parent_region(parent_function_op)
+    #         # TODO: add the checks
+    #     end
+    # end
+    return op
+end
+
 function create_operation(args...; kwargs...)
-    res = create_operation_common(args...; kwargs...)
+    res = create_operation_common_with_checks(args...; kwargs...)
     if _has_block()
         push!(block(), res)
     end
     return res
 end
 
 function create_operation_at_front(args...; kwargs...)
-    res = create_operation_common(args...; kwargs...)
+    res = create_operation_common_with_checks(args...; kwargs...)
     Base.pushfirst!(block(), res)
     return res
 end
+
+function get_parent_of_type_function_op(op::Operation)
+    GC.@preserve op begin
+        funcop = @ccall API.mlir_c.mlirGetParentOfTypeFunctionOp(
+            op::API.MlirOperation
+        )::API.MlirOperation
+    end
+    funcop.ptr == C_NULL && return C_NULL
+    return Operation(funcop, false)
+end