Adds newly renamed "affine_apply" operation to StandardOps.

Breaks "core operations" tests out into their own test file.

PiperOrigin-RevId: 205848090

mlir/include/mlir/IR/StandardOps.h

@@ -116,6 +116,36 @@ private:
   explicit DimOp(const Operation *state) : Base(state) {}
 };
 
+// The "affine_apply" operation applies an affine map to a list of operands,
+// yielding a list of results. The operand and result list sizes must be the
+// same. All operands and results are of type 'AffineInt'. This operation
+// requires a single affine map attribute named "map".
+// For example:
+//
+//   %y = "affine_apply" (%x) { map: (d0) -> (d0 + 1) } :
+//          (affineint) -> (affineint)
+//
+// TODO Add VariadicOperands/VariadicResults trait, and use them here
+// for getOperand/getResult accessors.
+class AffineApplyOp : public OpImpl::Base<AffineApplyOp> {
+public:
+  // Returns the affine map to be applied by this operation.
+  AffineMap *getAffineMap() const {
+    return getAttrOfType<AffineMapAttr>("map")->getValue();
+  }
+
+  /// Methods for support type inquiry through isa, cast, and dyn_cast.
+  static StringRef getOperationName() { return "affine_apply"; }
+
+  // Hooks to customize behavior of this op.
+  const char *verify() const;
+  void print(raw_ostream &os) const;
+
+private:
+  friend class Operation;
+  explicit AffineApplyOp(const Operation *state) : Base(state) {}
+};
+
 /// Install the standard operations in the specified operation set.
 void registerStandardOperations(OperationSet &opSet);
 

mlir/lib/IR/StandardOps.cpp

@@ -16,6 +16,7 @@
 // =============================================================================
 
 #include "mlir/IR/StandardOps.h"
+#include "mlir/IR/AffineMap.h"
 #include "mlir/IR/OperationSet.h"
 #include "mlir/IR/SSAValue.h"
 #include "mlir/IR/Types.h"
@@ -89,7 +90,23 @@ const char *DimOp::verify() const {
   return nullptr;
 }
 
+void AffineApplyOp::print(raw_ostream &os) const {
+  os << "affine_apply map: ";
+  getAffineMap()->print(os);
+}
+
+const char *AffineApplyOp::verify() const {
+  // TODO: Check input and output dimensions match.
+
+  // Check that affine map attribute was specified
+  auto affineMapAttr = getAttrOfType<AffineMapAttr>("map");
+  if (!affineMapAttr)
+    return "requires an affine map.";
+
+  return nullptr;
+}
+
 /// Install the standard operations in the specified operation set.
 void mlir::registerStandardOperations(OperationSet &opSet) {
-  opSet.addOperations<AddFOp, ConstantOp, DimOp>(/*prefix=*/"");
+  opSet.addOperations<AddFOp, ConstantOp, DimOp, AffineApplyOp>(/*prefix=*/"");
 }

mlir/test/IR/invalid-ops.mlir

@@ -33,3 +33,12 @@ bb:
   return
 }
 
+// -----
+
+cfgfunc @affine_apply_no_map() {
+bb0:
+  // TODO Make constant work with affineint.
+  %i = "const"() {value: 0} : () -> affineint
+  %x = "affine_apply" (%i) { } : (affineint) -> (affineint) //  expected-error {{'affine_apply' op requires an affine map.}}
+  return
+}

mlir/test/IR/parser-core-ops.mlir

@@ -0,0 +1,55 @@
+// RUN: %S/../../mlir-opt %s -o - | FileCheck %s
+
+// CHECK: #map{{[0-9]+}} = (d0, d1) -> ((d0 + 1), (d1 + 2))
+#map5 = (d0, d1) -> (d0 + 1, d1 + 2)
+
+// CHECK-LABEL: cfgfunc @cfgfunc_with_ops(f32) {
+cfgfunc @cfgfunc_with_ops(f32) {
+bb0(%a : f32):
+  // CHECK: %1 = "getTensor"() : () -> tensor<4x4x?xf32>
+  %t = "getTensor"() : () -> tensor<4x4x?xf32>
+
+  // CHECK: dim xxx, 2 : sometype
+  %t2 = "dim"(%t){index: 2} : (tensor<4x4x?xf32>) -> affineint
+
+  // CHECK: addf xx, yy : sometype
+  %x = "addf"(%a, %a) : (f32,f32) -> (f32)
+
+  // CHECK:   return
+  return
+}
+
+// CHECK-LABEL: cfgfunc @standard_instrs() {
+cfgfunc @standard_instrs() {
+bb42:       // CHECK: bb0:
+  // CHECK: %0 = "getTensor"() : () -> tensor<4x4x?xf32>
+  %42 = "getTensor"() : () -> tensor<4x4x?xf32>
+
+  // CHECK: dim xxx, 2 : sometype
+  %a = "dim"(%42){index: 2} : (tensor<4x4x?xf32>) -> affineint
+
+  %f = "Const"(){value: 1} : () -> f32
+  // CHECK: addf xx, yy : sometype
+  "addf"(%f, %f) : (f32,f32) -> f32
+
+  // TODO:   CHECK: FIXME: IMPLEMENT DEFAULT PRINTER
+  %x = "constant"(){value: 42} : () -> i32
+  return
+}
+
+// CHECK-LABEL: cfgfunc @affine_apply() {
+cfgfunc @affine_apply() {
+bb0:
+  // TODO: Make constant work with affineint.
+  %i = "const"() {value: 0} : () -> affineint
+  %j = "const"() {value: 1} : () -> affineint
+
+  // CHECK: affine_apply map: (d0) -> ((d0 + 1))
+  %x = "affine_apply" (%i) { map: (d0) -> (d0 + 1) } :
+    (affineint) -> (affineint)
+
+  // CHECK: affine_apply map: (d0, d1) -> ((d0 + 1), (d1 + 2))
+  %y = "affine_apply" (%i, %j) { map: #map5 } :
+    (affineint, affineint) -> (affineint, affineint)
+  return
+}

mlir/test/IR/parser.mlir

@@ -112,22 +112,6 @@ mlfunc @mlfunc_with_args(%a : f16) {
   return  %a  // CHECK: return
 }
 
-// CHECK-LABEL: cfgfunc @cfgfunc_with_ops(f32) {
-cfgfunc @cfgfunc_with_ops(f32) {
-bb0(%a : f32):
-  // CHECK: %1 = "getTensor"() : () -> tensor<4x4x?xf32>
-  %t = "getTensor"() : () -> tensor<4x4x?xf32>
-
-  // CHECK: dim xxx, 2 : sometype
-  %t2 = "dim"(%t){index: 2} : (tensor<4x4x?xf32>) -> affineint
-
-  // CHECK: addf xx, yy : sometype
-  %x = "addf"(%a, %a) : (f32,f32) -> (f32)
-
-  // CHECK:   return
-  return
-}
-
 // CHECK-LABEL: mlfunc @loops() {
 mlfunc @loops() {
   // CHECK: for x = 1 to 100 step 2 {
@@ -175,24 +159,6 @@ bb42:       // CHECK: bb0:
   return
 }
 
-// CHECK-LABEL: cfgfunc @standard_instrs() {
-cfgfunc @standard_instrs() {
-bb42:       // CHECK: bb0:
-  // CHECK: %0 = "getTensor"() : () -> tensor<4x4x?xf32>
-  %42 = "getTensor"() : () -> tensor<4x4x?xf32>
-
-  // CHECK: dim xxx, 2 : sometype
-  %a = "dim"(%42){index: 2} : (tensor<4x4x?xf32>) -> affineint
-
-  %f = "Const"(){value: 1} : () -> f32
-  // CHECK: addf xx, yy : sometype
-  "addf"(%f, %f) : (f32,f32) -> f32
-
-  // TODO:   CHECK: FIXME: IMPLEMENT DEFAULT PRINTER
-  %x = "constant"(){value: 42} : () -> i32
-  return
-}
-
 // CHECK-LABEL: cfgfunc @ssa_values() -> (i16, i8) {
 cfgfunc @ssa_values() -> (i16, i8) {
 bb0:       // CHECK: bb0: