Hello, I’m trying to implement an Operation like ArgMax to get the index of max value based on axis. If we use a 1D tensor as input, the return should be one interger value. I guess basic idea is like:
%A = alloc() : memref<10xf64>
%1 = constant 0 :index
%2 = dim %A, %1 : memref<10xf64>
%3 = alloc() : memref<1xi64>
affine.for %arg0 = 0 to %2 {
// do compare and store
}
return %3 : memref<1xi64>
I wish to get value of %arg0 during loop and store it to %3 if %A[%arg0] is the visited largest one, however it can not store index type to memref<1xi64>. Index Type is a signless integer so can it be converted to integer type for storing? Or there is a better way to implement ArgMax? Thanks a lot the help.