Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
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
1 change: 1 addition & 0 deletions mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
#define MLIR_DIALECT_AFFINE_IR_AFFINEOPS_H

#include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h"
#include "mlir/Dialect/Affine/IR/AffineTraits.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/AffineMap.h"
Expand Down
1 change: 1 addition & 0 deletions mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
#ifndef AFFINE_OPS
#define AFFINE_OPS

include "mlir/Dialect/Affine/IR/AffineTraits.td"
include "mlir/Dialect/Arith/IR/ArithBase.td"
include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td"
include "mlir/Interfaces/ControlFlowInterfaces.td"
Expand Down
27 changes: 27 additions & 0 deletions mlir/include/mlir/Dialect/Affine/IR/AffineTraits.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
//===- AffineTraits.h - MLIR Affine Traits --------------------------------===//
//
// Part of the LLVM Project, 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
//
//===----------------------------------------------------------------------===//
//
// This file defines traits brought in by the Affine dialect.
//
//===----------------------------------------------------------------------===//
#ifndef AFFINE_TRAITS_H
#define AFFINE_TRAITS_H

#include "mlir/IR/OpDefinition.h"

namespace mlir::OpTrait {

template <typename ConcreteType>
class AffineDim : public TraitBase<ConcreteType, AffineDim> {
public:
static LogicalResult verifyTrait(Operation *op) { return success(); }
};

} // namespace mlir::OpTrait

#endif // AFFINE_TRAITS_H
22 changes: 22 additions & 0 deletions mlir/include/mlir/Dialect/Affine/IR/AffineTraits.td
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
//===- AffineTraits.td - Affine dialect traits -------------*- tablegen -*-===//
//
// Part of the LLVM Project, 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
//
//===----------------------------------------------------------------------===//
//
// Defines traits brought in by the MLIR Affine dialect.
//
//===----------------------------------------------------------------------===//
#ifndef AFFINE_TRAITS
#define AFFINE_TRAITS

include "mlir/IR/OpBase.td"

// Trait to declare that an op result is an affine dimension identifier.
// Prevents the result from being seen as a symbol into AffineMaps
// and IntegerSets.
def AffineDim : NativeOpTrait<"AffineDim">;

#endif // AFFINE_TRAITS
1 change: 1 addition & 0 deletions mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#define MLIR_DIALECT_LINALG_IR_LINALG_H

#include "mlir/Bytecode/BytecodeOpInterface.h"
#include "mlir/Dialect/Affine/IR/AffineTraits.h"
#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/AffineExpr.h"
Expand Down
3 changes: 2 additions & 1 deletion mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
#ifndef LINALG_OPS
#define LINALG_OPS

include "mlir/Dialect/Affine/IR/AffineTraits.td"
include "mlir/Dialect/Linalg/IR/LinalgBase.td"
include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
include "mlir/Interfaces/ControlFlowInterfaces.td"
Expand Down Expand Up @@ -46,7 +47,7 @@ def Linalg_YieldOp : Linalg_Op<"yield", [Pure, ReturnLike, Terminator]>,
let hasVerifier = 1;
}

def Linalg_IndexOp : Linalg_Op<"index", [Pure]>,
def Linalg_IndexOp : Linalg_Op<"index", [Pure, AffineDim]>,
Arguments<(ins ConfinedAttr<I64Attr, [IntMinValue<0>]>:$dim)>,
Results<(outs Index:$result)> {
let summary = "linalg index operation";
Expand Down
9 changes: 9 additions & 0 deletions mlir/lib/Dialect/Affine/IR/AffineOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
Expand Down Expand Up @@ -312,6 +313,10 @@ bool mlir::affine::isValidDim(Value value, Region *region) {
return isa<AffineForOp, AffineParallelOp>(parentOp);
}

// Remove me: linalg.index ops are valid affine dim identifiers
if (op->hasTrait<OpTrait::AffineDim>())
return true;

// Affine apply operation is ok if all of its operands are ok.
if (auto applyOp = dyn_cast<AffineApplyOp>(op))
return applyOp.isValidDim(region);
Expand Down Expand Up @@ -439,6 +444,10 @@ bool mlir::affine::isValidSymbol(Value value, Region *region) {
return false;
}

// Remove me: linalg.index ops are not valid affine symbols
if (defOp->hasTrait<OpTrait::AffineDim>())
return false;

// Constant operation is ok.
Attribute operandCst;
if (matchPattern(defOp, m_Constant(&operandCst)))
Expand Down
24 changes: 12 additions & 12 deletions mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,9 @@ module attributes {transform.with_named_sequence} {
// CHECK: %[[KINDEX:.+]] = linalg.index 2 : index

// Compute input channel/convolved indices.
// CHECK: %[[ICINDEX:.+]] = affine.apply affine_map<()[s0] -> (s0 mod 4)>()[%[[KINDEX]]]
// CHECK: %[[CONVH:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 floordiv 14 + s1 floordiv 12)>()[%[[MINDEX]], %[[KINDEX]]]
// CHECK: %[[CONVW:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 mod 14 + (s1 mod 12) floordiv 4)>()[%[[MINDEX]], %[[KINDEX]]]
// CHECK: %[[ICINDEX:.+]] = affine.apply affine_map<(d0) -> (d0 mod 4)>(%[[KINDEX]])
// CHECK: %[[CONVH:.+]] = affine.apply affine_map<(d0, d1) -> (d0 floordiv 14 + d1 floordiv 12)>(%[[MINDEX]], %[[KINDEX]])
// CHECK: %[[CONVW:.+]] = affine.apply affine_map<(d0, d1) -> (d0 mod 14 + (d1 mod 12) floordiv 4)>(%[[MINDEX]], %[[KINDEX]])

// Extract from the input tensor.
// CHECK: %[[EXTRACTED_INPUT:.+]] = tensor.extract
Expand Down Expand Up @@ -227,9 +227,9 @@ module attributes {transform.with_named_sequence} {
// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>

// Im2col maps
// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0] -> (s0 floordiv 9)>
// CHECK-DAG: #[[MAP7:.+]] = affine_map<()[s0, s1] -> (s0 floordiv 14 + (s1 mod 9) floordiv 3)>
// CHECK-DAG: #[[MAP8:.+]] = affine_map<()[s0, s1] -> (s0 + s1 - (s0 floordiv 14) * 14 - (s1 floordiv 3) * 3)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0 floordiv 9)>
// CHECK-DAG: #[[MAP7:.+]] = affine_map<(d0, d1) -> (d0 floordiv 14 + (d1 mod 9) floordiv 3)>
// CHECK-DAG: #[[MAP8:.+]] = affine_map<(d0, d1) -> (d0 + d1 - (d0 floordiv 14) * 14 - (d1 floordiv 3) * 3)>


// CHECK-DAG: #[[LHSMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)>
Expand All @@ -251,9 +251,9 @@ module attributes {transform.with_named_sequence} {
// CHECK: %[[NINDEX:.+]] = linalg.index 2 : index

// Compute input channel/convolved indices.
// CHECK: %[[ICINDEX:.+]] = affine.apply #[[MAP1]]()[%[[KINDEX]]]
// CHECK: %[[CONVH:.+]] = affine.apply #[[MAP7]]()[%[[NINDEX]], %[[KINDEX]]]
// CHECK: %[[CONVW:.+]] = affine.apply #[[MAP8]]()[%[[NINDEX]], %[[KINDEX]]]
// CHECK: %[[ICINDEX:.+]] = affine.apply #[[MAP1]](%[[KINDEX]])
// CHECK: %[[CONVH:.+]] = affine.apply #[[MAP7]](%[[NINDEX]], %[[KINDEX]])
// CHECK: %[[CONVW:.+]] = affine.apply #[[MAP8]](%[[NINDEX]], %[[KINDEX]])

// Extract from the input tensor.
// CHECK: %[[EXTRACTED_INPUT:.+]] = tensor.extract
Expand Down Expand Up @@ -300,9 +300,9 @@ module attributes {transform.with_named_sequence} {
// CHECK: %[[KINDEX:.+]] = linalg.index 2 : index

// Compute input channel/convolved indices.
// CHECK: %[[ICINDEX:.+]] = affine.apply affine_map<()[s0] -> (s0 mod 4)>()[%[[KINDEX]]]
// CHECK: %[[CONVH:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 floordiv 14 + s1 floordiv 12)>()[%[[MINDEX]], %[[KINDEX]]]
// CHECK: %[[CONVW:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 mod 14 + (s1 mod 12) floordiv 4)>()[%[[MINDEX]], %[[KINDEX]]]
// CHECK: %[[ICINDEX:.+]] = affine.apply affine_map<(d0) -> (d0 mod 4)>(%[[KINDEX]])
// CHECK: %[[CONVH:.+]] = affine.apply affine_map<(d0, d1) -> (d0 floordiv 14 + d1 floordiv 12)>(%[[MINDEX]], %[[KINDEX]])
// CHECK: %[[CONVW:.+]] = affine.apply affine_map<(d0, d1) -> (d0 mod 14 + (d1 mod 12) floordiv 4)>(%[[MINDEX]], %[[KINDEX]])

// Extract from the input tensor.
// CHECK: %[[EXTRACTED_INPUT:.+]] = tensor.extract
Expand Down
4 changes: 2 additions & 2 deletions mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -986,13 +986,13 @@ module {
#map3 = affine_map<(d0, d1) -> (d0, d1)>

// CHECK-DAG: [[$MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-DAG: [[$MAP1:#[a-zA-Z0-9_]*]] = affine_map<()[s0] -> (s0 floordiv 4)>
// CHECK-DAG: [[$MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0) -> (d0 floordiv 4)>

func.func @fuse_and_collapse(%arg0: tensor<3x4xindex>) -> tensor<2x12xindex> {
%1 = tensor.empty() : tensor<2x3x4xindex>
// CHECK: linalg.generic {
// CHECK: %[[INDEX1:[a-zA-Z0-9_]+]] = linalg.index 1 : index
// CHECK-NEXT: %[[MAP:[a-zA-Z0-9_]+]] = affine.apply #map1()[%[[INDEX1]]]
// CHECK-NEXT: %[[MAP:[a-zA-Z0-9_]+]] = affine.apply #map1(%[[INDEX1]])
// CHECK-NEXT: linalg.yield %[[MAP]] : index
%2 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg0: tensor<3x4xindex>) outs(%1 : tensor<2x3x4xindex>) {
^bb0(%in: index, %out: index):
Expand Down
10 changes: 5 additions & 5 deletions mlir/test/Dialect/Linalg/tile-indexed.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,13 @@ module attributes {transform.with_named_sequence} {
}
}

// TILE-10n25-DAG: [[$MAP:#[a-zA-Z0-9_]*]] = affine_map<(d0)[s0] -> (d0 + s0)>
// TILE-10n25-DAG: [[$MAP:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0 + d1)>
// TILE-10n25-LABEL: func @indexed_vector
// TILE-10n25: %[[C10:.*]] = arith.constant 10 : index
// TILE-10n25: scf.for %[[J:.*]] = {{.*}} step %[[C10]]
// TILE-10n25: linalg.generic
// TILE-10n25: %[[I:.*]] = linalg.index 0 : index
// TILE-10n25: %[[NEW_I:.*]] = affine.apply [[$MAP]](%[[J]])[%[[I]]]
// TILE-10n25: %[[NEW_I:.*]] = affine.apply [[$MAP]](%[[I]], %[[J]])
// TILE-10n25: linalg.yield %[[NEW_I]] : index

// -----
Expand All @@ -51,16 +51,16 @@ module attributes {transform.with_named_sequence} {
}
}

// TILE-10n25-DAG: [[$MAP:#[a-zA-Z0-9_]*]] = affine_map<(d0)[s0] -> (d0 + s0)>
// TILE-10n25-DAG: [[$MAP:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0 + d1)>
// TILE-10n25-LABEL: func @indexed_matrix
// TILE-10n25-DAG: %[[C25:.*]] = arith.constant 25 : index
// TILE-10n25-DAG: %[[C10:.*]] = arith.constant 10 : index
// TILE-10n25: scf.for %[[K:.*]] = {{.*}} step %[[C10]]
// TILE-10n25: scf.for %[[L:.*]] = {{.*}} step %[[C25]]
// TILE-10n25: linalg.generic
// TILE-10n25: %[[I:.*]] = linalg.index 0 : index
// TILE-10n25: %[[NEW_I:.*]] = affine.apply [[$MAP]](%[[K]])[%[[I]]]
// TILE-10n25: %[[NEW_I:.*]] = affine.apply [[$MAP]](%[[I]], %[[K]])
// TILE-10n25: %[[J:.*]] = linalg.index 1 : index
// TILE-10n25: %[[NEW_J:.*]] = affine.apply [[$MAP]](%[[L]])[%[[J]]]
// TILE-10n25: %[[NEW_J:.*]] = affine.apply [[$MAP]](%[[J]], %[[L]])
// TILE-10n25: %[[SUM:.*]] = arith.addi %[[NEW_I]], %[[NEW_J]] : index
// TILE-10n25: linalg.yield %[[SUM]] : index
4 changes: 2 additions & 2 deletions mlir/test/Dialect/Linalg/transform-op-split.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ module attributes {transform.with_named_sequence} {

func.func private @elem(%arg0: f32, %arg1: index, %arg2: index) -> f32

// CHECK: #[[$ADD_42_MAP:.+]] = affine_map<()[s0] -> (s0 + 42)>
// CHECK: #[[$ADD_42_MAP:.+]] = affine_map<(d0) -> (d0 + 42)>

// CHECK-LABEL: @one_d_static
// CHECK-SAME: %[[IN:.+]]: tensor<100xf32>, %[[OUT:.+]]: tensor<100xf32>
Expand All @@ -30,7 +30,7 @@ func.func @one_d_static(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tenso
// CHECK: ins(%[[IN_SLICE_HIGH]]
// CHECK: outs(%[[OUT_SLICE_HIGH]]
// CHECK: %[[IDX:.+]] = linalg.index 0
// CHECK: affine.apply #[[$ADD_42_MAP]]()[%[[IDX]]]
// CHECK: affine.apply #[[$ADD_42_MAP]](%[[IDX]])
// CHECK: func.call @elem
// CHECK: %[[RES:.+]] = tensor.insert_slice %[[RES_SLICE_HIGH]] into %[[RES_PARTIAL]][42] [58] [1]
%0 = linalg.generic {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -259,14 +259,14 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
// CHECK: #[[$MAP_ADD:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
// CHECK: #[[$MAP_ADD:.+]] = affine_map<(d0, d1) -> (d0 + d1)>
// CHECK-LABEL: @indexed_semantics
// CHECK: scf.for %[[I0:.+]] = %{{.*}} to %{{.*}} step %{{.*}}
// CHECK: scf.for %[[I1:.+]] = %{{.*}} to %{{.*}} step %{{.*}}
// CHECK: %[[INDEX0:.+]] = linalg.index 0
// CHECK: %[[INDEX0_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[I0]])[%[[INDEX0]]]
// CHECK: %[[INDEX0_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[INDEX0]], %[[I0]])
// CHECK: %[[INDEX1:.+]] = linalg.index 1
// CHECK: %[[INDEX1_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[I1]])[%[[INDEX1]]]
// CHECK: %[[INDEX1_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[INDEX1]], %[[I1]])
// CHECK: arith.addi %[[INDEX0_AMENDED]], %[[INDEX1_AMENDED]]

// -----
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -205,7 +205,7 @@ module attributes {transform.with_named_sequence} {

// -----

// CHECK: #[[$MAP_ADD:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
// CHECK: #[[$MAP_ADD:.+]] = affine_map<(d0, d1) -> (d0 + d1)>

func.func @indexed_semantics(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
// Check that we correctly amend "linalg.index" results.
Expand Down Expand Up @@ -241,9 +241,9 @@ module attributes {transform.with_named_sequence} {
// CHECK-LABEL: @indexed_semantics
// CHECK: scf.forall (%[[I0:.+]], %[[I1:.+]]) =
// CHECK: %[[INDEX0:.+]] = linalg.index 0
// CHECK: %[[INDEX0_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[I0]])[%[[INDEX0]]]
// CHECK: %[[INDEX0_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[INDEX0]], %[[I0]])
// CHECK: %[[INDEX1:.+]] = linalg.index 1
// CHECK: %[[INDEX1_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[I1]])[%[[INDEX1]]]
// CHECK: %[[INDEX1_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[INDEX1]], %[[I1]])
// CHECK: arith.addi %[[INDEX0_AMENDED]], %[[INDEX1_AMENDED]]

// -----
Expand Down