tests/e2e/convolution/generate_e2e_conv2d_tests.py - 3p/openxla/iree - Git at Google

 #!/usr/bin/env python3
 # Copyright 2024 The IREE Authors
 #
 # Licensed under the Apache License v2.0 with LLVM Exceptions.
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 """Generator for e2e conv2d tests.
 """

 import argparse
 import enum
 import dataclasses
 import typing
 import math


 # Data type of kernel entries. The string values must match MLIR data types.
 @enum.unique
 class KernelElemTypeId(enum.Enum):
     NONE = ""
     F32 = "f32"
     F16 = "f16"


 # Data type of input entries. The string values must match MLIR data types.
 @enum.unique
 class InputElemTypeId(enum.Enum):
     NONE = ""
     F32 = "f32"
     F16 = "f16"


 # Enumerates of the collections of shapes that we can generate tests for.
 # The values are the accepted values for the --shapes= flag.
 @enum.unique
 class ShapesId(enum.Enum):
     SMALL = "small"
     MEDIUM = "medium"
     LARGE = "large"


 # Enumerates ways to construct MLIR tensor types.
 # TODO: Enable dynamic dimensions once the tests start passing.
 @enum.unique
 class Dynamicity(enum.Enum):
     DYNAMIC = "dynamic"  # Use '?' everywhere. Example: tensor<?x?xf32>.
     STATIC = "static"  # Use fixed values everywhere. Example: tensor<4x6xf32>.
     MIXED = "mixed"  # Randomly mix '?' and values. Example: tensor<?x4xf32>.


 # Enumerates ways to initialize input buffer contents.
 @enum.unique
 class InputGenerator(enum.Enum):
     ZERO = "zero"  # Fill with zeros
     RANDOM = "random"  # Fill with (deterministic) pseudorandom values.


 # Enumerates ways to initialize kernel buffer contents.
 @enum.unique
 class KernelGenerator(enum.Enum):
     ZERO = "zero"  # Fill with zeros
     RANDOM = "random"  # Fill with (deterministic) pseudorandom values.


 # TODO: Add more input layouts as needed. The layout determines the dim of input and kernel.
 @enum.unique
 class InputLayout(enum.Enum):
     NCHW = "nchw"
     NHWC = "nhwc"


 # TODO: Add more kernel layouts as needed.
 @enum.unique
 class KernelLayout(enum.Enum):
     FCHW = "fchw"
     HWCF = "hwcf"


 # Describes the shape of a tensor conv2d in the usual convention:
 # the input is {n}x{c}x{h}x{w}, the kernel is {f}x{c}x{kh}x{kw}, the accumulator/result is
 # {n}x{f}x{oh}x{ow}.
 # The extra `accumulate` boolean tells whether the conv2d is accumulating into
 # an existing accumulator (C += A * B) or just overwriting the result
 # (C = A * B).
 @dataclasses.dataclass
 class TestShape:
     n: int
     c: int
     h: int
     w: int
     kh: int
     kw: int
     f: int
     accumulate: bool


 # Attributes for the linalg.conv2d operation.
 @dataclasses.dataclass
 class ConvAttrs:
     STRIDE: typing.Tuple[int, int] = (1, 1)
     DILATION: typing.Tuple[int, int] = (1, 1)


 # Returns the list of TestShape's to use for the collection of shapes
 # identified by shapes_id.
 def get_test_shapes(shapes_id: ShapesId):
     # Notes:
     # 1. Be conservative in adding more shapes, as that can increase both the
     #    build and execution latency of tests. The build latency is nearly the
     #    same for all shapes, while execution latency grows linearly with
     #    n*f*ow*oh*kh*kw.

     if shapes_id == ShapesId.SMALL:
         return [
             TestShape(n=1, c=1, h=1, w=1, kh=1, kw=1, f=1, accumulate=True),
             TestShape(n=1, c=1, h=16, w=16, kh=2, kw=2, f=1, accumulate=True),
             TestShape(n=2, c=2, h=32, w=32, kh=3, kw=3, f=2, accumulate=True),
         ]
     if shapes_id == ShapesId.MEDIUM:
         return [
             TestShape(n=2, h=32, w=32, c=32, kh=3, kw=3, f=64, accumulate=True),
         ]
     if shapes_id == ShapesId.LARGE:
         return [
             TestShape(n=2, c=4, h=128, w=128, kh=3, kw=3, f=8, accumulate=True),
             TestShape(n=2, c=3, h=128, w=128, kh=3, kw=3, f=12, accumulate=True),
         ]

     raise ValueError(shapes_id)


 # Returns the list of Dynamicity's to use for the collection of shapes
 # identified by shapes_id.
 def get_dynamicities(shapes_id: ShapesId):
     if shapes_id == ShapesId.LARGE:
         return [
             Dynamicity.STATIC,
         ]
     # TODO: Enable dynamic dimensions once the tests start passing.
     else:
         return [
             Dynamicity.STATIC,
         ]
     raise ValueError(shapes_id)


 # Intentionally fixed seed! We want full reproducibility here, both across runs
 # and across machines.
 # Intentionally not shared with pseudorandom_generator_seed to limit the ways
 # in which shuffling testcases changes which random values are generated.
 local_pseudorandom_state = 1


 # A shape dimension value, i.e. a size value that could appear in a MLIR type
 # such as 'tensor<?x4xf32>'. None means a dynamic size, similar to '?' in MLIR.
 @dataclasses.dataclass
 class DimSize:
     value: typing.Optional[int]


 # Generates a compile-time MLIR size value, i.e. either a fixed positive integer
 # or None (which maps to MLIR '?') depending on dynamicity.
 def shape_dim(x: int, dynamicity: Dynamicity):
     if dynamicity == Dynamicity.DYNAMIC:
         return DimSize(None)
     elif dynamicity == Dynamicity.STATIC:
         return DimSize(x)
     else:
         raise ValueError(dynamicity)


 # Stringification used for generating MLIR types, e.g. tensor<?x?xf32>.
 def int_or_question_mark(s: DimSize):
     return s.value or "?"


 # Stringification used for generating alphanumeric identifiers, e.g.
 # func.func @somefunction_DYNxDYNxf32, where we can't use "?" characters.
 def int_or_DYN(s: DimSize):
     return s.value or "DYN"


 # Determines the shape of input and kernel tensors.
 @dataclasses.dataclass
 class TestInputTensorShapes:
     n: DimSize
     c: DimSize
     h: DimSize
     w: DimSize
     kh: DimSize
     kw: DimSize
     f: DimSize


 # Helper for generate_function. Generates TestInputTensorShapes, i.e.
 # converts from the runtime shape dimensions in TestShape and given dynamicity to
 # the set of shapes to be used in a test function's input tensors.
 def generate_shapes(shape: TestShape, dynamicity: Dynamicity):
     n = shape_dim(shape.n, dynamicity)
     c = shape_dim(shape.c, dynamicity)
     h = shape_dim(shape.h, dynamicity)
     w = shape_dim(shape.w, dynamicity)
     kh = shape_dim(shape.kh, dynamicity)
     kw = shape_dim(shape.kw, dynamicity)
     f = shape_dim(shape.f, dynamicity)
     shapes = TestInputTensorShapes(
         n=n,
         c=c,
         h=h,
         w=w,
         kh=kh,
         kw=kw,
         f=f,
     )
     return shapes


 # Helper to calculate the output shape based on the input shape, kernel shape,
 # dilation and stride.
 def calc_out_shape(i_shape: int, k_shape: int, dilation_val: int, stride_val: int):
     x = (k_shape - 1) * (dilation_val - 1)
     x = i_shape - k_shape - x
     return math.floor(x / stride_val) + 1


 # Helper to return input, kernel and output shapes based on the layout and Conv2dParams.
 def get_tensor_shape(
     shapes: TestShape,
     kernel_layout: KernelLayout,
     input_layout: InputLayout,
     conv_attr: ConvAttrs,
 ):
     n = shapes.n
     c = shapes.c
     h = shapes.h
     w = shapes.w
     kh = shapes.kh
     kw = shapes.kw
     f = shapes.f

     # Extract input dimensions
     input_height, input_width = h, w

     # Extract kernel dimensions
     kernel_height, kernel_width = kh, kw

     # Get the dilation and stride
     dilation = conv_attr.DILATION
     stride = conv_attr.STRIDE

     # Calculate output height.
     oh = calc_out_shape(input_height, kernel_height, dilation[0], stride[0])
     # Calculate output width.
     ow = calc_out_shape(input_width, kernel_width, dilation[1], stride[1])

     input_tensor_shape, kernel_tensor_shape, output_tensor_shape = [], [], []

     if input_layout == InputLayout.NCHW:
         input_tensor_shape = [n, c, h, w]
         output_tensor_shape = [n, f, oh, ow]
     elif input_layout == InputLayout.NHWC:
         input_tensor_shape = [n, h, w, c]
         output_tensor_shape = [n, oh, ow, f]
     else:
         raise ValueError(input_layout)

     if kernel_layout == KernelLayout.FCHW:
         kernel_tensor_shape = [f, c, kh, kw]
     elif kernel_layout == KernelLayout.HWCF:
         kernel_tensor_shape = [f, c, kh, kw]
     else:
         raise ValueError(kernel_layout)

     return input_tensor_shape, kernel_tensor_shape, output_tensor_shape


 # Helper for generate_function.
 # Generates a name for a test function in the generated MLIR code.
 def generate_function_name(
     input_type: InputElemTypeId,
     kernel_type: KernelElemTypeId,
     output_type: InputElemTypeId,
     shapes: TestInputTensorShapes,
     accumulate: bool,
 ):
     input_t = input_type.value
     kernel_t = kernel_type.value
     acc_t = output_type.value
     n = int_or_DYN(shapes.n)
     c = int_or_DYN(shapes.c)
     h = int_or_DYN(shapes.h)
     w = int_or_DYN(shapes.w)
     kh = int_or_DYN(shapes.kh)
     kw = int_or_DYN(shapes.kw)
     f = int_or_DYN(shapes.f)

     conv2d_kind = "conv2d_accumulate" if accumulate else "conv2d"
     return (
         f"{conv2d_kind}_{n}_{c}_{h}_{w}_times_"
         + f"{kh}_{kw}_{f}_dtype_{input_t}_{kernel_t}_{acc_t}"
     )


 # Represents a generated test function.
 @dataclasses.dataclass
 class MLIRFunction:
     name: str
     signature: str
     import_declaration: str
     definition: str


 # Generates a test function in the generated MLIR code.
 # The generated function will take the same arguments as linalg.conv2d variants
 # and will just call linalg.conv2d variants with them, returning its result.
 def generate_function(
     input_type: InputElemTypeId,
     input_layout: InputLayout,
     kernel_type: KernelElemTypeId,
     kernel_layout: KernelLayout,
     acc_type: InputElemTypeId,
     conv2d_attr: ConvAttrs,
     shape: TestShape,
     dynamicity: Dynamicity,
 ):
     shapes = generate_shapes(shape, dynamicity)
     func_name = generate_function_name(
         input_type,
         kernel_type,
         acc_type,
         shapes,
         shape.accumulate,
     )

     input_shape, kernel_shape, output_shape = get_tensor_shape(
         shape, kernel_layout, input_layout, conv2d_attr
     )
     input_tensor_type = f"tensor<{input_shape[0]}x{input_shape[1]}x{input_shape[2]}x{input_shape[3]}x{input_type.value}>"
     kernel_tensor_type = f"tensor<{kernel_shape[0]}x{kernel_shape[1]}x{kernel_shape[2]}x{kernel_shape[3]}x{kernel_type.value}>"

     acc_tensor_type = f"tensor<{output_shape[0]}x{output_shape[1]}x{output_shape[2]}x{output_shape[3]}x{input_type.value}>"

     op_name = None
     if input_layout == InputLayout.NCHW:
         if kernel_layout == KernelLayout.FCHW:
             op_name = "linalg.conv_2d_nchw_fchw"
         if kernel_layout == KernelLayout.HWCF:
             op_name = "linalg.conv_2d_nchw_hwcf"
     elif input_layout == InputLayout.NHWC:
         if kernel_layout == KernelLayout.HWCF:
             op_name = "linalg.conv_2d_nhwc_hwcf"

     conv_attr = f"{{dilations = dense<{list(conv2d_attr.DILATION)}> : tensor<2xi64>, strides = dense<{list(conv2d_attr.STRIDE)}> : tensor<2xi64>}}"

     # Compilation info is optional; prints empty string by default.
     func_definition = ""

     signature = f"({input_tensor_type}, {kernel_tensor_type}, {acc_tensor_type}) -> {acc_tensor_type}"
     import_declaration = f"func.func private @module.{func_name}(%input: !hal.buffer_view, %kernel: !hal.buffer_view, %acc: !hal.buffer_view) -> !hal.buffer_view"
     func_definition = func_definition + (
         f"func.func @{func_name}(%lhs: {input_tensor_type}, %rhs: {kernel_tensor_type}, %acc: {acc_tensor_type}) -> {acc_tensor_type} {{\n"
         f"  %result = {op_name} {conv_attr} ins(%lhs, %rhs: {input_tensor_type}, {kernel_tensor_type}) outs(%acc: {acc_tensor_type}) -> {acc_tensor_type}\n"
         f"  return %result: {acc_tensor_type}\n"
         f"}}\n"
     )

     return MLIRFunction(
         name=func_name,
         signature=signature,
         import_declaration=import_declaration,
         definition=func_definition,
     )


 # Represents a call to a generated test function.
 @dataclasses.dataclass
 class TestCall:
     function: MLIRFunction
     op: str


 # Enumerates ways to initialize tensor buffer contents.
 @enum.unique
 class TensorGenerator(enum.Enum):
     ZERO = "zero"  # Fill with zeros
     RANDOM = "random"  # Fill with (deterministic) pseudorandom values.


 # Intentionally fixed seed! We want full reproducibility here, both across runs
 # and across machines.
 # Intentionally not shared with local_pseudorandom_state to limit the ways
 # in which shuffling testcases changes which random values are generated.
 pseudorandom_generator_seed = 1


 def contents_generator_tag(generator: TensorGenerator):
     if generator == TensorGenerator.ZERO:
         return ""
     elif generator == TensorGenerator.RANDOM:
         global pseudorandom_generator_seed
         pseudorandom_generator_seed = pseudorandom_generator_seed + 1
         return f"!tag:iree:fully_specified_pseudorandom {pseudorandom_generator_seed}"
     else:
         raise ValueError(generator)


 # Generate a 4d tensor function argument of the given size as `%name`.
 def generate_random_4d_tensor(
     name: str,
     tensor_shape: list,
     element_type: typing.Union[InputElemTypeId, KernelElemTypeId],
 ):
     global pseudorandom_generator_seed
     pseudorandom_generator_seed = pseudorandom_generator_seed + 1
     return (
         f"  %{name}_dim0 = arith.constant {tensor_shape[0]} : i64\n"
         f"  %{name}_dim1 = arith.constant {tensor_shape[1]} : i64\n"
         f"  %{name}_dim2 = arith.constant {tensor_shape[2]} : i64\n"
         f"  %{name}_dim3 = arith.constant {tensor_shape[3]} : i64\n"
         f"  %{name}_element_type = hal.element_type<{element_type.value}> : i32\n"
         f"  %{name}_seed = arith.constant {pseudorandom_generator_seed} : i32\n"
         f"  %{name} = call @conv2d_test.generate_random_tensor(%device, %{name}_dim0, %{name}_dim1, %{name}_dim2, %{name}_dim3, %{name}_element_type, %{name}_seed) : (!hal.device, i64, i64, i64, i64, i32, i32) -> !hal.buffer_view\n"
     )


 call_id = 0


 def generate_call(
     function: MLIRFunction,
     input_type: InputElemTypeId,
     input_layout: InputLayout,
     kernel_type: KernelElemTypeId,
     kernel_layout: KernelLayout,
     conv2d_attr: ConvAttrs,
     acc_type: InputElemTypeId,
     shape: TestShape,
 ):
     global call_id
     func_name = f"{function.name}_{shape.n}_{shape.c}_{shape.h}_{shape.w}_{shape.f}_{shape.kh}_{shape.kw}"
     if shape.accumulate:
         func_name = f"{func_name}_acc"
     func_name = f"{func_name}_{call_id}"
     call_id = call_id + 1

     description = f"Conv2d shape (NxCxHxWxFxKHxKW): {shape.n}x{shape.c}x{shape.h}x{shape.w}x{shape.f}x{shape.kh}x{shape.kw}"
     op = (
         f"func.func @{func_name}() attributes {{\n"
         f'  iree.reflection = {{description = "{description}"}}\n'
         "} {\n"
         "  %device_index = arith.constant 0 : index\n"
         "  %device = hal.devices.get %device_index : !hal.device\n"
     )

     inp_shape, kernel_shape, out_shape = get_tensor_shape(
         shape,
         kernel_layout,
         input_layout,
         conv2d_attr,
     )

     op = op + generate_random_4d_tensor("input", inp_shape, input_type)
     op = op + generate_random_4d_tensor("kernel", kernel_shape, kernel_type)
     if shape.accumulate:
         op = op + generate_random_4d_tensor("acc", out_shape, acc_type)
         # TODO(#16168): there's a bug with in-place input->output aliasing and
         # we work around it here by passing in a unique copy.
         global pseudorandom_generator_seed
         pseudorandom_generator_seed = pseudorandom_generator_seed - 1
         op = op + generate_random_4d_tensor("acc_copy", out_shape, acc_type)
         op = op + (
             f"  %result = call @module.{function.name}(%input, %kernel, %acc_copy) : (!hal.buffer_view, !hal.buffer_view, !hal.buffer_view) -> !hal.buffer_view\n"
         )
     else:
         op = op + (
             f"  %acc = util.null : !hal.buffer_view\n"
             f"  %result = call @module.{function.name}(%input, %kernel) : (!hal.buffer_view, !hal.buffer_view) -> !hal.buffer_view\n"
         )

     op = op + (
         f"  %n = arith.constant {shape.n} : i64\n"
         f"  %c = arith.constant {shape.c} : i64\n"
         f"  %h = arith.constant {shape.h} : i64\n"
         f"  %w = arith.constant {shape.w} : i64\n"
         f"  %f = arith.constant {shape.f} : i64\n"
         f"  %kh = arith.constant {shape.kh} : i64\n"
         f"  %kw = arith.constant {shape.kw} : i64\n"
         f"  %sh = arith.constant {conv2d_attr.STRIDE[0]} : i64\n"
         f"  %sw = arith.constant {conv2d_attr.STRIDE[1]} : i64\n"
         f"  %dh = arith.constant {conv2d_attr.DILATION[0]} : i64\n"
         f"  %dw = arith.constant {conv2d_attr.DILATION[1]} : i64\n"
         f"  call @conv2d_test.check_conv2d_results(%device, %n, %c, %h, %w, %f, %kh, %kw, %sh, %sw, %dh, %dw, %input, %kernel, %acc, %result) : (!hal.device, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view) -> ()\n"
     )

     op = op + "  return\n"
     op = op + "}\n"

     return TestCall(function=function, op=op)


 # Generates all output files' contents as strings.
 def generate(
     input_elem_type: InputElemTypeId,
     input_layout: InputLayout,
     kernel_elem_type: KernelElemTypeId,
     kernel_layout: KernelLayout,
     conv2d_attr: ConvAttrs,
     acc_type: InputElemTypeId,
     shapes_id: ShapesId,
 ):
     functions = {}
     calls = []

     for shape in get_test_shapes(shapes_id):
         for dynamicity in get_dynamicities(shapes_id):
             function = generate_function(
                 input_elem_type,
                 input_layout,
                 kernel_elem_type,
                 kernel_layout,
                 acc_type,
                 conv2d_attr,
                 shape,
                 dynamicity,
             )
             # Different testcases may differ only by runtime parameters but
             # share the same code. For example, dynamic-shapes testcases
             # share the same code involing tensor<?x?xf32> even though the runtime
             # value in the trace are different. That's why we append conditionally
             # to calls, but unconditionally to function_definitions.
             if function.name not in functions:
                 functions[function.name] = function
             calls.append(
                 generate_call(
                     function,
                     input_elem_type,
                     input_layout,
                     kernel_elem_type,
                     kernel_layout,
                     conv2d_attr,
                     acc_type,
                     shape,
                 )
             )

     return (functions, calls)


 def parse_arguments():
     parser = argparse.ArgumentParser(description="Generator of e2e conv2d tests")
     parser.add_argument(
         "--output_conv2d_mlir",
         type=str,
         help="Path of output .mlir file containing the generated conv2d functions",
         required=True,
     )
     parser.add_argument(
         "--output_calls_mlir",
         type=str,
         help="Path of output .mlir file containing the calls",
         required=True,
     )
     parser.add_argument(
         "--input_type",
         type=str,
         choices=["f32", "f16"],
         help="Numeric type of input tensors",
         required=True,
     )
     parser.add_argument(
         "--input_layout",
         type=str,
         default="nchw",
         choices=["nchw", "nhwc"],
         help="Layout of the input tensor. Currently, only nchw is supported.",
         required=False,
     )
     parser.add_argument(
         "--kernel_type",
         type=str,
         choices=["f32", "f16"],
         help="Numeric type of input tensors",
         required=True,
     )
     parser.add_argument(
         "--kernel_layout",
         type=str,
         default="fchw",
         choices=["fchw", "hwcf"],
         help="Layout of kernel tensor. Currently, only fchw is supported.",
         required=False,
     )
     parser.add_argument(
         "--acc_type",
         type=str,
         choices=["f32", "f16"],
         help="Numeric type of input tensors",
         default="",
         required=False,
     )
     parser.add_argument(
         "--shapes",
         type=str,
         choices=[s.value for s in ShapesId],
         help="Collection of tensor shapes to test",
         required=True,
     )
     parser.add_argument(
         "--dilation",
         type=str,
         default="1,1",
         help="The dilation factor for the convolution operation. Comma-separated. As in 1,1",
         required=False,
     )
     parser.add_argument(
         "--stride",
         type=str,
         default="1,1",
         help="The stride factor for the convolution operation. Comma-separated. As in 1,1",
         required=False,
     )
     parser.add_argument(
         "--requirements",
         type=str,
         help="Target requirements for this module. Comma-separated. As in -iree-llvmcpu-target-cpu-features. If the target device does not meet all of the requirements, the test will be skipped.",
         required=False,
     )
     return parser.parse_args()


 def write_code_file(functions, filename):
     with open(filename, "w") as file:
         for function in functions.values():
             file.write(function.definition + "\n")


 def write_calls_file(functions, calls, filename, requirements):
     # Module-level reflection information used to control the test tool.
     reflection = ""
     if requirements:
         reflection = (
             "iree.reflection = {"
             'target_features = "'
             + ",".join([req.lstrip("+") for req in requirements.split(",")])
             + '"'
             "}"
         )
     module_definition = (
         f"builtin.module @calls attributes {{\n" f"  {reflection}\n" f"}} {{\n\n"
     )

     # Declare the custom module that generates arguments.
     module_definition = module_definition + (
         "func.func private @conv2d_test.generate_random_tensor(%device: !hal.device, %dim0: i64, %dim1: i64, %dim2: i64, %dim3: i64, %element_type: i32, %seed: i32) -> !hal.buffer_view\n"
         "func.func private @conv2d_test.check_conv2d_results(%device: !hal.device, %n: i64, %c: i64, %h: i64, %w: i64, %f:i64, %kh:i64, %kw:i64, %sh:i64, %sw:i64, %dh:i64, %dw:i64, %input: !hal.buffer_view, %kernel: !hal.buffer_view, %acc: !hal.buffer_view, %actual_result: !hal.buffer_view)\n"
         "\n"
     )

     # Declare the functions that will be called.
     for function in functions.values():
         module_definition = module_definition + function.import_declaration + "\n"
     module_definition = module_definition + "\n"

     # Emit the test cases for each call.
     for call in calls:
         module_definition = module_definition + call.op + "\n"

     module_definition = module_definition + "\n}\n"

     with open(filename, "w") as file:
         file.write(module_definition)


 def main(args):
     input_type = InputElemTypeId(args.input_type)
     input_layout = InputLayout(args.input_layout)
     kernel_type = KernelElemTypeId(args.kernel_type)
     kernel_layout = KernelLayout(args.kernel_layout)
     # TODO: The output type is same as the input type for now.
     acc_type = input_type
     shapes_id = ShapesId(args.shapes)
     conv2d_attr = ConvAttrs(
         tuple(map(int, args.stride.split(","))),
         tuple(map(int, args.dilation.split(","))),
     )

     (functions, calls) = generate(
         input_type,
         input_layout,
         kernel_type,
         kernel_layout,
         conv2d_attr,
         acc_type,
         shapes_id,
     )

     write_code_file(functions, args.output_conv2d_mlir)
     write_calls_file(
         functions,
         calls,
         args.output_calls_mlir,
         args.requirements,
     )


 if __name__ == "__main__":
     main(parse_arguments())
	#!/usr/bin/env python3
	# Copyright 2024 The IREE Authors
	#
	# Licensed under the Apache License v2.0 with LLVM Exceptions.
	# See https://llvm.org/LICENSE.txt for license information.
	# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
	"""Generator for e2e conv2d tests.
	"""

	import argparse
	import enum
	import dataclasses
	import typing
	import math


	# Data type of kernel entries. The string values must match MLIR data types.
	@enum.unique
	class KernelElemTypeId(enum.Enum):
	NONE = ""
	F32 = "f32"
	F16 = "f16"


	# Data type of input entries. The string values must match MLIR data types.
	@enum.unique
	class InputElemTypeId(enum.Enum):
	NONE = ""
	F32 = "f32"
	F16 = "f16"


	# Enumerates of the collections of shapes that we can generate tests for.
	# The values are the accepted values for the --shapes= flag.
	@enum.unique
	class ShapesId(enum.Enum):
	SMALL = "small"
	MEDIUM = "medium"
	LARGE = "large"


	# Enumerates ways to construct MLIR tensor types.
	# TODO: Enable dynamic dimensions once the tests start passing.
	@enum.unique
	class Dynamicity(enum.Enum):
	DYNAMIC = "dynamic" # Use '?' everywhere. Example: tensor<?x?xf32>.
	STATIC = "static" # Use fixed values everywhere. Example: tensor<4x6xf32>.
	MIXED = "mixed" # Randomly mix '?' and values. Example: tensor<?x4xf32>.


	# Enumerates ways to initialize input buffer contents.
	@enum.unique
	class InputGenerator(enum.Enum):
	ZERO = "zero" # Fill with zeros
	RANDOM = "random" # Fill with (deterministic) pseudorandom values.


	# Enumerates ways to initialize kernel buffer contents.
	@enum.unique
	class KernelGenerator(enum.Enum):
	ZERO = "zero" # Fill with zeros
	RANDOM = "random" # Fill with (deterministic) pseudorandom values.


	# TODO: Add more input layouts as needed. The layout determines the dim of input and kernel.
	@enum.unique
	class InputLayout(enum.Enum):
	NCHW = "nchw"
	NHWC = "nhwc"


	# TODO: Add more kernel layouts as needed.
	@enum.unique
	class KernelLayout(enum.Enum):
	FCHW = "fchw"
	HWCF = "hwcf"


	# Describes the shape of a tensor conv2d in the usual convention:
	# the input is {n}x{c}x{h}x{w}, the kernel is {f}x{c}x{kh}x{kw}, the accumulator/result is
	# {n}x{f}x{oh}x{ow}.
	# The extra `accumulate` boolean tells whether the conv2d is accumulating into
	# an existing accumulator (C += A * B) or just overwriting the result
	# (C = A * B).
	@dataclasses.dataclass
	class TestShape:
	n: int
	c: int
	h: int
	w: int
	kh: int
	kw: int
	f: int
	accumulate: bool


	# Attributes for the linalg.conv2d operation.
	@dataclasses.dataclass
	class ConvAttrs:
	STRIDE: typing.Tuple[int, int] = (1, 1)
	DILATION: typing.Tuple[int, int] = (1, 1)


	# Returns the list of TestShape's to use for the collection of shapes
	# identified by shapes_id.
	def get_test_shapes(shapes_id: ShapesId):
	# Notes:
	# 1. Be conservative in adding more shapes, as that can increase both the
	# build and execution latency of tests. The build latency is nearly the
	# same for all shapes, while execution latency grows linearly with
	# nfowohkh*kw.

	if shapes_id == ShapesId.SMALL:
	return [
	TestShape(n=1, c=1, h=1, w=1, kh=1, kw=1, f=1, accumulate=True),
	TestShape(n=1, c=1, h=16, w=16, kh=2, kw=2, f=1, accumulate=True),
	TestShape(n=2, c=2, h=32, w=32, kh=3, kw=3, f=2, accumulate=True),
	]
	if shapes_id == ShapesId.MEDIUM:
	return [
	TestShape(n=2, h=32, w=32, c=32, kh=3, kw=3, f=64, accumulate=True),
	]
	if shapes_id == ShapesId.LARGE:
	return [
	TestShape(n=2, c=4, h=128, w=128, kh=3, kw=3, f=8, accumulate=True),
	TestShape(n=2, c=3, h=128, w=128, kh=3, kw=3, f=12, accumulate=True),
	]

	raise ValueError(shapes_id)


	# Returns the list of Dynamicity's to use for the collection of shapes
	# identified by shapes_id.
	def get_dynamicities(shapes_id: ShapesId):
	if shapes_id == ShapesId.LARGE:
	return [
	Dynamicity.STATIC,
	]
	# TODO: Enable dynamic dimensions once the tests start passing.
	else:
	return [
	Dynamicity.STATIC,
	]
	raise ValueError(shapes_id)


	# Intentionally fixed seed! We want full reproducibility here, both across runs
	# and across machines.
	# Intentionally not shared with pseudorandom_generator_seed to limit the ways
	# in which shuffling testcases changes which random values are generated.
	local_pseudorandom_state = 1


	# A shape dimension value, i.e. a size value that could appear in a MLIR type
	# such as 'tensor<?x4xf32>'. None means a dynamic size, similar to '?' in MLIR.
	@dataclasses.dataclass
	class DimSize:
	value: typing.Optional[int]


	# Generates a compile-time MLIR size value, i.e. either a fixed positive integer
	# or None (which maps to MLIR '?') depending on dynamicity.
	def shape_dim(x: int, dynamicity: Dynamicity):
	if dynamicity == Dynamicity.DYNAMIC:
	return DimSize(None)
	elif dynamicity == Dynamicity.STATIC:
	return DimSize(x)
	else:
	raise ValueError(dynamicity)


	# Stringification used for generating MLIR types, e.g. tensor<?x?xf32>.
	def int_or_question_mark(s: DimSize):
	return s.value or "?"


	# Stringification used for generating alphanumeric identifiers, e.g.
	# func.func @somefunction_DYNxDYNxf32, where we can't use "?" characters.
	def int_or_DYN(s: DimSize):
	return s.value or "DYN"


	# Determines the shape of input and kernel tensors.
	@dataclasses.dataclass
	class TestInputTensorShapes:
	n: DimSize
	c: DimSize
	h: DimSize
	w: DimSize
	kh: DimSize
	kw: DimSize
	f: DimSize


	# Helper for generate_function. Generates TestInputTensorShapes, i.e.
	# converts from the runtime shape dimensions in TestShape and given dynamicity to
	# the set of shapes to be used in a test function's input tensors.
	def generate_shapes(shape: TestShape, dynamicity: Dynamicity):
	n = shape_dim(shape.n, dynamicity)
	c = shape_dim(shape.c, dynamicity)
	h = shape_dim(shape.h, dynamicity)
	w = shape_dim(shape.w, dynamicity)
	kh = shape_dim(shape.kh, dynamicity)
	kw = shape_dim(shape.kw, dynamicity)
	f = shape_dim(shape.f, dynamicity)
	shapes = TestInputTensorShapes(
	n=n,
	c=c,
	h=h,
	w=w,
	kh=kh,
	kw=kw,
	f=f,
	)
	return shapes


	# Helper to calculate the output shape based on the input shape, kernel shape,
	# dilation and stride.
	def calc_out_shape(i_shape: int, k_shape: int, dilation_val: int, stride_val: int):
	x = (k_shape - 1) * (dilation_val - 1)
	x = i_shape - k_shape - x
	return math.floor(x / stride_val) + 1


	# Helper to return input, kernel and output shapes based on the layout and Conv2dParams.
	def get_tensor_shape(
	shapes: TestShape,
	kernel_layout: KernelLayout,
	input_layout: InputLayout,
	conv_attr: ConvAttrs,
	):
	n = shapes.n
	c = shapes.c
	h = shapes.h
	w = shapes.w
	kh = shapes.kh
	kw = shapes.kw
	f = shapes.f

	# Extract input dimensions
	input_height, input_width = h, w

	# Extract kernel dimensions
	kernel_height, kernel_width = kh, kw

	# Get the dilation and stride
	dilation = conv_attr.DILATION
	stride = conv_attr.STRIDE

	# Calculate output height.
	oh = calc_out_shape(input_height, kernel_height, dilation[0], stride[0])
	# Calculate output width.
	ow = calc_out_shape(input_width, kernel_width, dilation[1], stride[1])

	input_tensor_shape, kernel_tensor_shape, output_tensor_shape = [], [], []

	if input_layout == InputLayout.NCHW:
	input_tensor_shape = [n, c, h, w]
	output_tensor_shape = [n, f, oh, ow]
	elif input_layout == InputLayout.NHWC:
	input_tensor_shape = [n, h, w, c]
	output_tensor_shape = [n, oh, ow, f]
	else:
	raise ValueError(input_layout)

	if kernel_layout == KernelLayout.FCHW:
	kernel_tensor_shape = [f, c, kh, kw]
	elif kernel_layout == KernelLayout.HWCF:
	kernel_tensor_shape = [f, c, kh, kw]
	else:
	raise ValueError(kernel_layout)

	return input_tensor_shape, kernel_tensor_shape, output_tensor_shape


	# Helper for generate_function.
	# Generates a name for a test function in the generated MLIR code.
	def generate_function_name(
	input_type: InputElemTypeId,
	kernel_type: KernelElemTypeId,
	output_type: InputElemTypeId,
	shapes: TestInputTensorShapes,
	accumulate: bool,
	):
	input_t = input_type.value
	kernel_t = kernel_type.value
	acc_t = output_type.value
	n = int_or_DYN(shapes.n)
	c = int_or_DYN(shapes.c)
	h = int_or_DYN(shapes.h)
	w = int_or_DYN(shapes.w)
	kh = int_or_DYN(shapes.kh)
	kw = int_or_DYN(shapes.kw)
	f = int_or_DYN(shapes.f)

	conv2d_kind = "conv2d_accumulate" if accumulate else "conv2d"
	return (
	f"{conv2d_kind}_{n}_{c}_{h}_{w}_times_"
	+ f"{kh}_{kw}_{f}_dtype_{input_t}_{kernel_t}_{acc_t}"
	)


	# Represents a generated test function.
	@dataclasses.dataclass
	class MLIRFunction:
	name: str
	signature: str
	import_declaration: str
	definition: str


	# Generates a test function in the generated MLIR code.
	# The generated function will take the same arguments as linalg.conv2d variants
	# and will just call linalg.conv2d variants with them, returning its result.
	def generate_function(
	input_type: InputElemTypeId,
	input_layout: InputLayout,
	kernel_type: KernelElemTypeId,
	kernel_layout: KernelLayout,
	acc_type: InputElemTypeId,
	conv2d_attr: ConvAttrs,
	shape: TestShape,
	dynamicity: Dynamicity,
	):
	shapes = generate_shapes(shape, dynamicity)
	func_name = generate_function_name(
	input_type,
	kernel_type,
	acc_type,
	shapes,
	shape.accumulate,
	)

	input_shape, kernel_shape, output_shape = get_tensor_shape(
	shape, kernel_layout, input_layout, conv2d_attr
	)
	input_tensor_type = f"tensor<{input_shape[0]}x{input_shape[1]}x{input_shape[2]}x{input_shape[3]}x{input_type.value}>"
	kernel_tensor_type = f"tensor<{kernel_shape[0]}x{kernel_shape[1]}x{kernel_shape[2]}x{kernel_shape[3]}x{kernel_type.value}>"

	acc_tensor_type = f"tensor<{output_shape[0]}x{output_shape[1]}x{output_shape[2]}x{output_shape[3]}x{input_type.value}>"

	op_name = None
	if input_layout == InputLayout.NCHW:
	if kernel_layout == KernelLayout.FCHW:
	op_name = "linalg.conv_2d_nchw_fchw"
	if kernel_layout == KernelLayout.HWCF:
	op_name = "linalg.conv_2d_nchw_hwcf"
	elif input_layout == InputLayout.NHWC:
	if kernel_layout == KernelLayout.HWCF:
	op_name = "linalg.conv_2d_nhwc_hwcf"

	conv_attr = f"{{dilations = dense<{list(conv2d_attr.DILATION)}> : tensor<2xi64>, strides = dense<{list(conv2d_attr.STRIDE)}> : tensor<2xi64>}}"

	# Compilation info is optional; prints empty string by default.
	func_definition = ""

	signature = f"({input_tensor_type}, {kernel_tensor_type}, {acc_tensor_type}) -> {acc_tensor_type}"
	import_declaration = f"func.func private @module.{func_name}(%input: !hal.buffer_view, %kernel: !hal.buffer_view, %acc: !hal.buffer_view) -> !hal.buffer_view"
	func_definition = func_definition + (
	f"func.func @{func_name}(%lhs: {input_tensor_type}, %rhs: {kernel_tensor_type}, %acc: {acc_tensor_type}) -> {acc_tensor_type} {{\n"
	f" %result = {op_name} {conv_attr} ins(%lhs, %rhs: {input_tensor_type}, {kernel_tensor_type}) outs(%acc: {acc_tensor_type}) -> {acc_tensor_type}\n"
	f" return %result: {acc_tensor_type}\n"
	f"}}\n"
	)

	return MLIRFunction(
	name=func_name,
	signature=signature,
	import_declaration=import_declaration,
	definition=func_definition,
	)


	# Represents a call to a generated test function.
	@dataclasses.dataclass
	class TestCall:
	function: MLIRFunction
	op: str


	# Enumerates ways to initialize tensor buffer contents.
	@enum.unique
	class TensorGenerator(enum.Enum):
	ZERO = "zero" # Fill with zeros
	RANDOM = "random" # Fill with (deterministic) pseudorandom values.


	# Intentionally fixed seed! We want full reproducibility here, both across runs
	# and across machines.
	# Intentionally not shared with local_pseudorandom_state to limit the ways
	# in which shuffling testcases changes which random values are generated.
	pseudorandom_generator_seed = 1


	def contents_generator_tag(generator: TensorGenerator):
	if generator == TensorGenerator.ZERO:
	return ""
	elif generator == TensorGenerator.RANDOM:
	global pseudorandom_generator_seed
	pseudorandom_generator_seed = pseudorandom_generator_seed + 1
	return f"!tag:iree:fully_specified_pseudorandom {pseudorandom_generator_seed}"
	else:
	raise ValueError(generator)


	# Generate a 4d tensor function argument of the given size as `%name`.
	def generate_random_4d_tensor(
	name: str,
	tensor_shape: list,
	element_type: typing.Union[InputElemTypeId, KernelElemTypeId],
	):
	global pseudorandom_generator_seed
	pseudorandom_generator_seed = pseudorandom_generator_seed + 1
	return (
	f" %{name}_dim0 = arith.constant {tensor_shape[0]} : i64\n"
	f" %{name}_dim1 = arith.constant {tensor_shape[1]} : i64\n"
	f" %{name}_dim2 = arith.constant {tensor_shape[2]} : i64\n"
	f" %{name}_dim3 = arith.constant {tensor_shape[3]} : i64\n"
	f" %{name}_element_type = hal.element_type<{element_type.value}> : i32\n"
	f" %{name}_seed = arith.constant {pseudorandom_generator_seed} : i32\n"
	f" %{name} = call @conv2d_test.generate_random_tensor(%device, %{name}_dim0, %{name}_dim1, %{name}_dim2, %{name}_dim3, %{name}_element_type, %{name}_seed) : (!hal.device, i64, i64, i64, i64, i32, i32) -> !hal.buffer_view\n"
	)


	call_id = 0


	def generate_call(
	function: MLIRFunction,
	input_type: InputElemTypeId,
	input_layout: InputLayout,
	kernel_type: KernelElemTypeId,
	kernel_layout: KernelLayout,
	conv2d_attr: ConvAttrs,
	acc_type: InputElemTypeId,
	shape: TestShape,
	):
	global call_id
	func_name = f"{function.name}_{shape.n}_{shape.c}_{shape.h}_{shape.w}_{shape.f}_{shape.kh}_{shape.kw}"
	if shape.accumulate:
	func_name = f"{func_name}_acc"
	func_name = f"{func_name}_{call_id}"
	call_id = call_id + 1

	description = f"Conv2d shape (NxCxHxWxFxKHxKW): {shape.n}x{shape.c}x{shape.h}x{shape.w}x{shape.f}x{shape.kh}x{shape.kw}"
	op = (
	f"func.func @{func_name}() attributes {{\n"
	f' iree.reflection = {{description = "{description}"}}\n'
	"} {\n"
	" %device_index = arith.constant 0 : index\n"
	" %device = hal.devices.get %device_index : !hal.device\n"
	)

	inp_shape, kernel_shape, out_shape = get_tensor_shape(
	shape,
	kernel_layout,
	input_layout,
	conv2d_attr,
	)

	op = op + generate_random_4d_tensor("input", inp_shape, input_type)
	op = op + generate_random_4d_tensor("kernel", kernel_shape, kernel_type)
	if shape.accumulate:
	op = op + generate_random_4d_tensor("acc", out_shape, acc_type)
	# TODO(#16168): there's a bug with in-place input->output aliasing and
	# we work around it here by passing in a unique copy.
	global pseudorandom_generator_seed
	pseudorandom_generator_seed = pseudorandom_generator_seed - 1
	op = op + generate_random_4d_tensor("acc_copy", out_shape, acc_type)
	op = op + (
	f" %result = call @module.{function.name}(%input, %kernel, %acc_copy) : (!hal.buffer_view, !hal.buffer_view, !hal.buffer_view) -> !hal.buffer_view\n"
	)
	else:
	op = op + (
	f" %acc = util.null : !hal.buffer_view\n"
	f" %result = call @module.{function.name}(%input, %kernel) : (!hal.buffer_view, !hal.buffer_view) -> !hal.buffer_view\n"
	)

	op = op + (
	f" %n = arith.constant {shape.n} : i64\n"
	f" %c = arith.constant {shape.c} : i64\n"
	f" %h = arith.constant {shape.h} : i64\n"
	f" %w = arith.constant {shape.w} : i64\n"
	f" %f = arith.constant {shape.f} : i64\n"
	f" %kh = arith.constant {shape.kh} : i64\n"
	f" %kw = arith.constant {shape.kw} : i64\n"
	f" %sh = arith.constant {conv2d_attr.STRIDE[0]} : i64\n"
	f" %sw = arith.constant {conv2d_attr.STRIDE[1]} : i64\n"
	f" %dh = arith.constant {conv2d_attr.DILATION[0]} : i64\n"
	f" %dw = arith.constant {conv2d_attr.DILATION[1]} : i64\n"
	f" call @conv2d_test.check_conv2d_results(%device, %n, %c, %h, %w, %f, %kh, %kw, %sh, %sw, %dh, %dw, %input, %kernel, %acc, %result) : (!hal.device, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view) -> ()\n"
	)

	op = op + " return\n"
	op = op + "}\n"

	return TestCall(function=function, op=op)


	# Generates all output files' contents as strings.
	def generate(
	input_elem_type: InputElemTypeId,
	input_layout: InputLayout,
	kernel_elem_type: KernelElemTypeId,
	kernel_layout: KernelLayout,
	conv2d_attr: ConvAttrs,
	acc_type: InputElemTypeId,
	shapes_id: ShapesId,
	):
	functions = {}
	calls = []

	for shape in get_test_shapes(shapes_id):
	for dynamicity in get_dynamicities(shapes_id):
	function = generate_function(
	input_elem_type,
	input_layout,
	kernel_elem_type,
	kernel_layout,
	acc_type,
	conv2d_attr,
	shape,
	dynamicity,
	)
	# Different testcases may differ only by runtime parameters but
	# share the same code. For example, dynamic-shapes testcases
	# share the same code involing tensor<?x?xf32> even though the runtime
	# value in the trace are different. That's why we append conditionally
	# to calls, but unconditionally to function_definitions.
	if function.name not in functions:
	functions[function.name] = function
	calls.append(
	generate_call(
	function,
	input_elem_type,
	input_layout,
	kernel_elem_type,
	kernel_layout,
	conv2d_attr,
	acc_type,
	shape,
	)
	)

	return (functions, calls)


	def parse_arguments():
	parser = argparse.ArgumentParser(description="Generator of e2e conv2d tests")
	parser.add_argument(
	"--output_conv2d_mlir",
	type=str,
	help="Path of output .mlir file containing the generated conv2d functions",
	required=True,
	)
	parser.add_argument(
	"--output_calls_mlir",
	type=str,
	help="Path of output .mlir file containing the calls",
	required=True,
	)
	parser.add_argument(
	"--input_type",
	type=str,
	choices=["f32", "f16"],
	help="Numeric type of input tensors",
	required=True,
	)
	parser.add_argument(
	"--input_layout",
	type=str,
	default="nchw",
	choices=["nchw", "nhwc"],
	help="Layout of the input tensor. Currently, only nchw is supported.",
	required=False,
	)
	parser.add_argument(
	"--kernel_type",
	type=str,
	choices=["f32", "f16"],
	help="Numeric type of input tensors",
	required=True,
	)
	parser.add_argument(
	"--kernel_layout",
	type=str,
	default="fchw",
	choices=["fchw", "hwcf"],
	help="Layout of kernel tensor. Currently, only fchw is supported.",
	required=False,
	)
	parser.add_argument(
	"--acc_type",
	type=str,
	choices=["f32", "f16"],
	help="Numeric type of input tensors",
	default="",
	required=False,
	)
	parser.add_argument(
	"--shapes",
	type=str,
	choices=[s.value for s in ShapesId],
	help="Collection of tensor shapes to test",
	required=True,
	)
	parser.add_argument(
	"--dilation",
	type=str,
	default="1,1",
	help="The dilation factor for the convolution operation. Comma-separated. As in 1,1",
	required=False,
	)
	parser.add_argument(
	"--stride",
	type=str,
	default="1,1",
	help="The stride factor for the convolution operation. Comma-separated. As in 1,1",
	required=False,
	)
	parser.add_argument(
	"--requirements",
	type=str,
	help="Target requirements for this module. Comma-separated. As in -iree-llvmcpu-target-cpu-features. If the target device does not meet all of the requirements, the test will be skipped.",
	required=False,
	)
	return parser.parse_args()


	def write_code_file(functions, filename):
	with open(filename, "w") as file:
	for function in functions.values():
	file.write(function.definition + "\n")


	def write_calls_file(functions, calls, filename, requirements):
	# Module-level reflection information used to control the test tool.
	reflection = ""
	if requirements:
	reflection = (
	"iree.reflection = {"
	'target_features = "'
	+ ",".join([req.lstrip("+") for req in requirements.split(",")])
	+ '"'
	"}"
	)
	module_definition = (
	f"builtin.module @calls attributes {{\n" f" {reflection}\n" f"}} {{\n\n"
	)

	# Declare the custom module that generates arguments.
	module_definition = module_definition + (
	"func.func private @conv2d_test.generate_random_tensor(%device: !hal.device, %dim0: i64, %dim1: i64, %dim2: i64, %dim3: i64, %element_type: i32, %seed: i32) -> !hal.buffer_view\n"
	"func.func private @conv2d_test.check_conv2d_results(%device: !hal.device, %n: i64, %c: i64, %h: i64, %w: i64, %f:i64, %kh:i64, %kw:i64, %sh:i64, %sw:i64, %dh:i64, %dw:i64, %input: !hal.buffer_view, %kernel: !hal.buffer_view, %acc: !hal.buffer_view, %actual_result: !hal.buffer_view)\n"
	"\n"
	)

	# Declare the functions that will be called.
	for function in functions.values():
	module_definition = module_definition + function.import_declaration + "\n"
	module_definition = module_definition + "\n"

	# Emit the test cases for each call.
	for call in calls:
	module_definition = module_definition + call.op + "\n"

	module_definition = module_definition + "\n}\n"

	with open(filename, "w") as file:
	file.write(module_definition)


	def main(args):
	input_type = InputElemTypeId(args.input_type)
	input_layout = InputLayout(args.input_layout)
	kernel_type = KernelElemTypeId(args.kernel_type)
	kernel_layout = KernelLayout(args.kernel_layout)
	# TODO: The output type is same as the input type for now.
	acc_type = input_type
	shapes_id = ShapesId(args.shapes)
	conv2d_attr = ConvAttrs(
	tuple(map(int, args.stride.split(","))),
	tuple(map(int, args.dilation.split(","))),
	)

	(functions, calls) = generate(
	input_type,
	input_layout,
	kernel_type,
	kernel_layout,
	conv2d_attr,
	acc_type,
	shapes_id,
	)

	write_code_file(functions, args.output_conv2d_mlir)
	write_calls_file(
	functions,
	calls,
	args.output_calls_mlir,
	args.requirements,
	)


	if __name__ == "__main__":
	main(parse_arguments())