Skip to content

[mlir][affine] implement inferType for delinearize #74644

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

Merged
merged 2 commits into from
Dec 7, 2023
Merged

[mlir][affine] implement inferType for delinearize #74644

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8effd1a
[mlir][affine] implement inferType for delinearize
makslevental Dec 6, 2023
7ace5a5
incorporate comments
makslevental Dec 7, 2023
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
14 changes: 2 additions & 12 deletions mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
include "mlir/Dialect/Arith/IR/ArithBase.td"
include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td"
include "mlir/Interfaces/ControlFlowInterfaces.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/LoopLikeInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"

Expand Down Expand Up @@ -63,10 +64,6 @@ def AffineApplyOp : Affine_Op<"apply", [Pure]> {
// has a constant builder. That way we wouldn't need to explicitly specify the
// result types here.
let builders = [
OpBuilder<(ins "AffineMap":$map, "ValueRange":$mapOperands),
[{
build($_builder, $_state, $_builder.getIndexType(), map, mapOperands);
}]>,
Comment on lines -66 to -69
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't quite understand the relationship between InferTypeOpInterface and these being autogenerated but that's what's happening here 🤷.

OpBuilder<(ins "ArrayRef<AffineExpr> ":$exprList,"ValueRange":$mapOperands),
[{
build($_builder, $_state, $_builder.getIndexType(),
Expand Down Expand Up @@ -541,13 +538,6 @@ class AffineMinMaxOpBase<string mnemonic, list<Trait> traits = []> :
let arguments = (ins AffineMapAttr:$map, Variadic<Index>:$operands);
let results = (outs Index);

let builders = [
OpBuilder<(ins "AffineMap":$affineMap, "ValueRange":$mapOperands),
[{
build($_builder, $_state, $_builder.getIndexType(), affineMap, mapOperands);
}]>
];

let extraClassDeclaration = [{
static StringRef getMapAttrStrName() { return "map"; }
AffineMap getAffineMap() { return getMap(); }
Expand Down Expand Up @@ -1068,7 +1058,7 @@ def AffineVectorStoreOp : AffineStoreOpBase<"vector_store"> {
//===----------------------------------------------------------------------===//

def AffineDelinearizeIndexOp : Affine_Op<"delinearize_index",
[Pure]> {
[Pure, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
let summary = "delinearize an index";
let description = [{
The `affine.delinearize_index` operation takes a single index value and
Expand Down
11 changes: 11 additions & 0 deletions mlir/lib/Dialect/Affine/IR/AffineOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4474,6 +4474,17 @@ LogicalResult AffineVectorStoreOp::verify() {
// DelinearizeIndexOp
//===----------------------------------------------------------------------===//

LogicalResult AffineDelinearizeIndexOp::inferReturnTypes(
MLIRContext *context, std::optional<::mlir::Location> location,
ValueRange operands, DictionaryAttr attributes, OpaqueProperties properties,
RegionRange regions, SmallVectorImpl<Type> &inferredReturnTypes) {
AffineDelinearizeIndexOpAdaptor adaptor(operands, attributes, properties,
regions);
inferredReturnTypes.assign(adaptor.getBasis().size(),
IndexType::get(context));
return success();
}

void AffineDelinearizeIndexOp::build(OpBuilder &builder, OperationState &result,
Value linearIndex,
ArrayRef<OpFoldResult> basis) {
Expand Down
1 change: 1 addition & 0 deletions mlir/lib/Dialect/Affine/IR/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ add_mlir_dialect_library(MLIRAffineDialect
MLIRArithDialect
MLIRDialectUtils
MLIRIR
MLIRInferTypeOpInterface
MLIRLoopLikeInterface
MLIRMemRefDialect
MLIRShapedOpInterfaces
Expand Down
12 changes: 12 additions & 0 deletions mlir/test/python/dialects/affine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from mlir.dialects import arith
from mlir.dialects import memref
from mlir.dialects import affine
import mlir.extras.types as T


def constructAndPrintInModule(f):
Expand Down Expand Up @@ -43,6 +44,17 @@ def affine_store_test(arg0):
return mem


# CHECK-LABEL: TEST: testAffineDelinearizeInfer
@constructAndPrintInModule
def testAffineDelinearizeInfer():
# CHECK: %[[C0:.*]] = arith.constant 0 : index
c0 = arith.ConstantOp(T.index(), 0)
# CHECK: %[[C1:.*]] = arith.constant 1 : index
c1 = arith.ConstantOp(T.index(), 1)
# CHECK: %{{.*}}:2 = affine.delinearize_index %[[C1:.*]] into (%[[C1:.*]], %[[C0:.*]]) : index, index
two_indices = affine.AffineDelinearizeIndexOp(c1, [c1, c0])


# CHECK-LABEL: TEST: testAffineLoadOp
@constructAndPrintInModule
def testAffineLoadOp():
Expand Down