[MLIR][Transform][SMT] Allow for declarative computations in schedules (#160895)

By allowing `transform.smt.constrain_params`'s region to yield SMT-vars,
op instances can declare relationships, through constraints, on incoming
params-as-SMT-vars and outgoing SMT-vars-as-params. This makes it
possible to declare that computations on params should be performed.

The semantics are that the yielded SMT-vars should be from any valid
satisfying assignment/model of the constraints in the region.

Author: rolfmorel

mlir/include/mlir/Dialect/SMT/IR/SMTOps.td

@@ -220,8 +220,6 @@ def YieldOp : SMTOp<"yield", [
   Pure,
   Terminator,
   ReturnLike,
-  ParentOneOf<["smt::SolverOp", "smt::CheckOp",
-               "smt::ForallOp", "smt::ExistsOp"]>,
 ]> {
   let summary = "terminator operation for various regions of SMT operations";
   let arguments = (ins Variadic<AnyType>:$values);

mlir/include/mlir/Dialect/Transform/SMTExtension/SMTExtensionOps.h

@@ -10,6 +10,7 @@
 #define MLIR_DIALECT_TRANSFORM_SMTEXTENSION_SMTEXTENSIONOPS_H
 
 #include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/Dialect/SMT/IR/SMTOps.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
 #include "mlir/IR/OpDefinition.h"

mlir/include/mlir/Dialect/Transform/SMTExtension/SMTExtensionOps.td

@@ -16,22 +16,28 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
 def ConstrainParamsOp : Op<Transform_Dialect, "smt.constrain_params", [
   DeclareOpInterfaceMethods<TransformOpInterface>,
   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-  NoTerminator
+  SingleBlockImplicitTerminator<"::mlir::smt::YieldOp">
 ]> {
   let cppNamespace = [{ mlir::transform::smt }];
 
   let summary = "Express contraints on params interpreted as symbolic values";
   let description = [{
     Allows expressing constraints on params using the SMT dialect.
 
-    Each Transform dialect param provided as an operand has a corresponding
+    Each Transform-dialect param provided as an operand has a corresponding
     argument of SMT-type in the region. The SMT-Dialect ops in the region use
-    these arguments as operands.
+    these params-as-SMT-vars as operands, thereby expressing relevant
+    constraints on their allowed values.
+
+    Computations w.r.t. passed-in params can also be expressed through the
+    region's SMT-ops. Namely, the constraints express relationships to other
+    SMT-variables which can then be yielded from the region (with `smt.yield`).
 
     The semantics of this op is that all the ops in the region together express
     a constraint on the params-interpreted-as-smt-vars. The op fails in case the
     expressed constraint is not satisfiable per SMTLIB semantics. Otherwise the
-    op succeeds.
+    op succeeds and any one satisfying assignment is used to map the
+    SMT-variables yielded in the region to `transform.param`s.
 
     ---
 
@@ -42,9 +48,10 @@ def ConstrainParamsOp : Op<Transform_Dialect, "smt.constrain_params", [
   }];
 
   let arguments = (ins Variadic<TransformParamTypeInterface>:$params);
+  let results = (outs Variadic<TransformParamTypeInterface>:$results);
   let regions = (region SizedRegion<1>:$body);
   let assemblyFormat =
-      "`(` $params `)` attr-dict `:` type(operands) $body";
+      "`(` $params `)` attr-dict `:` functional-type(operands, results) $body";
 
   let hasVerifier = 1;
 }

mlir/lib/Dialect/Transform/SMTExtension/SMTExtensionOps.cpp

@@ -8,8 +8,8 @@
 
 #include "mlir/Dialect/Transform/SMTExtension/SMTExtensionOps.h"
 #include "mlir/Dialect/SMT/IR/SMTDialect.h"
-#include "mlir/Dialect/Transform/IR/TransformOps.h"
-#include "mlir/Dialect/Transform/SMTExtension/SMTExtension.h"
+#include "mlir/Dialect/SMT/IR/SMTOps.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
 
 using namespace mlir;
 
@@ -23,6 +23,7 @@ using namespace mlir;
 void transform::smt::ConstrainParamsOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   onlyReadsHandle(getParamsMutable(), effects);
+  producesHandle(getResults(), effects);
 }
 
 DiagnosedSilenceableFailure
@@ -37,19 +38,95 @@ transform::smt::ConstrainParamsOp::apply(transform::TransformRewriter &rewriter,
   //       and allow for users to attach their own implementation, which would,
   //       e.g., translate the ops to SMTLIB and hand that over to the user's
   //       favourite solver. This requires changes to the dialect's verifier.
-  return emitDefiniteFailure() << "op does not have interpreted semantics yet";
+  return emitSilenceableFailure(getLoc())
+         << "op does not have interpreted semantics yet";
 }
 
 LogicalResult transform::smt::ConstrainParamsOp::verify() {
+  auto yieldTerminator =
+      dyn_cast<mlir::smt::YieldOp>(getRegion().front().back());
+  if (!yieldTerminator)
+    return emitOpError() << "expected '"
+                         << mlir::smt::YieldOp::getOperationName()
+                         << "' as terminator";
+
+  auto checkTypes = [](size_t idx, Type smtType, StringRef smtDesc,
+                       Type paramType, StringRef paramDesc,
+                       auto *atOp) -> InFlightDiagnostic {
+    if (!isa<mlir::smt::BoolType, mlir::smt::IntType, mlir::smt::BitVectorType>(
+            smtType))
+      return atOp->emitOpError() << "the type of " << smtDesc << " #" << idx
+                                 << " is expected to be either a !smt.bool, a "
+                                    "!smt.int, or a !smt.bv";
+
+    assert(isa<TransformParamTypeInterface>(paramType) &&
+           "ODS specifies params' type should implement param interface");
+    if (isa<transform::AnyParamType>(paramType))
+      return {}; // No further checks can be done.
+
+    // NB: This cast must succeed as long as the only implementors of
+    //     TransformParamTypeInterface are AnyParamType and ParamType.
+    Type typeWrappedByParam = cast<ParamType>(paramType).getType();
+
+    if (isa<mlir::smt::IntType>(smtType)) {
+      if (!isa<IntegerType>(typeWrappedByParam))
+        return atOp->emitOpError()
+               << "the type of " << smtDesc << " #" << idx
+               << " is !smt.int though the corresponding " << paramDesc
+               << " type (" << paramType << ") is not wrapping an integer type";
+    } else if (isa<mlir::smt::BoolType>(smtType)) {
+      auto wrappedIntType = dyn_cast<IntegerType>(typeWrappedByParam);
+      if (!wrappedIntType || wrappedIntType.getWidth() != 1)
+        return atOp->emitOpError()
+               << "the type of " << smtDesc << " #" << idx
+               << " is !smt.bool though the corresponding " << paramDesc
+               << " type (" << paramType << ") is not wrapping i1";
+    } else if (auto bvSmtType = dyn_cast<mlir::smt::BitVectorType>(smtType)) {
+      auto wrappedIntType = dyn_cast<IntegerType>(typeWrappedByParam);
+      if (!wrappedIntType || wrappedIntType.getWidth() != bvSmtType.getWidth())
+        return atOp->emitOpError()
+               << "the type of " << smtDesc << " #" << idx << " is " << smtType
+               << " though the corresponding " << paramDesc << " type ("
+               << paramType
+               << ") is not wrapping an integer type of the same bitwidth";
+    }
+
+    return {};
+  };
+
   if (getOperands().size() != getBody().getNumArguments())
     return emitOpError(
         "must have the same number of block arguments as operands");
 
+  for (auto [idx, operandType, blockArgType] :
+       llvm::enumerate(getOperandTypes(), getBody().getArgumentTypes())) {
+    InFlightDiagnostic typeCheckResult =
+        checkTypes(idx, blockArgType, "block arg", operandType, "operand",
+                   /*atOp=*/this);
+    if (LogicalResult(typeCheckResult).failed())
+      return typeCheckResult;
+  }
+
   for (auto &op : getBody().getOps()) {
     if (!isa<mlir::smt::SMTDialect>(op.getDialect()))
       return emitOpError(
           "ops contained in region should belong to SMT-dialect");
   }
 
+  if (yieldTerminator->getNumOperands() != getNumResults())
+    return yieldTerminator.emitOpError()
+           << "expected terminator to have as many operands as the parent op "
+              "has results";
+
+  for (auto [idx, termOperandType, resultType] : llvm::enumerate(
+           yieldTerminator->getOperands().getType(), getResultTypes())) {
+    InFlightDiagnostic typeCheckResult =
+        checkTypes(idx, termOperandType, "terminator operand",
+                   cast<transform::ParamType>(resultType), "result",
+                   /*atOp=*/&yieldTerminator);
+    if (LogicalResult(typeCheckResult).failed())
+      return typeCheckResult;
+  }
+
   return success();
 }

mlir/python/mlir/dialects/transform/smt.py

@@ -19,6 +19,7 @@
 class ConstrainParamsOp(ConstrainParamsOp):
     def __init__(
         self,
+        results: Sequence[Type],
         params: Sequence[transform.AnyParamType],
         arg_types: Sequence[Type],
         loc=None,
@@ -27,6 +28,7 @@ def __init__(
         if len(params) != len(arg_types):
             raise ValueError(f"{params=} not same length as {arg_types=}")
         super().__init__(
+            results,
             params,
             loc=loc,
             ip=ip,
@@ -36,3 +38,13 @@ def __init__(
     @property
     def body(self) -> Block:
         return self.regions[0].blocks[0]
+
+
+def constrain_params(
+    results: Sequence[Type],
+    params: Sequence[transform.AnyParamType],
+    arg_types: Sequence[Type],
+    loc=None,
+    ip=None,
+):
+    return ConstrainParamsOp(results, params, arg_types, loc=loc, ip=ip)

mlir/test/Dialect/Transform/test-smt-extension-invalid.mlir

@@ -1,11 +1,40 @@
 // RUN: mlir-opt %s --transform-interpreter --split-input-file --verify-diagnostics
 
+// CHECK-LABEL: @incorrect terminator
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @operands_not_one_to_one_with_vars(%arg0: !transform.any_op {transform.readonly}) {
+    %param_as_param = transform.param.constant 42 -> !transform.param<i64>
+    // expected-error@below {{op expected 'smt.yield' as terminator}}
+    transform.smt.constrain_params(%param_as_param) : (!transform.param<i64>) -> () {
+      ^bb0(%param_as_smt_var: !smt.int):
+      transform.yield
+    }
+    transform.yield
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @operands_not_one_to_one_with_vars
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @operands_not_one_to_one_with_vars(%arg0: !transform.any_op {transform.readonly}) {
+    %param_as_param = transform.param.constant 42 -> !transform.param<i64>
+    // expected-error@below {{must have the same number of block arguments as operands}}
+    transform.smt.constrain_params(%param_as_param) : (!transform.param<i64>) -> () {
+      ^bb0(%param_as_smt_var: !smt.int, %param_as_another_smt_var: !smt.int):
+    }
+    transform.yield
+  }
+}
+
+// -----
+
 // CHECK-LABEL: @constraint_not_using_smt_ops
 module attributes {transform.with_named_sequence} {
   transform.named_sequence @constraint_not_using_smt_ops(%arg0: !transform.any_op {transform.readonly}) {
     %param_as_param = transform.param.constant 42 -> !transform.param<i64>
     // expected-error@below {{ops contained in region should belong to SMT-dialect}}
-    transform.smt.constrain_params(%param_as_param) : !transform.param<i64> {
+    transform.smt.constrain_params(%param_as_param) : (!transform.param<i64>) -> () {
       ^bb0(%param_as_smt_var: !smt.int):
       %c4 = arith.constant 4 : i32
       // This is the kind of thing one might think works:
@@ -17,13 +46,90 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-// CHECK-LABEL: @operands_not_one_to_one_with_vars
+// CHECK-LABEL: @results_not_one_to_one_with_vars
 module attributes {transform.with_named_sequence} {
-  transform.named_sequence @operands_not_one_to_one_with_vars(%arg0: !transform.any_op {transform.readonly}) {
+  transform.named_sequence @results_not_one_to_one_with_vars(%arg0: !transform.any_op {transform.readonly}) {
     %param_as_param = transform.param.constant 42 -> !transform.param<i64>
-    // expected-error@below {{must have the same number of block arguments as operands}}
-    transform.smt.constrain_params(%param_as_param) : !transform.param<i64> {
+    transform.smt.constrain_params(%param_as_param, %param_as_param) : (!transform.param<i64>, !transform.param<i64>) -> () {
       ^bb0(%param_as_smt_var: !smt.int, %param_as_another_smt_var: !smt.int):
+      // expected-error@below {{expected terminator to have as many operands as the parent op has results}}
+      smt.yield %param_as_smt_var : !smt.int
+    }
+    transform.yield
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @non_smt_type_block_args
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @non_smt_type_block_args(%arg0: !transform.any_op {transform.readonly}) {
+    %param_as_param = transform.param.constant 42 -> !transform.param<i8>
+    // expected-error@below {{the type of block arg #0 is expected to be either a !smt.bool, a !smt.int, or a !smt.bv}}
+    transform.smt.constrain_params(%param_as_param) : (!transform.param<i8>) -> (!transform.param<i8>) {
+      ^bb0(%param_as_smt_var: !transform.param<i8>):
+      smt.yield %param_as_smt_var : !transform.param<i8>
+    }
+    transform.yield
+  }
+}
+
+
+// -----
+
+// CHECK-LABEL: @mismatched_arg_type_bool
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @mismatched_arg_type_bool(%arg0: !transform.any_op {transform.readonly}) {
+    %param_as_param = transform.param.constant 42 -> !transform.param<i64>
+    // expected-error@below {{the type of block arg #0 is !smt.bool though the corresponding operand type ('!transform.param<i64>') is not wrapping i1}}
+    transform.smt.constrain_params(%param_as_param) : (!transform.param<i64>) -> (!transform.param<i64>) {
+      ^bb0(%param_as_smt_var: !smt.bool):
+      smt.yield %param_as_smt_var : !smt.bool
+    }
+    transform.yield
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @mismatched_arg_type_bitvector
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @mismatched_arg_type_bitvector(%arg0: !transform.any_op {transform.readonly}) {
+    %param_as_param = transform.param.constant 42 -> !transform.param<i64>
+    // expected-error@below {{the type of block arg #0 is '!smt.bv<8>' though the corresponding operand type ('!transform.param<i64>') is not wrapping an integer type of the same bitwidth}}
+    transform.smt.constrain_params(%param_as_param) : (!transform.param<i64>) -> (!transform.param<i64>) {
+      ^bb0(%param_as_smt_var: !smt.bv<8>):
+      smt.yield %param_as_smt_var : !smt.bv<8>
+    }
+    transform.yield
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @mismatched_result_type_bool
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @mismatched_result_type_bool(%arg0: !transform.any_op {transform.readonly}) {
+    %param_as_param = transform.param.constant 1 -> !transform.param<i1>
+    transform.smt.constrain_params(%param_as_param) : (!transform.param<i1>) -> (!transform.param<i64>) {
+      ^bb0(%param_as_smt_var: !smt.bool):
+      // expected-error@below {{the type of terminator operand #0 is !smt.bool though the corresponding result type ('!transform.param<i64>') is not wrapping i1}}
+      smt.yield %param_as_smt_var : !smt.bool
+    }
+    transform.yield
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @mismatched_result_type_bitvector
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @mismatched_result_type_bitvector(%arg0: !transform.any_op {transform.readonly}) {
+    %param_as_param = transform.param.constant 42 -> !transform.param<i8>
+    transform.smt.constrain_params(%param_as_param) : (!transform.param<i8>) -> (!transform.param<i64>) {
+      ^bb0(%param_as_smt_var: !smt.bv<8>):
+      // expected-error@below {{the type of terminator operand #0 is '!smt.bv<8>' though the corresponding result type ('!transform.param<i64>') is not wrapping an integer type of the same bitwidth}}
+      smt.yield %param_as_smt_var : !smt.bv<8>
     }
     transform.yield
   }