[WIP] Sketching IR and parsing support for affine maps, affine expressions
Run test case:
$ mlir-opt test/IR/parser-affine-map.mlir
test/IR/parser-affine-map.mlir:3:30: error: expect '(' at start of map range
#hello_world2 (i, j) [s0] -> i+s0, j)
^
PiperOrigin-RevId: 202736856
diff --git a/include/mlir/IR/AffineExpr.h b/include/mlir/IR/AffineExpr.h
index 4d8516c..b5bf8f2 100644
--- a/include/mlir/IR/AffineExpr.h
+++ b/include/mlir/IR/AffineExpr.h
@@ -23,12 +23,229 @@
#ifndef MLIR_IR_AFFINE_EXPR_H
#define MLIR_IR_AFFINE_EXPR_H
+#include "mlir/Support/LLVM.h"
+
namespace mlir {
-class AffineExpr {
+class MLIRContext;
+
+/// A one-dimensional affine expression.
+/// AffineExpression's are immutable (like Type's)
+class AffineExpr {
public:
- AffineExpr();
- // TODO(andydavis,bondhugula) Implement affine expressions.
+ enum class Kind {
+ // Add.
+ Add,
+ // Mul.
+ Mul,
+ // Mod.
+ Mod,
+ // Floordiv
+ FloorDiv,
+ // Ceildiv
+ CeilDiv,
+
+ /// This is a marker for the last affine binary op. The range of binary op's
+ /// is expected to be this element and earlier.
+ LAST_AFFINE_BINARY_OP = CeilDiv,
+
+ // Unary op negation
+ Neg,
+
+ // Constant integer.
+ Constant,
+ // Dimensional identifier.
+ DimId,
+ // Symbolic identifier.
+ SymbolId,
+ };
+
+ /// Return the classification for this type.
+ Kind getKind() const { return kind; }
+
+ ~AffineExpr() = default;
+
+ void print(raw_ostream &os) const;
+ void dump() const;
+
+ protected:
+ explicit AffineExpr(Kind kind) : kind(kind) {}
+
+ private:
+ /// Classification of the subclass
+ const Kind kind;
+};
+
+/// Binary affine expression.
+class AffineBinaryOpExpr : public AffineExpr {
+ public:
+ static AffineBinaryOpExpr *get(Kind kind, AffineExpr *lhsOperand,
+ AffineExpr *rhsOperand, MLIRContext *context);
+
+ AffineExpr *getLeftOperand() const { return lhsOperand; }
+ AffineExpr *getRightOperand() const { return rhsOperand; }
+
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const AffineExpr *expr) {
+ return expr->getKind() <= Kind::LAST_AFFINE_BINARY_OP;
+ }
+
+ protected:
+ explicit AffineBinaryOpExpr(Kind kind, AffineExpr *lhsOperand,
+ AffineExpr *rhsOperand)
+ : AffineExpr(kind), lhsOperand(lhsOperand), rhsOperand(rhsOperand) {}
+
+ AffineExpr *const lhsOperand;
+ AffineExpr *const rhsOperand;
+};
+
+/// Binary affine add expression.
+class AffineAddExpr : public AffineBinaryOpExpr {
+ public:
+ static AffineAddExpr *get(AffineExpr *lhsOperand, AffineExpr *rhsOperand,
+ MLIRContext *context);
+
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const AffineExpr *expr) {
+ return expr->getKind() == Kind::Add;
+ }
+
+ private:
+ explicit AffineAddExpr(AffineExpr *lhsOperand, AffineExpr *rhsOperand)
+ : AffineBinaryOpExpr(Kind::Add, lhsOperand, rhsOperand) {}
+};
+
+/// Binary affine mul expression.
+class AffineMulExpr : public AffineBinaryOpExpr {
+ public:
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const AffineExpr *expr) {
+ return expr->getKind() == Kind::Mul;
+ }
+
+ private:
+ explicit AffineMulExpr(AffineExpr *lhsOperand, AffineExpr *rhsOperand)
+ : AffineBinaryOpExpr(Kind::Mul, lhsOperand, rhsOperand) {}
+};
+
+/// Binary affine mod expression.
+class AffineModExpr : public AffineBinaryOpExpr {
+ public:
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const AffineExpr *expr) {
+ return expr->getKind() == Kind::Mod;
+ }
+
+ private:
+ explicit AffineModExpr(AffineExpr *lhsOperand, AffineExpr *rhsOperand)
+ : AffineBinaryOpExpr(Kind::Mod, lhsOperand, rhsOperand) {}
+};
+
+/// Binary affine floordiv expression.
+class AffineFloorDivExpr : public AffineBinaryOpExpr {
+ public:
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const AffineExpr *expr) {
+ return expr->getKind() == Kind::FloorDiv;
+ }
+
+ private:
+ explicit AffineFloorDivExpr(AffineExpr *lhsOperand, AffineExpr *rhsOperand)
+ : AffineBinaryOpExpr(Kind::FloorDiv, lhsOperand, rhsOperand) {}
+};
+
+/// Binary affine ceildiv expression.
+class AffineCeilDivExpr : public AffineBinaryOpExpr {
+ public:
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const AffineExpr *expr) {
+ return expr->getKind() == Kind::CeilDiv;
+ }
+
+ private:
+ explicit AffineCeilDivExpr(AffineExpr *lhsOperand, AffineExpr *rhsOperand)
+ : AffineBinaryOpExpr(Kind::CeilDiv, lhsOperand, rhsOperand) {}
+};
+
+/// Unary affine expression.
+class AffineUnaryOpExpr : public AffineExpr {
+ public:
+ static AffineUnaryOpExpr *get(const AffineExpr &operand,
+ MLIRContext *context);
+
+ static AffineUnaryOpExpr *get(const AffineExpr &operand);
+ AffineExpr *getOperand() const { return operand; }
+
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const AffineExpr *expr) {
+ return expr->getKind() == Kind::Neg;
+ }
+
+ private:
+ explicit AffineUnaryOpExpr(Kind kind, AffineExpr *operand)
+ : AffineExpr(kind), operand(operand) {}
+
+ AffineExpr *operand;
+};
+
+/// A argument identifier appearing in an affine expression
+class AffineDimExpr : public AffineExpr {
+ public:
+ static AffineDimExpr *get(unsigned position, MLIRContext *context);
+
+ unsigned getPosition() const { return position; }
+
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const AffineExpr *expr) {
+ return expr->getKind() == Kind::DimId;
+ }
+
+ private:
+ explicit AffineDimExpr(unsigned position)
+ : AffineExpr(Kind::DimId), position(position) {}
+
+ /// Position of this identifier in the argument list.
+ unsigned position;
+};
+
+/// A symbolic identifier appearing in an affine expression
+class AffineSymbolExpr : public AffineExpr {
+ public:
+ static AffineSymbolExpr *get(unsigned position, MLIRContext *context);
+
+ unsigned getPosition() const { return position; }
+
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const AffineExpr *expr) {
+ return expr->getKind() == Kind::SymbolId;
+ }
+
+ private:
+ explicit AffineSymbolExpr(unsigned position)
+ : AffineExpr(Kind::SymbolId), position(position) {}
+
+ /// Position of this identifier in the symbol list.
+ unsigned position;
+};
+
+/// An integer constant appearing in affine expression.
+class AffineConstantExpr : public AffineExpr {
+ public:
+ static AffineConstantExpr *get(int64_t constant, MLIRContext *context);
+
+ int64_t getValue() const { return constant; }
+
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const AffineExpr *expr) {
+ return expr->getKind() == Kind::Constant;
+ }
+
+ private:
+ explicit AffineConstantExpr(int64_t constant)
+ : AffineExpr(Kind::Constant), constant(constant) {}
+
+ // The constant.
+ int64_t constant;
};
} // end namespace mlir
diff --git a/include/mlir/IR/AffineMap.h b/include/mlir/IR/AffineMap.h
index af1df96..6266050 100644
--- a/include/mlir/IR/AffineMap.h
+++ b/include/mlir/IR/AffineMap.h
@@ -26,30 +26,43 @@
#include <vector>
#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/ArrayRef.h"
namespace mlir {
+class MLIRContext;
class AffineExpr;
-class AffineMap {
+/// A multi-dimensional affine map
+/// Affine map's are immutable like Type's, and they are uniqued.
+/// Eg: (d0, d1) -> (d0/128, d0 mod 128, d1)
+/// The names used (d0, d1) don't matter - it's the mathematical function that
+/// is unique to this affine map.
+class AffineMap {
public:
- // Constructs an AffineMap with 'dimCount' dimension identifiers, and
- // 'symbolCount' symbols.
- // TODO(andydavis) Pass in ArrayRef<AffineExpr*> to populate list of exprs.
- AffineMap(unsigned dimCount, unsigned symbolCount);
+ static AffineMap *get(unsigned dimCount, unsigned symbolCount,
+ ArrayRef<AffineExpr *> exprs,
+ MLIRContext *context);
// Prints affine map to 'os'.
void print(raw_ostream &os) const;
+ void dump() const;
+
+ unsigned dimCount() const { return numDims; }
+ unsigned symbolCount() const { return numSymbols; }
private:
- // Number of dimensional indentifiers.
- const unsigned dimCount;
- // Number of symbols.
- const unsigned symbolCount;
- // TODO(andydavis) Do not use std::vector here (array size is not dynamic).
- std::vector<AffineExpr*> exprs;
+ AffineMap(unsigned dimCount, unsigned symbolCount,
+ ArrayRef<AffineExpr *> exprs);
+
+ const unsigned numDims;
+ const unsigned numSymbols;
+
+ /// The affine expressions for this (multi-dimensional) map.
+ /// TODO: use trailing objects for these
+ ArrayRef<AffineExpr *> exprs;
};
-} // end namespace mlir
+} // end namespace mlir
#endif // MLIR_IR_AFFINE_MAP_H
diff --git a/include/mlir/IR/Module.h b/include/mlir/IR/Module.h
index bf77078..7d22efd 100644
--- a/include/mlir/IR/Module.h
+++ b/include/mlir/IR/Module.h
@@ -22,10 +22,13 @@
#ifndef MLIR_IR_MODULE_H
#define MLIR_IR_MODULE_H
+#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Function.h"
-#include <vector>
namespace mlir {
+
+class AffineMap;
+
class Module {
public:
explicit Module();
@@ -33,6 +36,9 @@
// FIXME: wrong representation and API.
std::vector<Function*> functionList;
+ // FIXME: wrong representation and API.
+ // These affine maps are immutable
+ std::vector<const AffineMap *> affineMapList;
void print(raw_ostream &os) const;
void dump() const;