[mlir][linalg] Add unsigned min/max/cast function to OpDSL.

Update OpDSL to support unsigned integers by adding unsigned min/max/cast signatures. Add tests in OpDSL and on the C++ side to verify the proper signed and unsigned operations are emitted.

The patch addresses an issue brought up in https://reviews.llvm.org/D111170.

Reviewed By: rsuderman

Differential Revision: https://reviews.llvm.org/D111230

Author: Tobias Gysi

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

@@ -196,31 +196,40 @@ class RegionBuilderHelper {
   // If the cast cannot be performed, a warning will be issued and the
   // operand returned as-is (which will presumably yield a verification
   // issue downstream).
-  Value cast(Type toType, Value operand) {
+  Value cast(Type toType, Value operand, bool isUnsignedCast) {
     OpBuilder builder = getBuilder();
     auto loc = operand.getLoc();
 
     if (operand.getType() == toType)
       return operand;
     if (auto toIntType = toType.dyn_cast<IntegerType>()) {
       // If operand is floating point, cast directly to the int type.
-      if (operand.getType().isa<FloatType>())
+      if (operand.getType().isa<FloatType>()) {
+        if (isUnsignedCast)
+          return builder.create<FPToUIOp>(loc, toType, operand);
         return builder.create<FPToSIOp>(loc, toType, operand);
+      }
       // Cast index operands directly to the int type.
       if (operand.getType().isIndex())
         return builder.create<IndexCastOp>(loc, toType, operand);
       if (auto fromIntType = operand.getType().dyn_cast<IntegerType>()) {
-        // Either sign extend or truncate.
-        if (toIntType.getWidth() > fromIntType.getWidth())
+        // Either extend or truncate.
+        if (toIntType.getWidth() > fromIntType.getWidth()) {
+          if (isUnsignedCast)
+            return builder.create<ZeroExtendIOp>(loc, toType, operand);
           return builder.create<SignExtendIOp>(loc, toType, operand);
+        }
         if (toIntType.getWidth() < fromIntType.getWidth())
           return builder.create<TruncateIOp>(loc, toType, operand);
       }
     } else if (auto toFloatType = toType.dyn_cast<FloatType>()) {
       // If operand is integer, cast directly to the float type.
       // Note that it is unclear how to cast from BF16<->FP16.
-      if (operand.getType().isa<IntegerType>())
+      if (operand.getType().isa<IntegerType>()) {
+        if (isUnsignedCast)
+          return builder.create<UIToFPOp>(loc, toFloatType, operand);
         return builder.create<SIToFPOp>(loc, toFloatType, operand);
+      }
       if (auto fromFloatType = operand.getType().dyn_cast<FloatType>()) {
         if (toFloatType.getWidth() > fromFloatType.getWidth())
           return builder.create<FPExtOp>(loc, toFloatType, operand);
@@ -284,6 +293,15 @@ class RegionBuilderHelper {
     llvm_unreachable("unsupported non numeric type");
   }
 
+  Value applyfn__max_unsigned(Value lhs, Value rhs) {
+    OpBuilder builder = getBuilder();
+    if (isFloatingPoint(lhs))
+      return builder.create<MaxFOp>(lhs.getLoc(), lhs, rhs);
+    if (isInteger(lhs))
+      return builder.create<MaxUIOp>(lhs.getLoc(), lhs, rhs);
+    llvm_unreachable("unsupported non numeric type");
+  }
+
   Value applyfn__min(Value lhs, Value rhs) {
     OpBuilder builder = getBuilder();
     if (isFloatingPoint(lhs))
@@ -293,6 +311,15 @@ class RegionBuilderHelper {
     llvm_unreachable("unsupported non numeric type");
   }
 
+  Value applyfn__min_unsigned(Value lhs, Value rhs) {
+    OpBuilder builder = getBuilder();
+    if (isFloatingPoint(lhs))
+      return builder.create<MinFOp>(lhs.getLoc(), lhs, rhs);
+    if (isInteger(lhs))
+      return builder.create<MinUIOp>(lhs.getLoc(), lhs, rhs);
+    llvm_unreachable("unsupported non numeric type");
+  }
+
   void yieldOutputs(ValueRange values) {
     assert(!values.empty() && "linalg ops must yield outputs");
     if (values.empty())

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

@@ -340,6 +340,8 @@ class PrimFn:
   max = PrimFnType("max")
   min = PrimFnType("min")
   sub = PrimFnType("sub")
+  max_unsigned = PrimFnType("max_unsigned")
+  min_unsigned = PrimFnType("min_unsigned")
 
 
 class ReduceFnType:
@@ -365,6 +367,8 @@ class ReduceFn:
   mul = PrimFn.mul.reduce
   max = PrimFn.max.reduce
   min = PrimFn.min.reduce
+  max_unsigned = PrimFn.max_unsigned.reduce
+  min_unsigned = PrimFn.min_unsigned.reduce
 
 
 class PrimApply(TensorExpression):
@@ -438,8 +442,8 @@ def __init__(self, to_type: TypeVar, operand: TensorExpression):
     self.operand = operand
 
   def to_scalar_expression(self) -> ScalarExpression:
-    return ScalarSymbolicCast(self.to_type,
-                              self.operand.to_scalar_expression()).expr()
+    return ScalarSymbolicCast(self.to_type, self.operand.to_scalar_expression(),
+                              False).expr()
 
   def visit_tensor_exprs(self, callback):
     super().visit_tensor_exprs(callback)
@@ -449,6 +453,17 @@ def __repr__(self):
     return f"cast({self.to_type}, {repr(self.operand)})"
 
 
+class cast_unsigned(cast):
+  """Casts the element type to an unsigned type (typically symbolic TypeVar)."""
+
+  def to_scalar_expression(self) -> ScalarExpression:
+    return ScalarSymbolicCast(self.to_type, self.operand.to_scalar_expression(),
+                              True).expr()
+
+  def __repr__(self):
+    return f"cast_unsigned({self.to_type}, {repr(self.operand)})"
+
+
 class ReduceApply(TensorExpression):
   """Application of a reduction.
 

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

@@ -230,10 +230,12 @@ def expression(self, expr: ScalarExpression) -> Value:
       return fn(*operand_values)
     elif expr.symbolic_cast:
       operand_value = self.expression(expr.symbolic_cast.operand)
-      return self.cast(expr.symbolic_cast.to_type.name, operand_value)
+      return self.cast(expr.symbolic_cast.to_type.name, operand_value,
+                       expr.symbolic_cast.is_unsigned_cast)
     raise NotImplementedError(f"Unimplemented scalar body expression: {expr}")
 
-  def cast(self, type_var_name: str, operand: Value) -> Value:
+  def cast(self, type_var_name: str, operand: Value,
+           is_unsigned_cast: bool) -> Value:
     try:
       to_type = self.type_mapping[type_var_name]
     except KeyError:
@@ -242,29 +244,37 @@ def cast(self, type_var_name: str, operand: Value) -> Value:
     if operand.type == to_type:
       return operand
     if _is_integer_type(to_type):
-      return self._cast_to_integer(to_type, operand)
+      return self._cast_to_integer(to_type, operand, is_unsigned_cast)
     elif _is_floating_point_type(to_type):
-      return self._cast_to_floating_point(to_type, operand)
+      return self._cast_to_floating_point(to_type, operand, is_unsigned_cast)
 
-  def _cast_to_integer(self, to_type: Type, operand: Value) -> Value:
+  def _cast_to_integer(self, to_type: Type, operand: Value,
+                       is_unsigned_cast: bool) -> Value:
     to_width = IntegerType(to_type).width
     operand_type = operand.type
     if _is_floating_point_type(operand_type):
+      if is_unsigned_cast:
+        return std.FPToUIOp(to_type, operand).result
       return std.FPToSIOp(to_type, operand).result
     if _is_index_type(operand_type):
       return std.IndexCastOp(to_type, operand).result
     # Assume integer.
     from_width = IntegerType(operand_type).width
     if to_width > from_width:
+      if is_unsigned_cast:
+        return std.ZeroExtendIOp(to_type, operand).result
       return std.SignExtendIOp(to_type, operand).result
     elif to_width < from_width:
       return std.TruncateIOp(to_type, operand).result
     raise ValueError(f"Unable to cast body expression from {operand_type} to "
                      f"{to_type}")
 
-  def _cast_to_floating_point(self, to_type: Type, operand: Value) -> Value:
+  def _cast_to_floating_point(self, to_type: Type, operand: Value,
+                              is_unsigned_cast: bool) -> Value:
     operand_type = operand.type
     if _is_integer_type(operand_type):
+      if is_unsigned_cast:
+        return std.UIToFPOp(to_type, operand).result
       return std.SIToFPOp(to_type, operand).result
     # Assume FloatType.
     to_width = _get_floating_point_width(to_type)
@@ -324,13 +334,26 @@ def _eval_max(self, lhs: Value, rhs: Value) -> Value:
       return std.MaxSIOp(lhs.type, lhs, rhs).result
     raise NotImplementedError("Unsupported 'max' operand: {lhs}")
 
+  def _eval_max_unsigned(self, lhs: Value, rhs: Value) -> Value:
+    if _is_floating_point_type(lhs.type):
+      return std.MaxFOp(lhs.type, lhs, rhs).result
+    if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
+      return std.MaxUIOp(lhs.type, lhs, rhs).result
+    raise NotImplementedError("Unsupported 'max_unsigned' operand: {lhs}")
+
   def _eval_min(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return std.MinFOp(lhs.type, lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return std.MinSIOp(lhs.type, lhs, rhs).result
     raise NotImplementedError("Unsupported 'min' operand: {lhs}")
 
+  def _eval_min_unsigned(self, lhs: Value, rhs: Value) -> Value:
+    if _is_floating_point_type(lhs.type):
+      return std.MinFOp(lhs.type, lhs, rhs).result
+    if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
+      return std.MinUIOp(lhs.type, lhs, rhs).result
+    raise NotImplementedError("Unsupported 'min_unsigned' operand: {lhs}")
 
 def _infer_structured_outs(op_config: LinalgStructuredOpConfig,
                            in_arg_defs: Sequence[OperandDefConfig],

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

@@ -85,15 +85,17 @@ def __repr__(self):
 class ScalarSymbolicCast:
   """A type of ScalarExpression that symbolically casts an operand to a TypeVar."""
 
-  def __init__(self, to_type: TypeVar, operand: "ScalarExpression"):
+  def __init__(self, to_type: TypeVar, operand: "ScalarExpression",
+               is_unsigned_cast: bool):
     self.to_type = to_type
     self.operand = operand
+    self.is_unsigned_cast = is_unsigned_cast
 
   def expr(self) -> "ScalarExpression":
     return ScalarExpression(symbolic_cast=self)
 
   def __repr__(self):
-    return f"ScalarSymbolicCast({self.to_type}, {self.operand})"
+    return f"ScalarSymbolicCast({self.to_type}, {self.operand}, {self.is_unsigned_cast})"
 
 
 class ScalarExpression(YAMLObject):
@@ -144,7 +146,8 @@ def to_yaml_custom_dict(self):
       return dict(
           symbolic_cast=dict(
               type_var=self.symbolic_cast.to_type.name,
-              operands=[self.symbolic_cast.operand]))
+              operands=[self.symbolic_cast.operand],
+              is_unsigned_cast=self.symbolic_cast.is_unsigned_cast))
     else:
       raise ValueError(f"Unexpected ScalarExpression type: {self}")
 

mlir/python/mlir/dialects/linalg/opdsl/lang/scalar_expr.py

@@ -20,6 +20,20 @@ def matmul(
   implements(ContractionOpInterface)
   C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
 
+@linalg_structured_op
+def matmul_unsigned(
+    A=TensorDef(T1, S.M, S.K),
+    B=TensorDef(T2, S.K, S.N),
+    C=TensorDef(U, S.M, S.N, output=True)):
+  """Performs an unsigned matrix multiplication of two 2D inputs.
+
+  Numeric casting is performed on the operands to the inner multiply, promoting
+  them to the same data type as the accumulator/output.
+  """
+  domain(D.m, D.n, D.k)
+  implements(ContractionOpInterface)
+  C[D.m, D.n] += cast_unsigned(U, A[D.m, D.k]) * cast_unsigned(U, B[D.k, D.n])
+
 @linalg_structured_op
 def quantized_matmul(
     A=TensorDef(T1, S.M, S.K),
@@ -411,6 +425,24 @@ def pooling_nhwc_max(
       cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
                 D.c]))
 
+@linalg_structured_op
+def pooling_nhwc_max_unsigned(
+    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.C),
+    K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
+    O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
+    strides=AttributeDef(S.SH, S.SW),
+    dilations=AttributeDef(S.DH, S.DW)):
+  """Performs unsigned max pooling.
+
+  Numeric casting is performed on the input operand, promoting it to the same
+  data type as the accumulator/output.
+  """
+  implements(ConvolutionOpInterface)
+  domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
+  O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_unsigned(D.kh, D.kw)(
+      cast_unsigned(
+          U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
+
 @linalg_structured_op
 def pooling_nchw_max(
     I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW),
@@ -447,6 +479,23 @@ def pooling_nhwc_min(
       cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
                 D.c]))
 
+@linalg_structured_op
+def pooling_nhwc_min_unsigned(
+    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.C),
+    K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
+    O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
+    strides=AttributeDef(S.SH, S.SW),
+    dilations=AttributeDef(S.DH, S.DW)):
+  """Performs unsigned min pooling.
+
+  Numeric casting is performed on the input operand, promoting it to the same
+  data type as the accumulator/output.
+  """
+  implements(ConvolutionOpInterface)
+  domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
+  O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_unsigned(D.kh, D.kw)(
+      cast_unsigned(
+          U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 @linalg_structured_op
 def pooling_ndhwc_sum(