[tmtensor] Add support for i64 index type for tm_tensor.scatter (#4221)

using i32 indexing for large memory blocks is wrong

---------

Co-authored-by: Rob Suderman <[email protected]>

Author: Groverkss

lib/Conversion/TorchToTMTensor/TorchToTMTensor.cpp

@@ -841,7 +841,6 @@ class ConvertAtenIndexPutHackedTwinOp
     if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
       return failure();
     Location loc = op.getLoc();
-    MLIRContext *context = op->getContext();
     Value input = op.getSelf();
     Value values = op.getValues();
     auto inputType = cast<ValueTensorType>(input.getType());
@@ -962,11 +961,6 @@ class ConvertAtenIndexPutHackedTwinOp
     values =
         rewriter.create<AtenViewOp>(loc, valuesType, values, valuesDimsList);
 
-    // `TMTensor::ScatterOp` expects indices of element type i32.
-    indices = convertTensorToDtype(
-        rewriter, loc, indices,
-        mlir::IntegerType::get(context, 32, mlir::IntegerType::Signed));
-
     input = typeConverter->materializeTargetConversion(
         rewriter, loc, typeConverter->convertType(input.getType()), input);
     values = typeConverter->materializeTargetConversion(

lib/Dialect/TMTensor/IR/TMTensorOps.cpp

@@ -623,8 +623,9 @@ LogicalResult ScatterOp::verify() {
 
   auto indicesType = getIndicesType();
   if (indicesType.getRank() != 2 ||
-      !indicesType.getElementType().isInteger(32)) {
-    return emitOpError("expected indices to be of rank 2 of i32 element type");
+      !isa<IntegerType>(indicesType.getElementType())) {
+    return emitOpError(
+        "expected indices to be of rank 2 of integer element type");
   }
   auto indexDepth = getIndexDepth();
   if (ShapedType::isDynamic(indexDepth)) {