Repository: tinygrad/tinygrad

Issue:
Tensor.where() bug on LLVM
when either `input_` or `other` is of `dtypes.bool`, llvmir attempts to turn bool into float: `%".16" = fpext i1 %".15" to float` but this raises 
```
RuntimeError: LLVM IR parsing error
<string>:19:18: error: invalid cast opcode for cast from 'i1' to 'float'
```

to reproduce:
```
from tinygrad.lazy import Device
from tinygrad.tensor import Tensor
from tinygrad.helpers import dtypes
Device.DEFAULT = "LLVM"
cond = Tensor([0,1]) == Tensor([1,1]);
cond.where(Tensor([1,1], dtype=dtypes.bool), 0).realize()
```

prg:
```
; ModuleID = "/Users/zibo/fun/tinygrad/tinygrad/renderer/llvmir.py"                                                                              
target triple = "unknown-unknown-unknown"
target datalayout = ""                                                                                                                           
                                                                                                                                                 define void @"E_2n1"(float* %".1", float* %".2", float* %".3", float* %".4", i1* %".5") "no-nans-fp-math"="true"
{                                                                       
entry:
  br label %"loop_body_gidx0"                                                                                                                    
loop_body_gidx0:                   
  %"gidx0" = phi  i64 [0, %"entry"], [%".23", %"loop_body_gidx0"] 
  %".8" = getelementptr inbounds float, float* %".2", i64 %"gidx0"                                                                               
  %".9" = load float, float* %".8"                                      
  %".10" = getelementptr inbounds float, float* %".3", i64 %"gidx0"
  %".11" = load float, float* %".10" 
  %".12" = getelementptr inbounds float, float* %".4", i64 0                                                                                     
  %".13" = load float, float* %".12"                                                                                                             
  %".14" = getelementptr inbounds i1, i1* %".5", i64 %"gidx0"
  %".15" = load i1, i1* %".14"   
  %".16" = fpext i1 %".15" to float   # HERE IT IS !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  %".17" = fcmp fast oeq float %".9", %".11"                            
  %".18" = uitofp i1 %".17" to float 
  %".19" = fcmp fast une float %".18",              0x0                                                                                          
  %".20" = select fast i1 %".19", float %".13", float %".16"           
  %".21" = getelementptr inbounds float, float* %".1", i64 %"gidx0"                                                                              
  store float %".20", float* %".21"                                                                                                              
  %".23" = add i64 %"gidx0", 1
  %".24" = icmp eq i64 %".23", 2                                                                                                                 
  br i1 %".24", label %"loop_exit_gidx0", label %"loop_body_gidx0"
loop_exit_gidx0:                                                        
  ret void           
}
```



Discussion:

Resolving patch:
diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py
--- a/tinygrad/renderer/llvmir.py
+++ b/tinygrad/renderer/llvmir.py
@@ -104,8 +104,8 @@ def uops_to_llvm_ir(function_name:str, uops:List[UOp]) -> Tuple[str, Optional[Li
           val = bb[-1].load(bb[-1].gep(func.args[buf_index[args.name]], [idx], inbounds=True))
 
         if args.memory_dtype != newvar.dtype:
-          if dtypes.is_int(args.memory_dtype):
-            val = bb[-1].uitofp(val, ir.FloatType()) if dtypes.is_unsigned(args.memory_dtype) else bb[-1].sitofp(val, ir.FloatType())
+          if dtypes.is_int(args.memory_dtype) or args.memory_dtype == dtypes.bool:
+            val = bb[-1].uitofp(val, ir.FloatType()) if dtypes.is_unsigned(args.memory_dtype) or args.memory_dtype == dtypes.bool else bb[-1].sitofp(val, ir.FloatType())
           elif args.memory_dtype == dtypes.bfloat16:
             val = bb[-1].sext(val, ir.IntType(32))
             val = bb[-1].shl(val, ir.Constant(ir.IntType(32), 16))


Test patch:
diff --git a/test/test_ops.py b/test/test_ops.py
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -4,7 +4,7 @@
 import numpy as np
 import unittest
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import getenv, IMAGE, DEBUG, CI
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, dtypes
 from tinygrad.lazy import Device
 
 if CI:
@@ -150,6 +150,10 @@ def test_where_permute(self):
       lambda x: torch.where(x > 0.5, 4, 2).permute((1, 0)),
       lambda x: (x > 0.5).where(4, 2).permute((1, 0)), forward_only=True)
 
+  def test_where_bool(self): # Fixes #1479.
+    helper_test_op([(1,), (1,)], lambda x,y: torch.where(x==y, torch.tensor([1,1], dtype=torch.bool), 0),
+      lambda x,y: (x==y).where(Tensor([1,1], dtype=dtypes.bool), 0), forward_only=True, vals=[[0,1],[1,1]])
+
   def _test_cmp(self, fxn, reverse=True):
     for shps in [[(3, 4, 5), (3, 4, 5)], [(3, 4, 5), (5,)], [(5,), (3, 4, 5)]]:
       helper_test_op(shps, fxn, fxn, forward_only=True)