Fixing dynamic complex number splats in stream encoding. (#14100) The original PR #13363 had been using IndexCastOp instead of ExtUIOp. Fixes #14090.

diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/EncodeTensors.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/EncodeTensors.cpp
index c221ec0..d15a9db 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/EncodeTensors.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/EncodeTensors.cpp

@@ -146,8 +146,8 @@
     Value realInt = builder.create<arith::BitcastOp>(loc, bwElemType, real);
     Value imag = builder.create<mlir::complex::ImOp>(loc, pattern);
     Value imagInt = builder.create<arith::BitcastOp>(loc, bwElemType, imag);
-    realInt = builder.create<arith::IndexCastOp>(loc, bwType, realInt);
-    imagInt = builder.create<arith::IndexCastOp>(loc, bwType, imagInt);
+    realInt = builder.create<arith::ExtUIOp>(loc, bwType, realInt);
+    imagInt = builder.create<arith::ExtUIOp>(loc, bwType, imagInt);
     Value shiftReal =
         builder.create<arith::ShLIOp>(loc, bwType, realInt, shiftAmount);
     Value orImag = builder.create<arith::OrIOp>(loc, shiftReal, imagInt);
@@ -596,6 +596,7 @@
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<mlir::func::FuncDialect>();
     registry.insert<mlir::arith::ArithDialect>();
+    registry.insert<mlir::complex::ComplexDialect>();
     registry.insert<IREE::Stream::StreamDialect>();
     registry.insert<IREE::Util::UtilDialect>();
   }

diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors.mlir
index de2f759..00848d4 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors.mlir

@@ -109,15 +109,35 @@
 
 // -----
 
-// CHECK-LABEL: @denseTensorSplatComplexF32
-func.func @denseTensorSplatComplexF32(%arg0: !stream.resource<*>) -> (!stream.resource<*>) {
-    %cst = complex.constant [3.000000e+00 : f32, 1.000000e+01 : f32] : complex<f32>
-    %0 = stream.tensor.sizeof tensor<6xcomplex<f32>> : index
-    // CHECK: %[[I64NUMBER:.+]] = arith.constant 4629700418029486080
-    // CHECK: %[[SPLAT_RES:.+]] = stream.async.splat %[[I64NUMBER]]
-    %1 = stream.tensor.splat %cst : complex<f32> -> tensor<6xcomplex<f32>> in !stream.resource<*>{%0}
-    return %1 : !stream.resource<*>
-  }
+// CHECK-LABEL: @denseTensorSplatConstantComplexF32
+func.func @denseTensorSplatConstantComplexF32(%arg0: !stream.resource<*>) -> (!stream.resource<*>) {
+  %cst = complex.constant [3.000000e+00 : f32, 1.000000e+01 : f32] : complex<f32>
+  %0 = stream.tensor.sizeof tensor<6xcomplex<f32>> : index
+  // CHECK: %[[I64NUMBER:.+]] = arith.constant 4629700418029486080
+  // CHECK: %[[SPLAT_RES:.+]] = stream.async.splat %[[I64NUMBER]]
+  %1 = stream.tensor.splat %cst : complex<f32> -> tensor<6xcomplex<f32>> in !stream.resource<*>{%0}
+  // CHECK: return %[[SPLAT_RES]]
+  return %1 : !stream.resource<*>
+}
+
+// -----
+
+// CHECK-LABEL: @denseTensorSplatDynamicComplexF32
+func.func @denseTensorSplatDynamicComplexF32(%arg0: !stream.resource<*>, %arg1: complex<f32>) -> (!stream.resource<*>) {
+  %0 = stream.tensor.sizeof tensor<6xcomplex<f32>> : index
+  // CHECK-DAG: %[[REAL_F32:.+]] = complex.re %arg1 : complex<f32>
+  // CHECK-DAG: %[[REAL_I32:.+]] = arith.bitcast %[[REAL_F32]] : f32 to i32
+  // CHECK-DAG: %[[IMAG_F32:.+]] = complex.im %arg1 : complex<f32>
+  // CHECK-DAG: %[[IMAG_I32:.+]] = arith.bitcast %[[IMAG_F32]] : f32 to i32
+  // CHECK-DAG: %[[REAL_I64:.+]] = arith.extui %[[REAL_I32]] : i32 to i64
+  // CHECK-DAG: %[[IMAG_I64:.+]] = arith.extui %[[IMAG_I32]] : i32 to i64
+  // CHECK-DAG: %[[UPPER:.+]] = arith.shli %[[REAL_I64]], %c32
+  // CHECK-DAG: %[[VALUE:.+]] = arith.ori %[[UPPER]], %[[IMAG_I64]]
+  // CHECK: %[[SPLAT_RES:.+]] = stream.async.splat %[[VALUE]]
+  %1 = stream.tensor.splat %arg1 : complex<f32> -> tensor<6xcomplex<f32>> in !stream.resource<*>{%0}
+  // CHECK: return %[[SPLAT_RES]]
+  return %1 : !stream.resource<*>
+}
 
 // -----