Convert ml program to memref

class ConvertGlobalPattern(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(
        self, op: ml_program.GlobalOp, rewriter: PatternRewriter
    ) -> None:
        if op.value is None:
            raise NotImplementedError(
                "Converting ml_program.global with no value not implemented"
            )
        op_type = op.type
        assert isinstance(op_type, TensorType)
        op_type = cast(TensorType[Any], op_type)
        new_type = memref.MemRefType(op_type.element_type, op_type.shape)
        rewriter.replace(
            op,
            (
                memref.GlobalOp.get(
                    op.sym_name,
                    new_type,
                    op.value,
                    op.sym_visibility,
                    UnitAttr() if op.is_mutable is None else None,
                ),
            ),
        )

class ConvertGlobalLoadConst(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(
        self, op: ml_program.GlobalLoadConstantOp, rewriter: PatternRewriter
    ) -> None:
        op_type = op.result.type
        assert isinstance(op_type, TensorType)
        op_type = cast(TensorType[Any], op_type)
        new_type = memref.MemRefType(op_type.element_type, op_type.shape)
        rewriter.replace(
            op,
            (
                mem := memref.GetGlobalOp(op.global_attr, new_type),
                bufferization.ToTensorOp(mem.memref),
            ),
        )

class ConvertMlProgramToMemRefPass(ModulePass):
    """
    Converts operations in the `ml_program` dialect to `memref`.
    `ml_program` operations are at the `tensor` level of abstraction, so some of the
    rewrites insert `bufferization` ops to bridge the gap to existing consumers of global
    `tensor`s.
    """

    name = "convert-ml-program-to-memref"

    def apply(self, ctx: Context, op: ModuleOp) -> None:
        PatternRewriteWalker(
            GreedyRewritePatternApplier(
                [
                    ConvertGlobalPattern(),
                    ConvertGlobalLoadConst(),
                ]
            ),
            apply_recursively=False,
        ).rewrite_module(op)