Type materialization bug when lowering polynomial to standard

[j2kun]

The following code fails when run with --polynomial-to-standard and --heir-polynomial-to-llvm

// Regression test for https://github.com/google/heir/issues/505

// RUN: heir-opt --polynomial-to-standard %s | FileCheck %s

func.func @test_bin_ops(%arg0: memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>, %arg1: memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>) -> memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>> {
  %alloc = memref.alloc() {alignment = 64 : i64} : memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>
  affine.for %arg2 = 0 to 2 {
    %3 = affine.load %arg0[%arg2] : memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>>
    %4 = affine.load %arg1[%arg2] : memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>>
    %5 = polynomial.add(%3, %4) : !polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>
    affine.store %5, %alloc[%arg2] : memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>
  }
  return %alloc : memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>
}

Error:

/home/j2kun/fhe/heir/tests/polynomial/issue_505.mlir:8:22: error: use of value '%arg0' expects different type than prior uses: 'memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>' vs 'memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>'
    %3 = affine.load %arg0[%arg2] : memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>
                     ^
/home/j2kun/fhe/heir/tests/polynomial/issue_505.mlir:5:25: note: prior use here
func.func @test_bin_ops(%arg0: memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>, %arg1: memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>) -> memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>> {

[j2kun]

CC @AlexanderViand-Intel

[AlexanderViand-Intel]

Actually, looking at this a bit closer, the code you added above fails to parse because of the , strided<[?], offset: ?>> at the end of the type for %3 and %4 😉 but even after correcting that, this runs into the problem from #143, just for memref instead of for tensor. I extended my frankenstein verison of -polynomial-to-standard with an AddConversion for memrefs (https://github.com/AlexanderViand-Intel/heir/commit/aa5d3899e6c546ea54e13afe28f366accad08905) and running the example through that results in:

func.func @test_bin_ops(%arg0: memref<2x1024xi25>, %arg1: memref<2x1024xi25>) -> memref<2x1024xi25> {
  %0 = builtin.unrealized_conversion_cast %arg1 : memref<2x1024xi25> to memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>
  %1 = builtin.unrealized_conversion_cast %arg0 : memref<2x1024xi25> to memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>
  %alloc = memref.alloc() {alignment = 64 : i64} : memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>
  %2 = builtin.unrealized_conversion_cast %alloc : memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>> to memref<2x1024xi25>
  affine.for %arg2 = 0 to 2 {
    %3 = affine.load %1[%arg2] : memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>
    %4 = builtin.unrealized_conversion_cast %3 : !polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>> to tensor<1024xi25>
    %5 = affine.load %0[%arg2] : memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>
    %6 = builtin.unrealized_conversion_cast %5 : !polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>> to tensor<1024xi25>
    %cst = arith.constant dense<33538049> : tensor<1024xi26>
    %7 = arith.extsi %4 : tensor<1024xi25> to tensor<1024xi26>
    %8 = arith.extsi %6 : tensor<1024xi25> to tensor<1024xi26>
    %9 = arith.addi %7, %8 : tensor<1024xi26>
    %10 = arith.remsi %9, %cst : tensor<1024xi26>
    %11 = arith.trunci %10 : tensor<1024xi26> to tensor<1024xi25>
    %12 = builtin.unrealized_conversion_cast %11 : tensor<1024xi25> to !polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>
    affine.store %12, %alloc[%arg2] : memref<2x!polynomial.polynomial<<cmod=33538049, ideal=#polynomial.polynomial<1 + x**1024>>>>
  }
  return %2 : memref<2x1024xi25>
}

So another case where the unwanted "old" type persist in operations that aren't marked (dynamically) illegal.