feat: more constant checks

Author: avik-pal

src/Ops.jl

@@ -206,6 +206,12 @@ for (T, mlir_func) in (
 
             splatattr = MLIR.API.$mlir_func(tt, number)
             cst_op = stablehlo.constant(; output=tt, value=splatattr, location=location)
+
+            parent_func_op = MLIR.IR.get_parent_of_type_function_op(cst_op)
+            if parent_func_op == C_NULL
+                error("Constant must be created inside a Function Op.")
+            end
+
             cst = MLIR.IR.result(cst_op)
             ta = TracedRArray{$T,length(shape)}((), cst, shape)
             return ta
@@ -226,6 +232,12 @@ end
     tt = MLIR.IR.TensorType(shape, MLIR.IR.Type(T))
     splatattr = MLIR.API.mlirDenseElementsAttrSplatGet(tt, _fill_element_attr(element))
     cst_op = stablehlo.constant(; output=tt, value=splatattr, location=location)
+
+    parent_func_op = MLIR.IR.get_parent_of_type_function_op(cst_op)
+    if parent_func_op == C_NULL
+        error("Constant must be created inside a Function Op.")
+    end
+
     cst = MLIR.IR.result(cst_op)
     ta = TracedRArray{T,length(shape)}((), cst, shape)
     return ta

src/mlir/IR/Operation.jl

@@ -68,6 +68,12 @@ Gets the operation that owns this operation, returning null if the operation is
 parent_op(operation::Operation) =
     Operation(API.mlirOperationGetParentOperation(operation), false)
 
+"""
+    parent_region(op)
+Gets the region that owns this operation.
+"""
+parent_region(operation::Operation) = parent_region(block(operation))
+
 """
     rmfromparent!(op)
 
@@ -333,13 +339,14 @@ 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
+    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)
+            operand_region = parent_region.(operands)
+            # TODO: add the checks
+        end
+    end
     return op
 end
 

src/mlir/IR/Value.jl

@@ -121,3 +121,5 @@ function Base.show(io::IO, value::Value)
         API.mlirValuePrint(value, c_print_callback, ref)
     end
 end
+
+parent_region(value::Value) = parent_region(owner(value))