[RefE2E] Add interesting control flow example.

This also required adding a lowering for ForOp in our tensor->memref
conversion.

lib/E2E/TensorToMemref/LowerStructuralToMemref.cpp

@@ -50,6 +50,33 @@ public:
 };
 } // namespace
 
+namespace {
+// This is a type conversion similar to CallOpSignatureConversion.
+class LowerForOpTypes : public OpConversionPattern<scf::ForOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(scf::ForOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    SmallVector<Type, 6> newResultTypes;
+    for (auto type : op.getResultTypes()) {
+      Type newType = typeConverter->convertType(type);
+      if (!newType)
+        return rewriter.notifyMatchFailure(op, "not a 1:1 type conversion");
+      newResultTypes.push_back(newType);
+    }
+    rewriter.updateRootInPlace(op, [&] {
+      for (auto t : llvm::zip(op.getResults(), newResultTypes))
+        std::get<0>(t).setType(std::get<1>(t));
+      auto bodyArgs = op.getBody()->getArguments();
+      for (auto t : llvm::zip(llvm::drop_begin(bodyArgs, 1), newResultTypes))
+        std::get<0>(t).setType(std::get<1>(t));
+    });
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 // This is a type conversion similar to CallOpSignatureConversion.
 class LowerSelectOpTypes : public OpConversionPattern<SelectOp> {
@@ -151,6 +178,7 @@ class LowerStructuralToMemref
 
     patterns.insert<LowerSelectOpTypes>(typeConverter, context);
     patterns.insert<LowerIfOpTypes>(typeConverter, context);
+    patterns.insert<LowerForOpTypes>(typeConverter, context);
     patterns.insert<LowerTensorToMemrefOp>(typeConverter, context);
     patterns.insert<LowerMemrefToTensorOp>(typeConverter, context);
     target.addIllegalOp<tcp::TensorToMemrefOp>();

test/E2E/lower-structural-to-memref.mlir

@@ -42,6 +42,19 @@ func @if(%pred: i1, %true_val: tensor<?xf32>, %false_val: tensor<?xf32>) -> tens
   return %0 : tensor<?xf32>
 }
 
+// CHECK-LABEL: func @for(%arg0: memref<f32>, %arg1: index, %arg2: index, %arg3: index) -> memref<f32> {
+// CHECK-NEXT:    %[[RET:.*]] = scf.for %arg4 = %arg1 to %arg2 step %arg3 iter_args(%arg5 = %arg0) -> (memref<f32>) {
+// CHECK-NEXT:      scf.yield %arg5 : memref<f32>
+// CHECK-NEXT:    }
+// CHECK-NEXT:    return %[[RET]] : memref<f32>
+// CHECK-NEXT:  }
+func @for(%arg0: tensor<f32>, %lb: index, %ub: index, %step: index) -> tensor<f32> {
+  %ret = scf.for %iv = %lb to %ub step %step iter_args(%iter = %arg0) -> tensor<f32> {
+    scf.yield %iter : tensor<f32>
+  }
+  return %ret : tensor<f32>
+}
+
 
 // Test the interactions with materializations.
 // Note: this pass never actually expects IR with memref argument types.

test/npcomp-run-mlir/control-flow-basic.mlir

@@ -0,0 +1,29 @@
+// RUN: npcomp-run-mlir %s \
+// RUN:   -invoke pow2 \
+// RUN:   -arg-value="dense<8.0> : tensor<f32>" \
+// RUN:   -shared-libs=%npcomp_runtime_shlib 2>&1 \
+// RUN:   | FileCheck %s
+
+// 2^8 == 256
+// CHECK: output #0: dense<2.560000e+02> : tensor<f32>
+func @pow2(%arg0: tensor<f32>) -> tensor<f32> {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+
+  // Slight awkwardness: convert the tensor<f32> to an index.
+  // TODO: Allow passing plain integers/floats (not tensors) at
+  // calling convention boundaries.
+
+  %num_iters_float = extract_element %arg0[] : tensor<f32>
+  %num_iters_i32 = fptosi %num_iters_float : f32 to i32
+  %num_iters = index_cast %num_iters_i32 : i32 to index
+
+  // Repeatedly add the value to itself %num_iters times.
+  %tensor_c1 = constant dense<1.0> : tensor<f32>
+  %ret = scf.for %iv = %c0 to %num_iters step %c1 iter_args(%iter = %tensor_c1) -> tensor<f32> {
+    %doubled = tcf.add %iter, %iter : (tensor<f32>, tensor<f32>) -> tensor<f32>
+    scf.yield %doubled : tensor<f32>
+  }
+  return %ret : tensor<f32>
+}
+