tests/e2e/matmul/common.py - 3p/openxla/iree - Git at Google


            
            #!/usr/bin/env python3

            
            # Copyright 2025 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

            
            """common utilities and data structures for e2e matmul tests"""


            import enum

            
            import dataclasses

            
            import re

            
            import typing


            # Data type of matrix entries. The string values must match MLIR data types.

            
            # This is a superset of the values accepted for the --lhs_rhs_types= flag,

            
            # as this also includes accumulator-specific types like i32.

            
            @enum.unique

            
            class MatrixElemTypeId(enum.Enum):

            
                NONE = ""

            
                UI8 = "ui8"

            
                I8 = "i8"

            
                I32 = "i32"

            
                F64 = "f64"

            
                F32 = "f32"

            
                F16 = "f16"

            
                BF16 = "bf16"

            
                F8E5M2 = "f8E5M2"

            
                F8E4M3FN = "f8E4M3FN"

            
                F8E5M2FNUZ = "f8E5M2FNUZ"

            
                F8E4M3FNUZ = "f8E4M3FNUZ"

            
                F8E8M0FNU = "f8E8M0FNU"

            
                F6E3M2FN = "f6E3M2FN"

            
                F6E2M3FN = "f6E2M3FN"

            
                F4E2M1FN = "f4E2M1FN"


            def get_size_in_bits(type_id: MatrixElemTypeId):

            
                return int(re.search(r"\d+", str(type_id)).group())


            # 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):

            
                DEFAULT = "default"

            
                SMALL = "small"

            
                LARGE = "large"

            
                EASY_LARGE_STATIC = "easy_large_static"

            
                CUSTOM_MNK = "custom_mnk"  # Used for custom shapes specified by --mnk= flag.


                @classmethod

            
                def set_custom_mnk(cls, shapes_id, mnk_string):

            
                    """Parse and set custom MNK values from command line argument."""

            
                    if shapes_id != cls.CUSTOM_MNK:

            
                        if mnk_string:

            
                            raise ValueError("--mnk can only be used with --shapes=custom_mnk")

            
                        return


                    # shapes_id is CUSTOM_MNK

            
                    if not mnk_string:

            
                        raise ValueError("--mnk must be specified when using --shapes=custom_mnk")

            
                    try:

            
                        mnk_parts = mnk_string.split(",")

            
                        if len(mnk_parts) != 3:

            
                            raise ValueError("--mnk must have exactly 3 values: m,n,k")

            
                        cls.custom_mnk_values = tuple(int(x) for x in mnk_parts)

            
                    except ValueError as e:

            
                        raise ValueError(f"Invalid --mnk format: {e}")


                @classmethod

            
                def set_dynamicity(cls, shapes_id, dynamicity_string):

            
                    """Parse and set MNK dynamicity from command line argument."""

            
                    if shapes_id != cls.CUSTOM_MNK:

            
                        if dynamicity_string:

            
                            raise ValueError(

            
                                "--mnk_dynamicities can only be used with --shapes=custom_mnk"

            
                            )

            
                        return


                    if dynamicity_string:

            
                        try:

            
                            dynamicity_parts = dynamicity_string.split(",")

            
                            if len(dynamicity_parts) != 3:

            
                                raise ValueError(

            
                                    "--mnk_dynamicities must have exactly 3 values, each being either 'dynamic' or 'static'"

            
                                )

            
                            allowed_values = {"dynamic", "static"}

            
                            for x in dynamicity_parts:

            
                                if x.lower() not in allowed_values:

            
                                    raise ValueError(f"Invalid dynamicity value: {x}")


                            cls.custom_dynamicity = tuple(

            
                                Dynamicity.DYNAMIC if x.lower() == "dynamic" else Dynamicity.STATIC

            
                                for x in dynamicity_parts

            
                            )

            
                        except ValueError as e:

            
                            raise ValueError(f"Invalid --mnk_dynamicities format: {e}")


            # Class attribute to store custom MNK values

            
            ShapesId.custom_mnk_values = None

            
            # Class attribute to store custom dynamicities for MNK

            
            ShapesId.custom_dynamicity = None


            # 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.EASY_LARGE_STATIC:

            
                    return [(Dynamicity.STATIC, Dynamicity.STATIC, Dynamicity.STATIC)]

            
                elif shapes_id.custom_dynamicity:

            
                    return [shapes_id.custom_dynamicity]

            
                else:

            
                    return [

            
                        (Dynamicity.DYNAMIC, Dynamicity.DYNAMIC, Dynamicity.DYNAMIC),

            
                        (Dynamicity.STATIC, Dynamicity.STATIC, Dynamicity.STATIC),

            
                    ]

            
                raise ValueError(shapes_id)


            # Enumerates ways to construct MLIR tensor types.

            
            @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 matrix buffer contents.

            
            @enum.unique

            
            class MatrixGenerator(enum.Enum):

            
                ZERO = "zero"  # Fill with zeros

            
                RANDOM = "random"  # Fill with (deterministic) pseudorandom values.


            # Describes the shape of a matrix multiplication in the usual convention:

            
            # the LHS is {m}x{k}, the RHS is {k}x{n}, the accumulator/result is {m}x{n}.

            
            # The extra `accumulate` boolean tells whether the matmul is accumulating into

            
            # an existing accumulator (C += A * B) or just overwriting the result

            
            # (C = A * B).

            
            @dataclasses.dataclass

            
            class TestShape:

            
                m: int

            
                k: int

            
                n: int

            
                accumulate: bool


            # 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.

            
            # util.func @somefunction_DYNxDYNxf32, where we can't use "?" characters.

            
            def int_or_DYN(s: DimSize):

            
                return s.value or "DYN"


            # Describes the fully resolved shape dimensions of all 3 input matrices,

            
            # LHS, RHS, and Accumulator, in a testcase.

            
            # Each value is a string, which may either represent a positive integer such as "123",

            
            # or a "?" string, meaning a dynamic dimension as in MLIR.

            
            # These string values are used to generate MLIR function names and tensor shapes.

            
            @dataclasses.dataclass

            
            class TestInputMatricesShapes:

            
                lhs_rows: DimSize

            
                lhs_cols: DimSize

            
                rhs_rows: DimSize

            
                rhs_cols: DimSize

            
                acc_rows: DimSize

            
                acc_cols: DimSize


            # Represents a generated test function.

            
            @dataclasses.dataclass

            
            class MLIRFunction:

            
                name: str

            
                signature: str

            
                import_declaration: str

            
                definition: str


            # Represents a call to a generated test function.

            
            @dataclasses.dataclass

            
            class TestCall:

            
                function: MLIRFunction

            
                op: str


            # Helper for generate_function. Generates TestInputMatricesShapes, i.e.

            
            # converts from the runtime shape dimensions in TestShape and given dynamicities(m,n,k) to

            
            # the set of shapes to be used in a test function's input tensors.

            
            def generate_shapes(

            
                shape: TestShape,

            
                transpose_rhs: bool,

            
                dynamicities: tuple[Dynamicity, Dynamicity, Dynamicity],

            
            ):

            
                dynamicity_m, dynamicity_n, dynamicity_k = dynamicities


                lhs_rows = shape_dim(shape.m, dynamicity_m)

            
                lhs_cols = shape_dim(shape.k, dynamicity_k)

            
                acc_rows = shape_dim(shape.m, dynamicity_m)

            
                acc_cols = shape_dim(shape.n, dynamicity_n)

            
                if transpose_rhs:

            
                    rhs_rows = shape_dim(shape.n, dynamicity_n)

            
                    rhs_cols = shape_dim(shape.k, dynamicity_k)

            
                else:

            
                    rhs_rows = shape_dim(shape.k, dynamicity_k)

            
                    rhs_cols = shape_dim(shape.n, dynamicity_n)

            
                shapes = TestInputMatricesShapes(

            
                    lhs_rows=lhs_rows,

            
                    lhs_cols=lhs_cols,

            
                    rhs_rows=rhs_rows,

            
                    rhs_cols=rhs_cols,

            
                    acc_rows=acc_rows,

            
                    acc_cols=acc_cols,

            
                )

            
                return shapes


            random_matrix_seed = 0


            # Generate a matrix function argument of the given size as `%name`.

            
            def generate_random_matrix(

            
                name: str, matrix_shape: list, element_type: MatrixElemTypeId, increment_seed=True

            
            ):

            
                global random_matrix_seed

            
                if increment_seed:

            
                    random_matrix_seed += 1

            
                return (

            
                    f"  %{name}_dim0 = arith.constant {matrix_shape[0]} : i64\n"

            
                    f"  %{name}_dim1 = arith.constant {matrix_shape[1]} : i64\n"

            
                    f"  %{name}_element_type = hal.element_type<{element_type.value}> : i32\n"

            
                    f"  %{name}_seed = arith.constant {random_matrix_seed} : i32\n"

            
                    f"  %{name} = util.call @matmul_test.generate_random_matrix(%device, %{name}_dim0, %{name}_dim1, %{name}_element_type, %{name}_seed) : (!hal.device, i64, i64, i32, i32) -> !hal.buffer_view\n"

            
                )
	#!/usr/bin/env python3
	# Copyright 2025 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
	"""common utilities and data structures for e2e matmul tests"""

	import enum
	import dataclasses
	import re
	import typing


	# Data type of matrix entries. The string values must match MLIR data types.
	# This is a superset of the values accepted for the --lhs_rhs_types= flag,
	# as this also includes accumulator-specific types like i32.
	@enum.unique
	class MatrixElemTypeId(enum.Enum):
	NONE = ""
	UI8 = "ui8"
	I8 = "i8"
	I32 = "i32"
	F64 = "f64"
	F32 = "f32"
	F16 = "f16"
	BF16 = "bf16"
	F8E5M2 = "f8E5M2"
	F8E4M3FN = "f8E4M3FN"
	F8E5M2FNUZ = "f8E5M2FNUZ"
	F8E4M3FNUZ = "f8E4M3FNUZ"
	F8E8M0FNU = "f8E8M0FNU"
	F6E3M2FN = "f6E3M2FN"
	F6E2M3FN = "f6E2M3FN"
	F4E2M1FN = "f4E2M1FN"


	def get_size_in_bits(type_id: MatrixElemTypeId):
	return int(re.search(r"\d+", str(type_id)).group())


	# 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):
	DEFAULT = "default"
	SMALL = "small"
	LARGE = "large"
	EASY_LARGE_STATIC = "easy_large_static"
	CUSTOM_MNK = "custom_mnk" # Used for custom shapes specified by --mnk= flag.

	@classmethod
	def set_custom_mnk(cls, shapes_id, mnk_string):
	"""Parse and set custom MNK values from command line argument."""
	if shapes_id != cls.CUSTOM_MNK:
	if mnk_string:
	raise ValueError("--mnk can only be used with --shapes=custom_mnk")
	return

	# shapes_id is CUSTOM_MNK
	if not mnk_string:
	raise ValueError("--mnk must be specified when using --shapes=custom_mnk")
	try:
	mnk_parts = mnk_string.split(",")
	if len(mnk_parts) != 3:
	raise ValueError("--mnk must have exactly 3 values: m,n,k")
	cls.custom_mnk_values = tuple(int(x) for x in mnk_parts)
	except ValueError as e:
	raise ValueError(f"Invalid --mnk format: {e}")

	@classmethod
	def set_dynamicity(cls, shapes_id, dynamicity_string):
	"""Parse and set MNK dynamicity from command line argument."""
	if shapes_id != cls.CUSTOM_MNK:
	if dynamicity_string:
	raise ValueError(
	"--mnk_dynamicities can only be used with --shapes=custom_mnk"
	)
	return

	if dynamicity_string:
	try:
	dynamicity_parts = dynamicity_string.split(",")
	if len(dynamicity_parts) != 3:
	raise ValueError(
	"--mnk_dynamicities must have exactly 3 values, each being either 'dynamic' or 'static'"
	)
	allowed_values = {"dynamic", "static"}
	for x in dynamicity_parts:
	if x.lower() not in allowed_values:
	raise ValueError(f"Invalid dynamicity value: {x}")

	cls.custom_dynamicity = tuple(
	Dynamicity.DYNAMIC if x.lower() == "dynamic" else Dynamicity.STATIC
	for x in dynamicity_parts
	)
	except ValueError as e:
	raise ValueError(f"Invalid --mnk_dynamicities format: {e}")


	# Class attribute to store custom MNK values
	ShapesId.custom_mnk_values = None
	# Class attribute to store custom dynamicities for MNK
	ShapesId.custom_dynamicity = None


	# 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.EASY_LARGE_STATIC:
	return [(Dynamicity.STATIC, Dynamicity.STATIC, Dynamicity.STATIC)]
	elif shapes_id.custom_dynamicity:
	return [shapes_id.custom_dynamicity]
	else:
	return [
	(Dynamicity.DYNAMIC, Dynamicity.DYNAMIC, Dynamicity.DYNAMIC),
	(Dynamicity.STATIC, Dynamicity.STATIC, Dynamicity.STATIC),
	]
	raise ValueError(shapes_id)


	# Enumerates ways to construct MLIR tensor types.
	@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 matrix buffer contents.
	@enum.unique
	class MatrixGenerator(enum.Enum):
	ZERO = "zero" # Fill with zeros
	RANDOM = "random" # Fill with (deterministic) pseudorandom values.


	# Describes the shape of a matrix multiplication in the usual convention:
	# the LHS is {m}x{k}, the RHS is {k}x{n}, the accumulator/result is {m}x{n}.
	# The extra `accumulate` boolean tells whether the matmul is accumulating into
	# an existing accumulator (C += A * B) or just overwriting the result
	# (C = A * B).
	@dataclasses.dataclass
	class TestShape:
	m: int
	k: int
	n: int
	accumulate: bool


	# 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.
	# util.func @somefunction_DYNxDYNxf32, where we can't use "?" characters.
	def int_or_DYN(s: DimSize):
	return s.value or "DYN"


	# Describes the fully resolved shape dimensions of all 3 input matrices,
	# LHS, RHS, and Accumulator, in a testcase.
	# Each value is a string, which may either represent a positive integer such as "123",
	# or a "?" string, meaning a dynamic dimension as in MLIR.
	# These string values are used to generate MLIR function names and tensor shapes.
	@dataclasses.dataclass
	class TestInputMatricesShapes:
	lhs_rows: DimSize
	lhs_cols: DimSize
	rhs_rows: DimSize
	rhs_cols: DimSize
	acc_rows: DimSize
	acc_cols: DimSize


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


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


	# Helper for generate_function. Generates TestInputMatricesShapes, i.e.
	# converts from the runtime shape dimensions in TestShape and given dynamicities(m,n,k) to
	# the set of shapes to be used in a test function's input tensors.
	def generate_shapes(
	shape: TestShape,
	transpose_rhs: bool,
	dynamicities: tuple[Dynamicity, Dynamicity, Dynamicity],
	):
	dynamicity_m, dynamicity_n, dynamicity_k = dynamicities

	lhs_rows = shape_dim(shape.m, dynamicity_m)
	lhs_cols = shape_dim(shape.k, dynamicity_k)
	acc_rows = shape_dim(shape.m, dynamicity_m)
	acc_cols = shape_dim(shape.n, dynamicity_n)
	if transpose_rhs:
	rhs_rows = shape_dim(shape.n, dynamicity_n)
	rhs_cols = shape_dim(shape.k, dynamicity_k)
	else:
	rhs_rows = shape_dim(shape.k, dynamicity_k)
	rhs_cols = shape_dim(shape.n, dynamicity_n)
	shapes = TestInputMatricesShapes(
	lhs_rows=lhs_rows,
	lhs_cols=lhs_cols,
	rhs_rows=rhs_rows,
	rhs_cols=rhs_cols,
	acc_rows=acc_rows,
	acc_cols=acc_cols,
	)
	return shapes


	random_matrix_seed = 0


	# Generate a matrix function argument of the given size as `%name`.
	def generate_random_matrix(
	name: str, matrix_shape: list, element_type: MatrixElemTypeId, increment_seed=True
	):
	global random_matrix_seed
	if increment_seed:
	random_matrix_seed += 1
	return (
	f" %{name}_dim0 = arith.constant {matrix_shape[0]} : i64\n"
	f" %{name}_dim1 = arith.constant {matrix_shape[1]} : i64\n"
	f" %{name}_element_type = hal.element_type<{element_type.value}> : i32\n"
	f" %{name}_seed = arith.constant {random_matrix_seed} : i32\n"
	f" %{name} = util.call @matmul_test.generate_random_matrix(%device, %{name}_dim0, %{name}_dim1, %{name}_element_type, %{name}_seed) : (!hal.device, i64, i64, i32, i32) -> !hal.buffer_view\n"
	)