[WIP][mlir][EmitC] Model lvalues as a type in EmitC

Author: Simon Camphausen

mlir/include/mlir/Dialect/EmitC/IR/EmitC.td

@@ -835,6 +835,22 @@ def EmitC_LogicalOrOp : EmitC_BinaryOp<"logical_or", [CExpression]> {
   let assemblyFormat = "operands attr-dict `:` type(operands)";
 }
 
+def EmitC_LValueToRValueOp : EmitC_Op<"lvalue_to_rvalue", [
+  TypesMatchWith<"result type matches value type of 'operand'",
+                  "operand", "result",
+                  "::llvm::cast<LValueType>($_self).getValue()">
+]> {
+  let summary = "lvalue to rvalue conversion operation";
+  let description = [{}];
+
+  let arguments = (ins EmitC_LValueType:$operand);
+  let results = (outs AnyType:$result);
+
+  let assemblyFormat = "$operand attr-dict `:` type($operand)"; 
+
+  let hasVerifier = 1;
+}
+
 def EmitC_MulOp : EmitC_BinaryOp<"mul", [CExpression]> {
   let summary = "Multiplication operation";
   let description = [{
@@ -1009,7 +1025,7 @@ def EmitC_VariableOp : EmitC_Op<"variable", []> {
   }];
 
   let arguments = (ins EmitC_OpaqueOrTypedAttr:$value);
-  let results = (outs EmitCType);
+  let results = (outs EmitC_LValueType);
 
   let hasVerifier = 1;
 }
@@ -1137,7 +1153,7 @@ def EmitC_AssignOp : EmitC_Op<"assign", []> {
     ```
   }];
 
-  let arguments = (ins EmitCType:$var, EmitCType:$value);
+  let arguments = (ins EmitC_LValueType:$var, EmitCType:$value);
   let results = (outs);
 
   let hasVerifier = 1;
@@ -1243,7 +1259,7 @@ def EmitC_SubscriptOp : EmitC_Op<"subscript", []> {
       EmitC_PointerType]>,
     "the value to subscript">:$value,
     Variadic<EmitCType>:$indices);
-  let results = (outs EmitCType:$result);
+  let results = (outs EmitC_LValueType:$result);
 
   let builders = [
     OpBuilder<(ins "TypedValue<ArrayType>":$array, "ValueRange":$indices), [{

mlir/include/mlir/Dialect/EmitC/IR/EmitCTypes.td

@@ -83,6 +83,23 @@ def EmitC_ArrayType : EmitC_Type<"Array", "array", [ShapedTypeInterface]> {
   let hasCustomAssemblyFormat = 1;
 }
 
+def EmitC_LValueType : EmitC_Type<"LValue", "lvalue"> {
+  let summary = "EmitC lvalue type";
+
+  let description = [{
+    Values of this type can be assigned to and their address can be taken.
+  }];
+
+  let parameters = (ins "Type":$value);
+  let builders = [
+    TypeBuilderWithInferredContext<(ins "Type":$value), [{
+      return $_get(value.getContext(), value);
+    }]>
+  ];
+  let assemblyFormat = "`<` qualified($value) `>`";
+  let genVerifyDecl = 1;
+}
+
 def EmitC_OpaqueType : EmitC_Type<"Opaque", "opaque"> {
   let summary = "EmitC opaque type";
 

mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp

@@ -63,9 +63,10 @@ static SmallVector<Value> createVariablesForResults(T op,
 
   for (OpResult result : op.getResults()) {
     Type resultType = result.getType();
+    Type varType = emitc::LValueType::get(resultType);
     emitc::OpaqueAttr noInit = emitc::OpaqueAttr::get(context, "");
     emitc::VariableOp var =
-        rewriter.create<emitc::VariableOp>(loc, resultType, noInit);
+        rewriter.create<emitc::VariableOp>(loc, varType, noInit);
     resultVariables.push_back(var);
   }
 
@@ -76,57 +77,98 @@ static SmallVector<Value> createVariablesForResults(T op,
 // the current insertion point of given rewriter.
 static void assignValues(ValueRange values, SmallVector<Value> &variables,
                          PatternRewriter &rewriter, Location loc) {
-  for (auto [value, var] : llvm::zip(values, variables))
-    rewriter.create<emitc::AssignOp>(loc, var, value);
+  for (auto [value, var] : llvm::zip(values, variables)) {
+    assert(isa<emitc::LValueType>(var.getType()) &&
+           "expected var to be an lvalue type");
+    assert(!isa<emitc::LValueType>(value.getType()) &&
+           "expected value to not be an lvalue type");
+    auto assign = rewriter.create<emitc::AssignOp>(loc, var, value);
+
+    // TODO: Make sure this is safe, as this moves operations with memory
+    // effects.
+    if (auto op = dyn_cast_if_present<emitc::LValueToRValueOp>(
+            value.getDefiningOp())) {
+      rewriter.moveOpBefore(op, assign);
+    }
+  }
 }
 
-static void lowerYield(SmallVector<Value> &resultVariables,
-                       PatternRewriter &rewriter, scf::YieldOp yield) {
+static void lowerYield(SmallVector<Value> &variables, PatternRewriter &rewriter,
+                       scf::YieldOp yield) {
   Location loc = yield.getLoc();
   ValueRange operands = yield.getOperands();
 
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPoint(yield);
 
-  assignValues(operands, resultVariables, rewriter, loc);
+  assignValues(operands, variables, rewriter, loc);
 
   rewriter.create<emitc::YieldOp>(loc);
   rewriter.eraseOp(yield);
 }
 
+static void replaceUsers(PatternRewriter &rewriter,
+                         SmallVector<Value> fromValues,
+                         SmallVector<Value> toValues) {
+  OpBuilder::InsertionGuard guard(rewriter);
+  for (auto [from, to] : llvm::zip(fromValues, toValues)) {
+    assert(from.getType() == cast<emitc::LValueType>(to.getType()).getValue() &&
+           "expected types to match");
+
+    for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
+      Operation *op = operand.getOwner();
+      // Skip yield ops, as these get rewritten anyways.
+      if (isa<scf::YieldOp>(op)) {
+        continue;
+      }
+      Location loc = op->getLoc();
+
+      rewriter.setInsertionPoint(op);
+      Value rValue =
+          rewriter.create<emitc::LValueToRValueOp>(loc, from.getType(), to);
+      operand.set(rValue);
+    }
+  }
+}
+
 LogicalResult ForLowering::matchAndRewrite(ForOp forOp,
                                            PatternRewriter &rewriter) const {
   Location loc = forOp.getLoc();
 
-  // Create an emitc::variable op for each result. These variables will be
-  // assigned to by emitc::assign ops within the loop body.
-  SmallVector<Value> resultVariables =
-      createVariablesForResults(forOp, rewriter);
-  SmallVector<Value> iterArgsVariables =
-      createVariablesForResults(forOp, rewriter);
+  // Create an emitc::variable op for each result. These variables will be used
+  // for the results of the operations as well as the iter_args. They are
+  // assigned to by emitc::assign ops before the loop and at the end of the loop
+  // body.
+  SmallVector<Value> variables = createVariablesForResults(forOp, rewriter);
 
-  assignValues(forOp.getInits(), iterArgsVariables, rewriter, loc);
+  // Assign initial values to the iter arg variables.
+  assignValues(forOp.getInits(), variables, rewriter, loc);
 
-  emitc::ForOp loweredFor = rewriter.create<emitc::ForOp>(
-      loc, forOp.getLowerBound(), forOp.getUpperBound(), forOp.getStep());
+  // Replace users of the iter args with variables.
+  SmallVector<Value> iterArgs;
+  for (BlockArgument arg : forOp.getRegionIterArgs()) {
+    iterArgs.push_back(arg);
+  }
 
-  Block *loweredBody = loweredFor.getBody();
+  replaceUsers(rewriter, iterArgs, variables);
 
-  // Erase the auto-generated terminator for the lowered for op.
-  rewriter.eraseOp(loweredBody->getTerminator());
+  emitc::ForOp loweredFor = rewriter.create<emitc::ForOp>(
+      loc, forOp.getLowerBound(), forOp.getUpperBound(), forOp.getStep());
+  rewriter.eraseBlock(loweredFor.getBody());
 
-  SmallVector<Value> replacingValues;
-  replacingValues.push_back(loweredFor.getInductionVar());
-  replacingValues.append(iterArgsVariables.begin(), iterArgsVariables.end());
+  rewriter.inlineRegionBefore(forOp.getRegion(), loweredFor.getRegion(),
+                              loweredFor.getRegion().end());
+  Operation *terminator = loweredFor.getRegion().back().getTerminator();
+  lowerYield(variables, rewriter, cast<scf::YieldOp>(terminator));
 
-  rewriter.mergeBlocks(forOp.getBody(), loweredBody, replacingValues);
-  lowerYield(iterArgsVariables, rewriter,
-             cast<scf::YieldOp>(loweredBody->getTerminator()));
+  // Erase block arguments for iter_args.
+  loweredFor.getRegion().back().eraseArguments(1, variables.size());
 
-  // Copy iterArgs into results after the for loop.
-  assignValues(iterArgsVariables, resultVariables, rewriter, loc);
+  // Replace all users of the results with lazily created lvalue-to-rvalue
+  // ops.
+  replaceUsers(rewriter, forOp.getResults(), variables);
 
-  rewriter.replaceOp(forOp, resultVariables);
+  rewriter.eraseOp(forOp);
   return success();
 }
 
@@ -167,7 +209,7 @@ LogicalResult IfLowering::matchAndRewrite(IfOp ifOp,
 
   bool hasElseBlock = !elseRegion.empty();
 
-  auto loweredIf =
+  emitc::IfOp loweredIf =
       rewriter.create<emitc::IfOp>(loc, ifOp.getCondition(), false, false);
 
   Region &loweredThenRegion = loweredIf.getThenRegion();
@@ -178,7 +220,11 @@ LogicalResult IfLowering::matchAndRewrite(IfOp ifOp,
     lowerRegion(elseRegion, loweredElseRegion);
   }
 
-  rewriter.replaceOp(ifOp, resultVariables);
+  // Replace all users of the results with lazily created lvalue-to-rvalue
+  // ops.
+  replaceUsers(rewriter, ifOp.getResults(), resultVariables);
+
+  rewriter.eraseOp(ifOp);
   return success();
 }
 

mlir/lib/Dialect/EmitC/IR/EmitC.cpp

@@ -140,6 +140,8 @@ static LogicalResult verifyInitializationAttribute(Operation *op,
            << "string attributes are not supported, use #emitc.opaque instead";
 
   Type resultType = op->getResult(0).getType();
+  if (auto lType = dyn_cast<LValueType>(resultType))
+    resultType = lType.getValue();
   Type attrType = cast<TypedAttr>(value).getType();
 
   if (resultType != attrType)
@@ -203,18 +205,21 @@ LogicalResult ApplyOp::verify() {
 /// assigned-to variable type.
 LogicalResult emitc::AssignOp::verify() {
   Value variable = getVar();
-  Operation *variableDef = variable.getDefiningOp();
-  if (!variableDef ||
-      !llvm::isa<emitc::VariableOp, emitc::SubscriptOp>(variableDef))
+
+  if (!variable.getDefiningOp())
+    return emitOpError() << "cannot assign to block argument";
+  if (!llvm::isa<emitc::LValueType>(variable.getType()))
     return emitOpError() << "requires first operand (" << variable
-                         << ") to be a Variable or subscript";
-
-  Value value = getValue();
-  if (variable.getType() != value.getType())
-    return emitOpError() << "requires value's type (" << value.getType()
-                         << ") to match variable's type (" << variable.getType()
-                         << ")";
-  if (isa<ArrayType>(variable.getType()))
+                         << ") to be an lvalue";
+
+  Type valueType = getValue().getType();
+  Type variableType = variable.getType().cast<emitc::LValueType>().getValue();
+  if (variableType != valueType)
+    return emitOpError() << "requires value's type (" << valueType
+                         << ") to match variable's type (" << variableType
+                         << ")\n  variable: " << variable
+                         << "\n  value: " << getValue() << "\n";
+  if (isa<ArrayType>(variableType))
     return emitOpError() << "cannot assign to array type";
   return success();
 }
@@ -769,6 +774,47 @@ LogicalResult emitc::LiteralOp::verify() {
     return emitOpError() << "value must not be empty";
   return success();
 }
+
+//===----------------------------------------------------------------------===//
+// LValueToRValueOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult emitc::LValueToRValueOp::verify() {
+  Type operandType = getOperand().getType();
+  Type resultType = getResult().getType();
+  if (!llvm::isa<emitc::LValueType>(operandType))
+    return emitOpError("operand must be a lvalue");
+  if (llvm::cast<emitc::LValueType>(operandType).getValue() != resultType)
+    return emitOpError("types must match");
+
+  Value result = getResult();
+  if (!result.hasOneUse()) {
+    int numUses = std::distance(result.use_begin(), result.use_end());
+    return emitOpError("must have exactly one use, but got ") << numUses;
+  }
+
+  Block *block = result.getParentBlock();
+
+  Operation *user = *result.getUsers().begin();
+  Block *userBlock = user->getBlock();
+
+  if (block != userBlock) {
+    return emitOpError("user must be in the same block");
+  }
+
+  // for (auto it = block.begin(), e = std::prev(block.end()); it != e; it++) {
+  //   if (*it == this)
+  // }
+
+  // TODO: To model this op correctly as a memory read of the lvalue, we
+  // should additionally ensure that the single use of the op follows immediatly
+  // on this definition. Alternativly we could alter emitc ops to implicitly
+  // support lvalues. This would make it harder to do partial conversions and
+  // mix dialects though.
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // SubOp
 //===----------------------------------------------------------------------===//
@@ -964,6 +1010,20 @@ emitc::ArrayType::cloneWith(std::optional<ArrayRef<int64_t>> shape,
   return emitc::ArrayType::get(*shape, elementType);
 }
 
+//===----------------------------------------------------------------------===//
+// LValueType
+//===----------------------------------------------------------------------===//
+
+LogicalResult mlir::emitc::LValueType::verify(
+    llvm::function_ref<mlir::InFlightDiagnostic()> emitError,
+    mlir::Type value) {
+  if (llvm::isa<emitc::LValueType>(value)) {
+    return emitError()
+           << "!emitc.lvalue type cannot be nested inside another type";
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // OpaqueType
 //===----------------------------------------------------------------------===//