Skip to content

layout op does not work with tensor metadata #5500

@zasdfgbnm

Description

@zasdfgbnm

In layoutAllocationDomain, the following code:

Fuser/csrc/ops/indexing.cpp

Lines 344 to 362 in 06a21af

// pad row size: num_groups * (row_multiple - 1) + row_size
auto pad_to_max_extent = [&](IterDomain* id, int multiple) -> IterDomain* {
auto* maximum_pad_value_per_group =
IrBuilder::create<Val>(multiple - 1, DataType::Index);
// NOTE: we do not use `resize` to represent the padding.
//
// resize sounds good in theory, because transformation can propagate across
// it. In reality, we do not have a protocol to index this operation via the
// logical to allocation domain transform. I question how much a resize op
// provides in functionality. More importantly, using resize hits asserts in
// vectorization analysis (validateDeviceSplit ATM), which doesn't look easy
// to handle for me.
Val* padded_ext = SimplifyingIrBuilder::addExpr(
id->extent(),
SimplifyingIrBuilder::mulExpr(num_groups, maximum_pad_value_per_group));
return IterDomainBuilder(id).extent(padded_ext).build();
};
alloc_dom.push_back(pad_to_max_extent(logical_dom[0], row_multiple));

We are creating a dangling IterDomain in the allocation domain, which is not connected to the logical domain. However, in tensor_metadata.cpp, especially

ForwardTraverseFromLogicalToAlloc(ee, active_ids).run(tv, logical, alloc);
BackwardTraverseFromLogicalToAlloc(ee, active_ids).run(tv, logical, alloc);

we rely on the connection between logical domain and allocation domain to infer the shape and stride.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions