Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

remove linalg.pack #124

Merged
merged 1 commit into from
Feb 21, 2025
Merged

remove linalg.pack #124

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
182 changes: 0 additions & 182 deletions mlir/extras/dialects/ext/linalg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -534,185 +534,3 @@ def vecmat(y, A, x, *, loc=None, ip=None):
if loc is None:
loc = get_user_code_loc()
return linalg.vecmat(y, A, loc=loc, ip=ip, outs=[x])


@_cext.register_operation(linalg.Dialect)
class PackOp(ir.OpView):
OPERATION_NAME = "linalg.pack"

_ODS_OPERAND_SEGMENTS = [
1,
1,
0,
-1,
]

_ODS_REGIONS = (0, True)

def __init__(
self,
source,
dest,
inner_dims_pos,
inner_tiles,
static_inner_tiles,
*,
padding_value=None,
outer_dims_perm=None,
loc=None,
ip=None,
):
operands = []
results = []
attributes = {}
regions = None
operands.append(source)
operands.append(dest)
operands.append(padding_value)
operands.append(get_op_results_or_values(inner_tiles))
_ods_context = get_default_loc_context(loc)
if outer_dims_perm is not None:
attributes["outer_dims_perm"] = (
outer_dims_perm
if (
isinstance(outer_dims_perm, ir.Attribute)
or not ir.AttrBuilder.contains("DenseI64ArrayAttr")
)
else ir.AttrBuilder.get("DenseI64ArrayAttr")(
outer_dims_perm, context=_ods_context
)
)
attributes["inner_dims_pos"] = (
inner_dims_pos
if (
isinstance(inner_dims_pos, ir.Attribute)
or not ir.AttrBuilder.contains("DenseI64ArrayAttr")
)
else ir.AttrBuilder.get("DenseI64ArrayAttr")(
inner_dims_pos, context=_ods_context
)
)
attributes["static_inner_tiles"] = (
static_inner_tiles
if (
isinstance(static_inner_tiles, ir.Attribute)
or not ir.AttrBuilder.contains("DenseI64ArrayAttr")
)
else ir.AttrBuilder.get("DenseI64ArrayAttr")(
static_inner_tiles, context=_ods_context
)
)
_ods_successors = None
super().__init__(
self.OPERATION_NAME,
self._ODS_REGIONS,
self._ODS_OPERAND_SEGMENTS,
self._ODS_RESULT_SEGMENTS,
attributes=attributes,
operands=operands,
successors=_ods_successors,
regions=regions,
loc=loc,
ip=ip,
)

@property
def source(self):
operand_range = segmented_accessor(
self.operation.operands, self.operation.attributes["operandSegmentSizes"], 0
)
return operand_range[0]

@property
def dest(self):
operand_range = segmented_accessor(
self.operation.operands, self.operation.attributes["operandSegmentSizes"], 1
)
return operand_range[0]

@property
def padding_value(self):
operand_range = segmented_accessor(
self.operation.operands, self.operation.attributes["operandSegmentSizes"], 2
)
return operand_range[0] if len(operand_range) > 0 else None

@property
def inner_tiles(self):
operand_range = segmented_accessor(
self.operation.operands, self.operation.attributes["operandSegmentSizes"], 3
)
return operand_range

@property
def outer_dims_perm(self):
if "outer_dims_perm" not in self.operation.attributes:
return None
return self.operation.attributes["outer_dims_perm"]

@outer_dims_perm.setter
def outer_dims_perm(self, value):
if value is not None:
self.operation.attributes["outer_dims_perm"] = value
elif "outer_dims_perm" in self.operation.attributes:
del self.operation.attributes["outer_dims_perm"]

@outer_dims_perm.deleter
def outer_dims_perm(self):
del self.operation.attributes["outer_dims_perm"]

@property
def inner_dims_pos(self):
return self.operation.attributes["inner_dims_pos"]

@inner_dims_pos.setter
def inner_dims_pos(self, value):
if value is None:
raise ValueError("'None' not allowed as value for mandatory attributes")
self.operation.attributes["inner_dims_pos"] = value

@property
def static_inner_tiles(self):
return self.operation.attributes["static_inner_tiles"]

@static_inner_tiles.setter
def static_inner_tiles(self, value):
if value is None:
raise ValueError("'None' not allowed as value for mandatory attributes")
self.operation.attributes["static_inner_tiles"] = value

@property
def result(self):
return self.operation.results[0]


def pack(
source,
dest,
inner_dims_pos,
inner_tiles,
*,
padding_value=None,
outer_dims_perm=None,
loc=None,
ip=None,
) -> ir.Value:

(
dynamic_inner_tiles,
# packed here means %1:2 packing (results packing)
_inner_tiles,
static_inner_tiles,
) = _dispatch_mixed_values(inner_tiles)

return PackOp(
source=source,
dest=dest,
inner_dims_pos=inner_dims_pos,
inner_tiles=dynamic_inner_tiles,
static_inner_tiles=static_inner_tiles,
padding_value=padding_value,
outer_dims_perm=outer_dims_perm,
loc=loc,
ip=ip,
).result