Placeholder operations for defining dummy domains.

Introduce a new PlaceholderOp operation that defines a range of unknown size. This operation is used when introducing new operations during code generation. In many cases, the domain of new operations will not be used, either:
* because the operation is a ComputeOp with a loop-nest and loops sizes are already defined by other operations.
* because the operation is not a compute operation so its domain won't be used to introduce loops or buffers.

PiperOrigin-RevId: 361100458

Author: ulysseB

canonicalization_patterns.cc

@@ -417,6 +417,11 @@ void SairDynRangeOp::getCanonicalizationPatterns(
   patterns.insert<SimplifySairOperands>();
 }
 
+void SairPlaceholderOp::getCanonicalizationPatterns(
+    mlir::OwningRewritePatternList &patterns, mlir::MLIRContext *context) {
+  patterns.insert<SimplifySairOperands>();
+}
+
 void SairStaticRangeOp::getCanonicalizationPatterns(
     mlir::OwningRewritePatternList &patterns, mlir::MLIRContext *context) {
   patterns.insert<SimplifySairOperands>();

sair_ops.cc

@@ -201,6 +201,22 @@ ParseResult ParseStaticRangeOp(mlir::OpAsmParser &parser,
                  parser.addTypeToList(type, result.types));
 }
 
+// Parses the placeholder dimension. This operation has an iteration domain and
+// returns a range value. The syntax is the following.
+//
+// placeholder-op ::= `sair.placeholder` domain : range-type
+//
+ParseResult ParsePlaceholderOp(mlir::OpAsmParser &parser,
+                               mlir::OperationState &result) {
+  llvm::SmallVector<mlir::OpAsmParser::OperandType> domain;
+  RangeType type;
+
+  return mlir::failure(ParseDomain(parser, domain) ||
+                       parser.parseColonType<RangeType>(type) ||
+                       parser.addTypeToList(type, result.types) ||
+                       ResolveDomain(parser, type.Shape(), domain, result));
+}
+
 // Parses the copy operation. This operation has an iteration domain and
 // accesses a single Sair value. The syntax for the operation is the following.
 //
@@ -616,6 +632,12 @@ void Print(SairStaticRangeOp op, OpAsmPrinter &printer) {
   printer << " : " << op.getType();
 }
 
+static void Print(SairPlaceholderOp op, mlir::OpAsmPrinter &printer) {
+  printer << SairPlaceholderOp::getOperationName();
+  PrintDomain(op.domain(), printer);
+  printer << " : " << op.range().getType();
+}
+
 // Prints the copy operation.
 void Print(SairCopyOp op, OpAsmPrinter &printer) {
   printer << SairCopyOp::getOperationName();
@@ -892,7 +914,7 @@ llvm::SmallBitVector SairStoreToMemRefOp::DimsDependingOnOperand(
 
 ParseResult ParseDomain(
     mlir::OpAsmParser &parser,
-    llvm::SmallVector<mlir::OpAsmParser::OperandType, 4> &dimensions) {
+    llvm::SmallVectorImpl<mlir::OpAsmParser::OperandType> &dimensions) {
   if (failed(parser.parseOptionalLSquare())) return success();
   do {
     std::string dim_name = "d" + std::to_string(dimensions.size());
@@ -1559,6 +1581,10 @@ llvm::SmallVector<int, 2> SairDynRangeOp::SubDomains() {
   return {static_cast<int>(domain().size())};
 }
 
+llvm::SmallVector<int, 2> SairPlaceholderOp::SubDomains() {
+  return {static_cast<int>(domain().size())};
+}
+
 llvm::SmallVector<int, 2> SairCopyOp::SubDomains() {
   return {static_cast<int>(domain().size())};
 }
@@ -1672,6 +1698,14 @@ SairOp SairStaticRangeOp::ReCreateWithNewDomain(
       "not called by NormalizeLoops because the op defines a dimension");
 }
 
+SairOp SairPlaceholderOp::ReCreateWithNewDomain(
+    llvm::ArrayRef<llvm::SmallVector<mlir::Value, 4>> new_domains,
+    DomainShapeAttr new_shape, MappingAttr new_to_old_mapping,
+    mlir::OpBuilder &builder) {
+  llvm_unreachable(
+      "not called by NormalizeLoops because the op defines a dimension");
+}
+
 SairOp SairCopyOp::ReCreateWithNewDomain(
     llvm::ArrayRef<llvm::SmallVector<mlir::Value, 4>> new_domains,
     DomainShapeAttr new_shape, MappingAttr new_to_old_mapping,

sair_ops.h

@@ -53,7 +53,7 @@ namespace sair {
 // '['.
 ParseResult ParseDomain(
     mlir::OpAsmParser &parser,
-    llvm::SmallVector<mlir::OpAsmParser::OperandType, 4> &dimensions);
+    llvm::SmallVectorImpl<mlir::OpAsmParser::OperandType> &dimensions);
 // Resolves the operands that consitute the dimensions of an iteration domain
 // and registers them in 'result'.
 ParseResult ResolveDomain(mlir::OpAsmParser &parser,

sair_ops.td

@@ -106,6 +106,30 @@ def SairStaticRangeOp : SairOp<"static_range", [NoSideEffect, SairRangeOp]> {
   }];
 }
 
+def SairPlaceholderOp : SairOp<"placeholder", [NoSideEffect]> {
+  let summary = "Placeholder for an iteration dimension";
+
+  let description = [{
+    Defines an iteration dimension that will be replaced by an actual iteration
+    dimension during loop-normalization pass. This is used to introduce
+    operations in specific loop nests before dimensions defining loop ranges are
+    introduced.
+
+    If a dimension bound to a placeholder operation appears in loop iterators
+    mapping, the dimension must be defined by another operation with a
+    non-placeholder operation. In that sense, they behave similarly to `none` in
+    loop-nest description.
+  }];
+
+  let arguments = (ins Variadic<SairRange>:$domain);
+  let results = (outs SairRange:$range);
+
+  let parser = [{return ParsePlaceholderOp(parser, result);}];
+  let printer = [{return Print(*this, p);}];
+
+  DerivedAttr shape = SairResultDomainShapeAttr;
+}
+
 def SairCopyOp : SairOp<"copy", [
     NoSideEffect,
     SairComputeOp,

test/introduce_loops_invalid.mlir

@@ -135,3 +135,27 @@ func @size_not_in_register(%arg0: index) {
   }
   return
 }
+
+// -----
+
+func @placeholder() {
+  sair.program {
+    %0 = sair.static_range 8 : !sair.range
+    // expected-error @+1 {{placeholders must be replaced by actual dimensions before introducing loops}}
+    %1 = sair.placeholder : !sair.range
+    sair.map[d0:%1] attributes {
+      loop_nest = [{name = "loopA", iter = #sair.mapping_expr<d0>}]
+    } {
+      ^bb0(%arg1: index):
+        sair.return
+    } : #sair.shape<d0:range>, () -> ()
+    sair.map[d0:%0] attributes {
+      loop_nest = [{name = "loopA", iter = #sair.mapping_expr<d0>}]
+    } {
+      ^bb0(%arg1: index):
+        sair.return
+    } : #sair.shape<d0:range>, () -> ()
+    sair.exit
+  }
+  return
+}

test/invalid.mlir

@@ -1519,3 +1519,18 @@ func @map_reduce_init_result_storages_must_match(%arg0: f32) {
   }
   return
 }
+
+// -----
+
+func @placeholder_loop_nest_unspecified(%arg0: f32) {
+  sair.program {
+    %0 = sair.from_scalar %arg0 : !sair.value<(), f32>
+    %1 = sair.placeholder : !sair.range
+    // expected-error @+1 {{loop "loopA" iterator is not fully specified}}
+    %2 = sair.copy[d0:%1] %0 {
+      loop_nest = [{name = "loopA", iter = #sair.mapping_expr<d0>}]
+    } : !sair.value<d0:range, f32>
+    sair.exit
+  }
+  return
+}

test/roundtrip.mlir

@@ -627,3 +627,32 @@ func @storage_stripe(%arg0: f32) {
   }
   return
 }
+
+// CHECK-LABEL: @placeholder
+func @placeholder(%arg0: f32) {
+  sair.program {
+    %0 = sair.static_range 2 : !sair.range
+    %1 = sair.placeholder[d0:%0] : !sair.range<d0:range>
+    %2 = sair.from_scalar %arg0 : !sair.value<(), f32>
+    %3 = sair.copy[d0:%0, d1:%1] %2 : !sair.value<d0:range x d1:range(d0), f32>
+    sair.exit
+  }
+  return
+}
+
+// CHECK-LABEL: @placeholder_with_loop_nest
+func @placeholder_with_loop_nest(%arg0: f32) {
+  sair.program {
+    %0 = sair.static_range 2 : !sair.range
+    %1 = sair.placeholder : !sair.range
+    %2 = sair.from_scalar %arg0 : !sair.value<(), f32>
+    %3 = sair.copy[d0:%0] %2 {
+      loop_nest = [{name = "loopA", iter = #sair.mapping_expr<d0>}]
+    } : !sair.value<d0:range, f32>
+    %4 = sair.copy[d0:%1] %2 {
+      loop_nest = [{name = "loopA", iter = #sair.mapping_expr<d0>}]
+    } : !sair.value<d0:range, f32>
+    sair.exit
+  }
+  return
+}

transforms/introduce_loops.cc

@@ -496,7 +496,14 @@ mlir::LogicalResult IntroduceLoop(SairMapOp op,
   LoopAttr loop = loop_nest.back().cast<LoopAttr>();
 
   int dimension = loop.iter().cast<MappingDimExpr>().dimension();
-  RangeOp range = cast<RangeOp>(op.domain()[dimension].getDefiningOp());
+  mlir::Operation *dimension_op = op.domain()[dimension].getDefiningOp();
+  if (isa<SairPlaceholderOp>(dimension_op)) {
+    return dimension_op->emitError()
+           << "placeholders must be replaced by actual dimensions before "
+              "introducing loops";
+  }
+
+  RangeOp range = cast<RangeOp>(dimension_op);
   MappingAttr range_mapping =
       op.shape().Dimension(dimension).dependency_mapping().ResizeUseDomain(
           op.domain().size() - 1);

transforms/normalize_loops.cc

@@ -272,8 +272,8 @@ class NormalizeLoopsPass : public NormalizeLoopsPassBase<NormalizeLoopsPass> {
       llvm::DenseMap<mlir::Attribute, std::pair<mlir::Value, DomainShapeDim>>
           loop_range_cache;
       program.walk([&](SairOp op) {
-        // Do not normalize range operations.
-        if (isa<RangeOp>(op.getOperation())) return;
+        // Do not normalize range and placeholder operations.
+        if (isa<RangeOp, SairPlaceholderOp>(op.getOperation())) return;
         if (mlir::failed(NormalizeLoops(op, iteration_spaces, builder,
                                         loop_range_cache))) {
           signalPassFailure();

util.cc

@@ -16,6 +16,7 @@
 
 #include "llvm/ADT/STLExtras.h"
 #include "sair_op_interfaces.h"
+#include "sair_ops.h"
 
 namespace sair {
 
@@ -79,6 +80,10 @@ mlir::LogicalResult ResolveUnificationConstraint(
     llvm::SmallVectorImpl<ValueAccess> &target_domain) {
   mlir::MLIRContext *context = constraint.getContext();
 
+  // Ignore placeholders.
+  mlir::Operation *defining_op = dimension.value.getDefiningOp();
+  if (isa<SairPlaceholderOp>(defining_op)) return mlir::success();
+
   if (constraint.isa<MappingNoneExpr>()) {
     constraint = MappingDimExpr::get(target_domain.size(), context);
     target_domain.push_back(dimension);