Custom mask bshd attention variant (#665)

Added support for custom mask bshd kernel
Added `operator.eq` and `bitwise or`
Added implementation for `scaled_dot_product_attention` to verify the
numeric.

---------

Signed-off-by: Stanley Winata <[email protected]>
Signed-off-by: xintin <[email protected]>
Co-authored-by: Stanley Winata <[email protected]>
Signed-off-by: nithinsubbiah <[email protected]>

Author: 2 people

iree/turbine/aot/support/ir_utils.py

@@ -497,6 +497,10 @@ def _is_integer_like_type(type):
     return isinstance(type, (IntegerType, IndexType))
 
 
+def _is_signed_or_signless_type(type):
+    return getattr(type, "is_signed", False) or getattr(type, "is_signless", False)
+
+
 def _attribute_from_device_affinity(
     affinity: DeviceAffinity, context: Context
 ) -> Attribute:

iree/turbine/kernel/ops/wave_ops.py

@@ -202,6 +202,10 @@ def le(lhs: "Register", rhs: "Register") -> "Register":
     ...
 
 
+def eq(lhs: "Register", rhs: "Register") -> "Register":
+    ...
+
+
 def cast(src: "Register", dtype: DataType) -> "Register":
     ...
 
@@ -780,6 +784,7 @@ def infer_shape(self) -> Any:
 @define_py_op(operator.sub)
 @define_py_op(operator.mul)
 @define_py_op(operator.and_)
+@define_py_op(operator.or_)
 @define_py_op(operator.truediv)
 @define_interface_op("maximum")
 @define_interface_op("minimum")
@@ -789,10 +794,12 @@ def infer_type(self):
         self.type = Register[(*self.infer_shape(), get_custom(self.lhs).type.dtype)]
 
 
+@define_py_op(operator.eq)
 @define_py_op(operator.gt)
 @define_py_op(operator.ge)
 @define_py_op(operator.lt)
 @define_py_op(operator.le)
+@define_interface_op("eq")
 @define_interface_op("gt")
 @define_interface_op("ge")
 @define_interface_op("lt")

iree/turbine/kernel/wave/codegen/handlers.py

@@ -47,6 +47,7 @@
     _is_float_type,
     _is_index_type,
     _is_integer_like_type,
+    _is_signed_or_signless_type,
 )
 
 # TK infrastructure imports.
@@ -58,6 +59,7 @@
     broadcast,
     cast,
     conditional,
+    eq,
     exp2,
     extract,
     extract_slice,
@@ -444,8 +446,8 @@ def handle_div(lhs: Value, rhs: Value) -> OpResult:
     element_type = get_type_or_element_type(lhs.type)
     if _is_float_type(element_type):
         result = arith_d.divf(lhs, rhs)
-    elif _is_integer_like_type(element_type) and (
-        element_type.is_signed or element_type.is_signless
+    elif _is_integer_like_type(element_type) and _is_signed_or_signless_type(
+        element_type
     ):
         result = arith_d.divsi(lhs, rhs)
     else:
@@ -456,12 +458,28 @@ def handle_div(lhs: Value, rhs: Value) -> OpResult:
 @handle_binary_op(operator.and_)
 def handle_and(lhs: Value, rhs: Value) -> OpResult:
     element_type = get_type_or_element_type(lhs.type)
-    if _is_integer_like_type(element_type) and (
-        element_type.is_signed or element_type.is_signless
+    if _is_integer_like_type(element_type) and _is_signed_or_signless_type(
+        element_type
     ):
         result = arith_d.andi(lhs, rhs)
     else:
-        raise ValidationError(f"Found unhandled operand type for div: {element_type}")
+        raise ValidationError(
+            f"Found unhandled operand type for bitwise and: {element_type}"
+        )
+    return result
+
+
+@handle_binary_op(operator.or_)
+def handle_or(lhs: Value, rhs: Value) -> OpResult:
+    element_type = get_type_or_element_type(lhs.type)
+    if _is_integer_like_type(element_type) and _is_signed_or_signless_type(
+        element_type
+    ):
+        result = arith_d.ori(lhs, rhs)
+    else:
+        raise ValidationError(
+            f"Found unhandled operand type for bitwise or: {element_type}"
+        )
     return result
 
 
@@ -470,8 +488,8 @@ def handle_gt(lhs: Value, rhs: Value) -> OpResult:
     element_type = get_type_or_element_type(lhs.type)
     if _is_float_type(element_type):
         result = arith_d.cmpi(arith_d.CmpFPredicate.OGT, lhs, rhs)
-    elif _is_integer_like_type(element_type) and (
-        element_type.is_signed or element_type.is_signless
+    elif _is_integer_like_type(element_type) and _is_signed_or_signless_type(
+        element_type
     ):
         result = arith_d.cmpi(arith_d.CmpIPredicate.sgt, lhs, rhs)
     else:
@@ -484,8 +502,8 @@ def handle_ge(lhs: Value, rhs: Value) -> OpResult:
     element_type = get_type_or_element_type(lhs.type)
     if _is_float_type(element_type):
         result = arith_d.cmpi(arith_d.CmpFPredicate.OGE, lhs, rhs)
-    elif _is_integer_like_type(element_type) and (
-        element_type.is_signed or element_type.is_signless
+    elif _is_integer_like_type(element_type) and _is_signed_or_signless_type(
+        element_type
     ):
         result = arith_d.cmpi(arith_d.CmpIPredicate.sge, lhs, rhs)
     else:
@@ -498,8 +516,8 @@ def handle_lt(lhs: Value, rhs: Value) -> OpResult:
     element_type = get_type_or_element_type(lhs.type)
     if _is_float_type(element_type):
         result = arith_d.cmpi(arith_d.CmpFPredicate.OLT, lhs, rhs)
-    elif _is_integer_like_type(element_type) and (
-        element_type.is_signed or element_type.is_signless
+    elif _is_integer_like_type(element_type) and _is_signed_or_signless_type(
+        element_type
     ):
         result = arith_d.cmpi(arith_d.CmpIPredicate.slt, lhs, rhs)
     else:
@@ -512,22 +530,36 @@ def handle_le(lhs: Value, rhs: Value) -> OpResult:
     element_type = get_type_or_element_type(lhs.type)
     if _is_float_type(element_type):
         result = arith_d.cmpi(arith_d.CmpFPredicate.OLE, lhs, rhs)
-    elif _is_integer_like_type(element_type) and (
-        element_type.is_signed or element_type.is_signless
+    elif _is_integer_like_type(element_type) and _is_signed_or_signless_type(
+        element_type
     ):
         result = arith_d.cmpi(arith_d.CmpIPredicate.sle, lhs, rhs)
     else:
         raise ValidationError(f"Found unhandled operand type for le: {element_type}")
     return result
 
 
+@handle_binary_op([operator.eq, eq])
+def handle_le(lhs: Value, rhs: Value) -> OpResult:
+    element_type = get_type_or_element_type(lhs.type)
+    if _is_float_type(element_type):
+        result = arith_d.cmpf(arith_d.CmpFPredicate.OEQ, lhs, rhs)
+    elif _is_integer_like_type(element_type) and _is_signed_or_signless_type(
+        element_type
+    ):
+        result = arith_d.cmpi(arith_d.CmpIPredicate.eq, lhs, rhs)
+    else:
+        raise ValidationError(f"Found unhandled operand type for eq: {element_type}")
+    return result
+
+
 @handle_binary_op(maximum)
 def handle_maximum(lhs: Value, rhs: Value) -> OpResult:
     element_type = get_type_or_element_type(lhs.type)
     if _is_float_type(element_type):
         result = arith_d.maximumf(lhs, rhs)
-    elif _is_integer_like_type(element_type) and (
-        element_type.is_signed or element_type.is_signless
+    elif _is_integer_like_type(element_type) and _is_signed_or_signless_type(
+        element_type
     ):
         result = arith_d.maxsi(lhs, rhs)
     else:
@@ -542,8 +574,8 @@ def handle_minimum(lhs: Value, rhs: Value) -> OpResult:
     element_type = get_type_or_element_type(lhs.type)
     if _is_float_type(element_type):
         result = arith_d.minimumf(lhs, rhs)
-    elif _is_integer_like_type(element_type) and (
-        element_type.is_signed or element_type.is_signless
+    elif _is_integer_like_type(element_type) and _is_signed_or_signless_type(
+        element_type
     ):
         result = arith_d.minsi(lhs, rhs)
     else: