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"
	)