imporve doc, test, change trait check.

Author: linuxlonelyeagle

mlir/docs/Dialects/Affine.md

@@ -69,7 +69,7 @@ immediately enclosed by the latter),
 3. a value that dominates the `AffineScope` op enclosing the value's
 use,
 4. the result of a constant operation,
-5. the result of an `AffineSymbol` op,
+5. the result of an op with [`AffineSymbol`](Traits.md#AffineSymbol) trait,
 6. the result of an
 [`affine.apply` operation](#affineapply-mliraffineapplyop) that recursively takes as
 arguments any valid symbolic identifiers, or

mlir/docs/Traits/_index.md

@@ -223,6 +223,16 @@ affine.load, and affine.store. The polyhedral scope defined by an operation with
 this trait includes all operations in its region excluding operations that are
 nested inside of other operations that themselves have this trait.
 
+*   `OpTrait::AffineSymbol` -- `AffineSymbol`
+
+This trait will make the results of the operation that holds it become valid symbols.
+It complements the AffineScope trait, allowing SSA values in the region of operation
+to be used as operands for various Affine dialect operations depending on the AffineScope
+trait, without altering the properties of the operation. In contrast, the AffineSymbol
+trait directly changes the properties of the operations. Additionally, the AffineScope
+trait cannot make SSA values in the nested region of operations valid symbols, operations
+can hold the AffineSymbol trait, then their results whill be vaild symbols.
+
 ### AutomaticAllocationScope
 
 *   `OpTrait::AutomaticAllocationScope` -- `AutomaticAllocationScope`

mlir/include/mlir/IR/OpDefinition.h

@@ -1269,8 +1269,10 @@ class AffineScope : public TraitBase<ConcreteType, AffineScope> {
   }
 };
 
-/// A trait of operation. Any operation holds the AffineSymbol, and its result
-/// can be used as a symbol.
+/// A trait of operation. The results of an operation with the AffineSymbol
+/// trait can be used as symbols for the purposes for affine dialect purposes.
+/// This trait will make the results of the operation that holds it become valid
+/// symbols. For more details, see `Traits.md#AffineSymbol`.
 template <typename ConcreteType>
 class AffineSymbol : public TraitBase<ConcreteType, AffineSymbol> {
 public:

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -410,6 +410,7 @@ bool mlir::affine::isValidSymbol(Value value) {
 /// A value can be used as a symbol for `region` iff it meets one of the
 /// following conditions:
 /// *) It is a constant.
+/// *) It is a result, its defineOp is holded by `AffineSymbol` Trait.
 /// *) It is the result of an affine apply operation with symbol arguments.
 /// *) It is a result of the dim op on a memref whose corresponding size is
 ///    a valid symbol.
@@ -443,7 +444,7 @@ bool mlir::affine::isValidSymbol(Value value, Region *region) {
   if (matchPattern(defOp, m_Constant(&operandCst)))
     return true;
 
-  // AffineScope Op.
+  // AffineSymbol Op.
   if (defOp->hasTrait<OpTrait::AffineSymbol>())
     return true;
 

mlir/lib/IR/Operation.cpp

@@ -1391,10 +1391,17 @@ LogicalResult OpTrait::impl::verifyIsIsolatedFromAbove(Operation *isolatedOp) {
 }
 
 LogicalResult OpTrait::impl::verifyIndexResultType(Operation *op) {
-  if (op->getNumResults() != 1)
-    op->emitError("operation's result number should be 1.");
-  if (!mlir::isa<IndexType>(op->getResult(0).getType()))
-    op->emitError("operation's result type should be index.");
+  ResultRange results = op->getResults();
+  if (results.empty()) {
+    op->emitError("operation must have at least one result");
+    return failure();
+  }
+  for (OpResult result : results) {
+    if (!mlir::isa<IndexType>(result.getType())) {
+      op->emitError("operation's result type should be index");
+      return failure();
+    }
+  }
   return success();
 }
 

mlir/test/Dialect/Affine/ops.mlir

@@ -331,6 +331,8 @@ module attributes {gpu.container_module} {
 
 // CHECK: #[[$ATTR_0:.+]] = affine_map<()[s0] -> (s0 mod 32)>
 
+// CHECK-LABEL: @affine_thread_id
+
 module {
   gpu.module @gpu {
     gpu.func @affine_thread_id(%arg0: memref<?x?xf32>) kernel {
@@ -350,13 +352,12 @@ module {
   }
 }
 
-// CHECK-LABEL:     @affine_thread_id
-// CHECK-SAME:        (%[[VAL_0:.*]]: memref<?x?xf32>) kernel {
-// CHECK:             %[[VAL_1:.*]] = arith.constant 3 : index
-// CHECK:             %[[VAL_2:.*]] = memref.dim %[[VAL_0]], %[[VAL_1]] : memref<?x?xf32>
-// CHECK:             %[[VAL_3:.*]] = arith.constant 0 : index
-// CHECK:             affine.for %[[VAL_4:.*]] = %[[VAL_3]] to %[[VAL_2]] step 32 {
-// CHECK:               %[[VAL_5:.*]] = gpu.thread_id  x
-// CHECK:               %[[VAL_6:.*]] = affine.apply #[[$ATTR_0]](){{\[}}%[[VAL_5]]]
-// CHECK:               %[[VAL_7:.*]] = arith.constant 128 : index
-// CHECK:               affine.for %[[VAL_8:.*]] = %[[VAL_6]] to %[[VAL_7]] step 8 {
+// CHECK-SAME: (%[[VAL_0:.*]]: memref<?x?xf32>) kernel {
+// CHECK: %[[VAL_1:.*]] = arith.constant 3 : index
+// CHECK: %[[VAL_2:.*]] = memref.dim %[[VAL_0]], %[[VAL_1]] : memref<?x?xf32>
+// CHECK: %[[VAL_3:.*]] = arith.constant 0 : index
+// CHECK: affine.for %[[VAL_4:.*]] = %[[VAL_3]] to %[[VAL_2]] step 32 {
+// CHECK: %[[VAL_5:.*]] = gpu.thread_id  x
+// CHECK: %[[VAL_6:.*]] = affine.apply #[[$ATTR_0]](){{\[}}%[[VAL_5]]]
+// CHECK: %[[VAL_7:.*]] = arith.constant 128 : index
+// CHECK: affine.for %[[VAL_8:.*]] = %[[VAL_6]] to %[[VAL_7]] step 8 {